DataConnectCredentialsTokenManager.kt: awaitTokenProvider() added

Author: dconeybe

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt

@@ -45,8 +45,8 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.launch
 
 /** Base class that shares logic for managing the Auth token and AppCheck token. */
@@ -59,9 +59,6 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
   val instanceId: String
     get() = logger.nameWithId
 
-  private val _providerAvailable = MutableStateFlow(false)
-  val providerAvailable: StateFlow<Boolean> = _providerAvailable.asStateFlow()
-
   @Suppress("LeakingThis") private val weakThis = WeakReference(this)
 
   private val coroutineScope =
@@ -156,6 +153,28 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     setClosedState()
   }
 
+  /**
+   * Suspends until the token provider becomes available to this object.
+   *
+   * If [close] has been invoked, or is invoked _before_ a token provider becomes available, then
+   * this method returns normally, as if a token provider _had_ become available.
+   */
+  suspend fun awaitTokenProvider() {
+    logger.debug { "awaitTokenProvider() start" }
+    val currentState =
+      state
+        .filter {
+          when (it) {
+            State.Closed -> true
+            is State.New -> false
+            is State.Idle -> true
+            is State.Active -> true
+          }
+        }
+        .first()
+    logger.debug { "awaitTokenProvider() done: currentState=$currentState" }
+  }
+
   // This function must ONLY be called from close().
   private fun setClosedState() {
     while (true) {
@@ -357,8 +376,6 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
         break
       }
     }
-
-    _providerAvailable.value = true
   }
 
   /**

firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt

@@ -54,8 +54,6 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.collect
-import kotlinx.coroutines.flow.first
-import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -120,9 +118,6 @@ internal class FirebaseDataConnectImpl(
         }
     )
 
-  private val authProviderAvailable = MutableStateFlow(false)
-  private val appCheckProviderAvailable = MutableStateFlow(false)
-
   // Protects `closed`, `grpcClient`, `emulatorSettings`, and `queryManager`.
   private val mutex = Mutex()
 
@@ -141,17 +136,7 @@ internal class FirebaseDataConnectImpl(
     )
 
   override suspend fun awaitAuthReady() {
-    authProviderAvailable.first { it }
-  }
-
-  init {
-    val name = CoroutineName("DataConnectAuth isProviderAvailable pipe for $instanceId")
-    coroutineScope.launch(name) {
-      dataConnectAuth.providerAvailable.collect { isProviderAvailable ->
-        logger.debug { "authProviderAvailable=$isProviderAvailable" }
-        authProviderAvailable.value = isProviderAvailable
-      }
-    }
+    dataConnectAuth.awaitTokenProvider()
   }
 
   private val dataConnectAppCheck: DataConnectAppCheck =
@@ -163,17 +148,7 @@ internal class FirebaseDataConnectImpl(
     )
 
   override suspend fun awaitAppCheckReady() {
-    appCheckProviderAvailable.first { it }
-  }
-
-  init {
-    val name = CoroutineName("DataConnectAppCheck isProviderAvailable pipe for $instanceId")
-    coroutineScope.launch(name) {
-      dataConnectAppCheck.providerAvailable.collect { isProviderAvailable ->
-        logger.debug { "appCheckProviderAvailable=$isProviderAvailable" }
-        appCheckProviderAvailable.value = isProviderAvailable
-      }
-    }
+    dataConnectAppCheck.awaitTokenProvider()
   }
 
   private val lazyGrpcRPCs =