Refactor to allow multiple initialization listeners to be attached

Author: Ian Bird

sdk/src/main/java/com/uid2/UID2Manager.kt

@@ -41,6 +41,9 @@ import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.flow.retryWhen
 import kotlinx.coroutines.flow.single
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.sync.Mutex
+import kotlinx.coroutines.sync.withLock
 
 /**
  * A listener interface allowing the consumer to be notified when either the identity or status of the identity changes
@@ -99,22 +102,6 @@ public class UID2Manager internal constructor(
      */
     public var onIdentityChangedListener: UID2ManagerIdentityChangedListener? = null
 
-    /**
-     * Gets or sets a listener which can be used to determine if the [UID2Manager] instance has finished initializing.
-     * Initializing includes any time required to restore a previously persisted [UID2Identity] from storage.
-     *
-     * If this property is set *after* initialization is complete, the callback will be invoked immediately.
-     */
-    public var onInitialized: (() -> Unit)? = null
-        set(value) {
-            field = value
-
-            // If we've already finished initializing, we should immediately invoke the callback.
-            if (initialized.isCompleted) {
-                value?.invoke()
-            }
-        }
-
     private val _state = MutableStateFlow<UID2ManagerState>(Loading)
 
     /**
@@ -123,8 +110,10 @@ public class UID2Manager internal constructor(
     public val state: Flow<UID2ManagerState> = _state.asStateFlow()
 
     // The Job responsible for initialising the manager. This will include de-serialising our initial state from
-    // storage.
+    // storage. We allow consumers to attach a listener to detect when this Job is complete.
     private var initialized: Job
+    private val onInitializedListeners = mutableListOf<() -> Unit>()
+    private val initializedLock = Mutex()
 
     // An active Job that is scheduled to refresh the current identity
     private var refreshJob: Job? = null
@@ -177,6 +166,35 @@ public class UID2Manager internal constructor(
             checkIdentityRefresh()
         }
 
+    /**
+     * Adds a listener which can be used to determine if the [UID2Manager] instance has finished initializing.
+     * Initializing includes any time required to restore a previously persisted [UID2Identity] from storage.
+     *
+     * If a listener is added *after* initialization is complete, the callback will be invoked immediately.
+     */
+    public fun addOnInitializedListener(listener: () -> Unit): Unit = runBlocking {
+        initializedLock.withLock {
+            // If we've already finished initializing, we should immediately invoke the callback.
+            if (initialized.isCompleted) {
+                listener()
+            } else {
+                onInitializedListeners += listener
+            }
+        }
+    }
+
+    /**
+     * Removes a listener which was previously added via [addOnInitializedListener].
+     *
+     * If the listener has already been invoked after initialization was complete, it will already be removed internally
+     * to avoid leaking.
+     */
+    public fun removeOnInitializedListener(listener: () -> Unit): Unit = runBlocking {
+        initializedLock.withLock {
+            onInitializedListeners -= listener
+        }
+    }
+
     init {
         initialized = scope.launch {
             // Attempt to load the Identity from storage. If successful, we can notify any observers.
@@ -188,8 +206,7 @@ public class UID2Manager internal constructor(
                 validateAndSetIdentity(it.first, it.second, false)
             }
 
-            // If we have a callback provided, invoke it.
-            onInitialized?.invoke()
+            onInitialized()
         }
     }
 
@@ -309,6 +326,19 @@ public class UID2Manager internal constructor(
         }
     }
 
+    /**
+     * After initialization is complete, all the attached listeners will be invoked.
+     */
+    private suspend fun onInitialized() {
+        initializedLock.withLock {
+            onInitializedListeners.forEach { listener ->
+                listener()
+            }
+
+            onInitializedListeners.clear()
+        }
+    }
+
     private fun refreshIdentityInternal(identity: UID2Identity) = scope.launch {
         try {
             refreshToken(identity).retryWhen { _, attempt ->

sdk/src/test/java/com/uid2/UID2ManagerTest.kt

@@ -80,23 +80,30 @@ class UID2ManagerTest {
 
     @Test
     fun `reports when initialization is complete`() = runTest(testDispatcher) {
-        var isInitialized = false
-        val onInitialized = { isInitialized = true }
+        val listenerCount = 5
+        val listenerResults = ArrayList<Boolean>().apply {
+            for (i in 0..listenerCount) {
+                add(false)
+            }
+        }
 
         val manager = UID2Manager(client, storageManager, timeUtils, inputUtils, testDispatcher, false, logger).apply {
             this.checkExpiration = false
-            this.onInitialized = onInitialized
+
+            // Add the required listeners.
+            for (i in 0..listenerCount) {
+                addOnInitializedListener { listenerResults[i] = true }
+            }
         }
 
         // Verify that the manager invokes our callback after it's been able to load the identity from storage.
-        assertFalse(isInitialized)
+        assertTrue(listenerResults.all { !it })
         testDispatcher.scheduler.advanceUntilIdle()
-        assertTrue(isInitialized)
+        assertTrue(listenerResults.all { it })
 
-        // Reset our state and re-assign the manager's callback. Verify that even though initialization is complete, our
-        // callback is invoked immediately.
-        isInitialized = false
-        manager.onInitialized = onInitialized
+        // Create a new listener that we add after initialization is complete. Verify that it's called immediately.
+        var isInitialized = false
+        manager.addOnInitializedListener { isInitialized = true }
         assertTrue(isInitialized)
     }
 

securesignals-gma/src/main/java/com/uid2/securesignals/gma/UID2MediationAdapter.kt

@@ -48,7 +48,7 @@ public class UID2MediationAdapter : RtbAdapter() {
         // After we've asked to initialize the manager, we should wait until it's complete before reporting success.
         // This will potentially allow any previously persisted identity to be fully restored before we allow any
         // signals to be collected.
-        UID2Manager.getInstance().onInitialized = initializationCompleteCallback::onInitializationSucceeded
+        UID2Manager.getInstance().addOnInitializedListener(initializationCompleteCallback::onInitializationSucceeded)
     }
 
     /**

securesignals-ima/src/main/java/com/uid2/securesignals/ima/UID2SecureSignalsAdapter.kt

@@ -44,7 +44,7 @@ public class UID2SecureSignalsAdapter : SecureSignalsAdapter {
         // After we've asked to initialize the manager, we should wait until it's complete before reporting success.
         // This will potentially allow any previously persisted identity to be fully restored before we allow any
         // signals to be collected.
-        UID2Manager.getInstance().onInitialized = callback::onSuccess
+        UID2Manager.getInstance().addOnInitializedListener(callback::onSuccess)
     }
 
     /**