GrpcMetadataIntegrationTest.kt: Fix race condition waiting for auth/appcheck to be ready

Author: dconeybe

firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/GrpcMetadataIntegrationTest.kt

@@ -23,6 +23,7 @@ import com.google.android.gms.tasks.Tasks
 import com.google.firebase.appcheck.AppCheckProvider
 import com.google.firebase.appcheck.AppCheckProviderFactory
 import com.google.firebase.appcheck.FirebaseAppCheck
+import com.google.firebase.dataconnect.core.FirebaseDataConnectInternal
 import com.google.firebase.dataconnect.generated.GeneratedConnector
 import com.google.firebase.dataconnect.generated.GeneratedMutation
 import com.google.firebase.dataconnect.generated.GeneratedQuery
@@ -137,6 +138,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldNotSendAuthMetadataWhenNotLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val queryRef = dataConnect.query("qryfyk7yfppfe", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
 
@@ -149,6 +151,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldNotSendAuthMetadataWhenNotLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutckjpte9v9j", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -162,6 +165,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldSendAuthMetadataWhenLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val queryRef = dataConnect.query("qryyarwrxe2fv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
     firebaseAuthSignIn(dataConnect)
@@ -175,6 +179,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldSendAuthMetadataWhenLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutayn7as5k7d", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -189,6 +194,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldNotSendAuthMetadataAfterLogout() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val queryRef = dataConnect.query("qryyarwrxe2fv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob1 = async { grpcServer.metadatas.first() }
     val metadatasJob2 = async { grpcServer.metadatas.take(2).last() }
@@ -206,6 +212,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldNotSendAuthMetadataAfterLogout() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutvw945ag3vv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob1 = async { grpcServer.metadatas.first() }
@@ -226,6 +233,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
     // appcheck token is sent at all.
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
     val queryRef = dataConnect.query("qrybbeekpkkck", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
 
@@ -240,6 +248,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
     // appcheck token is sent at all.
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
     val mutationRef =
       dataConnect.mutation("mutbs7hhxk39c", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -253,6 +262,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldSendAppCheckMetadataWhenAppCheckIsEnabled() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
     val queryRef = dataConnect.query("qryyarwrxe2fv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
     val appCheck = FirebaseAppCheck.getInstance(dataConnect.app)
@@ -267,6 +277,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldSendAppCheckMetadataWhenAppCheckIsEnabled() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
     val mutationRef =
       dataConnect.mutation("mutz4hzqzpgb4", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }

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

@@ -45,6 +45,9 @@ import kotlinx.coroutines.SupervisorJob
 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.launch
 import kotlinx.coroutines.yield
 
@@ -58,6 +61,9 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, L : 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 =
@@ -448,6 +454,8 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any, L : Any>(
         break
       }
     }
+
+    _providerAvailable.value = true
   }
 
   /**

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

@@ -54,6 +54,8 @@ 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
@@ -72,6 +74,9 @@ internal interface FirebaseDataConnectInternal : FirebaseDataConnect {
 
   val lazyGrpcClient: SuspendingLazy<DataConnectGrpcClient>
   val lazyQueryManager: SuspendingLazy<QueryManager>
+
+  suspend fun awaitAuthReady()
+  suspend fun awaitAppCheckReady()
 }
 
 internal class FirebaseDataConnectImpl(
@@ -107,11 +112,18 @@ internal class FirebaseDataConnectImpl(
       SupervisorJob() +
         nonBlockingDispatcher +
         CoroutineName(instanceId) +
-        CoroutineExceptionHandler { _, throwable ->
-          logger.warn(throwable) { "uncaught exception from a coroutine" }
+        CoroutineExceptionHandler { context, throwable ->
+          logger.warn(throwable) {
+            val coroutineName = context[CoroutineName]?.name
+            "WARNING: uncaught exception from coroutine named \"$coroutineName\" " +
+              "(error code jszxcbe37k)"
+          }
         }
     )
 
+  private val authProviderAvailable = MutableStateFlow(false)
+  private val appCheckProviderAvailable = MutableStateFlow(false)
+
   // Protects `closed`, `grpcClient`, `emulatorSettings`, and `queryManager`.
   private val mutex = Mutex()
 
@@ -121,29 +133,49 @@ internal class FirebaseDataConnectImpl(
   // All accesses to this variable _must_ have locked `mutex`.
   private var closed = false
 
-  private val lazyDataConnectAuth =
-    SuspendingLazy(mutex) {
-      if (closed) throw IllegalStateException("FirebaseDataConnect instance has been closed")
-      DataConnectAuth(
-          deferredAuthProvider = deferredAuthProvider,
-          parentCoroutineScope = coroutineScope,
-          blockingDispatcher = blockingDispatcher,
-          logger = Logger("DataConnectAuth").apply { debug { "created by $instanceId" } },
-        )
-        .apply { initialize() }
+  private val dataConnectAuth: DataConnectAuth =
+    DataConnectAuth(
+      deferredAuthProvider = deferredAuthProvider,
+      parentCoroutineScope = coroutineScope,
+      blockingDispatcher = blockingDispatcher,
+      logger = Logger("DataConnectAuth").apply { debug { "created by $instanceId" } },
+    )
+
+  override suspend fun awaitAuthReady() {
+    authProviderAvailable.first { it }
+  }
+
+  init {
+    coroutineScope.launch(CoroutineName("DataConnectAuth initializer for $instanceId")) {
+      dataConnectAuth.initialize()
+      dataConnectAuth.providerAvailable.collect { isProviderAvailable ->
+        logger.debug { "authProviderAvailable=$isProviderAvailable" }
+        authProviderAvailable.value = isProviderAvailable
+      }
     }
+  }
 
-  private val lazyDataConnectAppCheck =
-    SuspendingLazy(mutex) {
-      if (closed) throw IllegalStateException("FirebaseDataConnect instance has been closed")
-      DataConnectAppCheck(
-          deferredAppCheckTokenProvider = deferredAppCheckProvider,
-          parentCoroutineScope = coroutineScope,
-          blockingDispatcher = blockingDispatcher,
-          logger = Logger("DataConnectAppCheck").apply { debug { "created by $instanceId" } },
-        )
-        .apply { initialize() }
+  private val dataConnectAppCheck: DataConnectAppCheck =
+    DataConnectAppCheck(
+      deferredAppCheckTokenProvider = deferredAppCheckProvider,
+      parentCoroutineScope = coroutineScope,
+      blockingDispatcher = blockingDispatcher,
+      logger = Logger("DataConnectAppCheck").apply { debug { "created by $instanceId" } },
+    )
+
+  override suspend fun awaitAppCheckReady() {
+    appCheckProviderAvailable.first { it }
+  }
+
+  init {
+    coroutineScope.launch(CoroutineName("DataConnectAppCheck initializer for $instanceId")) {
+      dataConnectAppCheck.initialize()
+      dataConnectAppCheck.providerAvailable.collect { isProviderAvailable ->
+        logger.debug { "appCheckProviderAvailable=$isProviderAvailable" }
+        appCheckProviderAvailable.value = isProviderAvailable
+      }
     }
+  }
 
   private val lazyGrpcRPCs =
     SuspendingLazy(mutex) {
@@ -181,8 +213,8 @@ internal class FirebaseDataConnectImpl(
       val grpcMetadata =
         DataConnectGrpcMetadata.forSystemVersions(
           firebaseApp = app,
-          dataConnectAuth = lazyDataConnectAuth.getLocked(),
-          dataConnectAppCheck = lazyDataConnectAppCheck.getLocked(),
+          dataConnectAuth = dataConnectAuth,
+          dataConnectAppCheck = dataConnectAppCheck,
           connectorLocation = config.location,
           parentLogger = logger,
         )
@@ -210,8 +242,8 @@ internal class FirebaseDataConnectImpl(
         projectId = projectId,
         connector = config,
         grpcRPCs = lazyGrpcRPCs.getLocked(),
-        dataConnectAuth = lazyDataConnectAuth.getLocked(),
-        dataConnectAppCheck = lazyDataConnectAppCheck.getLocked(),
+        dataConnectAuth = dataConnectAuth,
+        dataConnectAppCheck = dataConnectAppCheck,
         logger = Logger("DataConnectGrpcClient").apply { debug { "created by $instanceId" } },
       )
     }
@@ -397,8 +429,8 @@ internal class FirebaseDataConnectImpl(
 
     // Close Auth and AppCheck synchronously to avoid race conditions with auth callbacks.
     // Since close() is re-entrant, this is safe even if they have already been closed.
-    lazyDataConnectAuth.initializedValueOrNull?.close()
-    lazyDataConnectAppCheck.initializedValueOrNull?.close()
+    dataConnectAuth.close()
+    dataConnectAppCheck.close()
 
     // Start the job to asynchronously close the gRPC client.
     while (true) {