From e0feba939275c0ac3a388d22a648647d2b419827 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 06:34:08 +0000
Subject: [PATCH 01/10] DataConnectCredentialsTokenManager.kt: use
 MutableStateFlow instead of AtomicReference

---
 .../core/DataConnectCredentialsTokenManager.kt        | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index e6f4049c359..cf127d571e7 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
@@ -30,7 +30,6 @@ import com.google.firebase.inject.Provider
 import com.google.firebase.internal.api.FirebaseNoSignedInUserException
 import com.google.firebase.util.nextAlphanumericString
 import java.lang.ref.WeakReference
-import java.util.concurrent.atomic.AtomicReference
 import kotlin.coroutines.coroutineContext
 import kotlin.random.Random
 import kotlinx.coroutines.CancellationException
@@ -129,7 +128,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
    * coroutines to run on the thread.
    */
   private val state =
-    AtomicReference<State<T>>(State.Idle(provider = null, forceTokenRefresh = false))
+    MutableStateFlow<State<T>>(State.Idle(provider = null, forceTokenRefresh = false))
 
   /**
    * Adds the token listener to the given provider.
@@ -171,7 +170,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
   // This function must ONLY be called from close().
   private fun setClosedState() {
     while (true) {
-      val oldState = state.get()
+      val oldState = state.value
       val providerProvider: ProviderProvider<T> =
         when (oldState) {
           is State.Closed -> return
@@ -194,7 +193,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
   suspend fun forceRefresh() {
     logger.debug { "forceRefresh()" }
     while (true) {
-      val oldState = state.get()
+      val oldState = state.value
       val oldStateProviderProvider =
         when (oldState) {
           is State.Closed -> return
@@ -246,7 +245,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     logger.debug { "$invocationId getToken(requestId=$requestId)" }
     while (true) {
       val attemptSequenceNumber = nextSequenceNumber()
-      val oldState = state.get()
+      val oldState = state.value
 
       val newState: State.Active<T> =
         when (oldState) {
@@ -342,7 +341,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     addTokenListener(newProvider)
 
     while (true) {
-      val oldState = state.get()
+      val oldState = state.value
       val newState =
         when (oldState) {
           is State.Closed -> {

From 2adbdfe7f2282e81cc302d64f17f89d1e92c950d Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 07:09:30 +0000
Subject: [PATCH 02/10] DataConnectCredentialsTokenManager.kt: clean up State
 sealed interface

---
 .../DataConnectCredentialsTokenManager.kt     | 84 +++++++++----------
 .../dataconnect/core/DataConnectGrpcClient.kt |  2 +-
 2 files changed, 39 insertions(+), 47 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index cf127d571e7..de9638d82c3 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
@@ -48,7 +48,6 @@ import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.yield
 
 /** Base class that shares logic for managing the Auth token and AppCheck token. */
 internal sealed class DataConnectCredentialsTokenManager<T : Any>(
@@ -86,49 +85,39 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     }
   }
 
-  private interface ProviderProvider<T> {
-    val provider: T?
-  }
-
   private sealed interface State<out T> {
 
     /** State indicating that [close] has been invoked. */
     object Closed : State<Nothing>
 
-    /** State indicating that there is no outstanding "get token" request. */
-    class Idle<T>(
-
-      /**
-       * The [InternalAuthProvider] or [InteropAppCheckTokenProvider]; may be null if the deferred
-       * has not yet given us a provider.
-       */
-      override val provider: T?,
-
+    sealed interface StateWithForceTokenRefresh<out T> : State<T> {
       /** The value to specify for `forceRefresh` on the next invocation of [getToken]. */
       val forceTokenRefresh: Boolean
-    ) : State<T>, ProviderProvider<T>
+    }
 
-    /** State indicating that there _is_ an outstanding "get token" request. */
-    class Active<T>(
+    /** State indicating that the token provider is not (yet?) available. */
+    data class New(override val forceTokenRefresh: Boolean) : StateWithForceTokenRefresh<Nothing>
+
+    sealed interface StateWithProvider<out T> : State<T> {
+      /** The token provider, [InternalAuthProvider] or [InteropAppCheckTokenProvider] */
+      val provider: T
+    }
+
+    /** State indicating that there is no outstanding "get token" request. */
+    data class Idle<T>(override val provider: T, override val forceTokenRefresh: Boolean) :
+      StateWithProvider<T>, StateWithForceTokenRefresh<T>
 
-      /**
-       * The [InternalAuthProvider] or [InteropAppCheckTokenProvider] that is performing the "get
-       * token" request.
-       */
+    /** State indicating that there _is_ an outstanding "get token" request. */
+    data class Active<out T>(
       override val provider: T,
 
       /** The job that is performing the "get token" request. */
       val job: Deferred<SequencedReference<Result<GetTokenResult>>>
-    ) : State<T>, ProviderProvider<T>
+    ) : StateWithProvider<T>
   }
 
-  /**
-   * The current state of this object. The value should only be changed in a compare-and-swap loop
-   * in order to be thread-safe. Such a loop should call `yield()` on each iteration to allow other
-   * coroutines to run on the thread.
-   */
-  private val state =
-    MutableStateFlow<State<T>>(State.Idle(provider = null, forceTokenRefresh = false))
+  /** The current state of this object. */
+  private val state = MutableStateFlow<State<T>>(State.New(forceTokenRefresh = false))
 
   /**
    * Adds the token listener to the given provider.
@@ -171,15 +160,16 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
   private fun setClosedState() {
     while (true) {
       val oldState = state.value
-      val providerProvider: ProviderProvider<T> =
+      val provider: T? =
         when (oldState) {
           is State.Closed -> return
-          is State.Idle -> oldState
-          is State.Active -> oldState
+          is State.New -> null
+          is State.Idle -> oldState.provider
+          is State.Active -> oldState.provider
         }
 
       if (state.compareAndSet(oldState, State.Closed)) {
-        providerProvider.provider?.let { removeTokenListener(it) }
+        provider?.let { removeTokenListener(it) }
         break
       }
     }
@@ -190,27 +180,28 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
    *
    * If [close] has been called, this method does nothing.
    */
-  suspend fun forceRefresh() {
+  fun forceRefresh() {
     logger.debug { "forceRefresh()" }
     while (true) {
       val oldState = state.value
-      val oldStateProviderProvider =
+      val newState: State.StateWithForceTokenRefresh<T> =
         when (oldState) {
           is State.Closed -> return
-          is State.Idle -> oldState
+          is State.New -> oldState.copy(forceTokenRefresh = true)
+          is State.Idle -> oldState.copy(forceTokenRefresh = true)
           is State.Active -> {
             val message = "needs token refresh (wgrwbrvjxt)"
             oldState.job.cancel(message, ForceRefresh(message))
-            oldState
+            State.Idle(oldState.provider, forceTokenRefresh = true)
           }
         }
 
-      val newState = State.Idle(oldStateProviderProvider.provider, forceTokenRefresh = true)
+      check(newState.forceTokenRefresh) {
+        "newState.forceTokenRefresh should be true (error code gnvr2wx7nz)"
+      }
       if (state.compareAndSet(oldState, newState)) {
         break
       }
-
-      yield()
     }
   }
 
@@ -256,13 +247,13 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
             }
             throw CredentialsTokenManagerClosedException(this)
           }
-          is State.Idle -> {
-            if (oldState.provider === null) {
-              logger.debug {
-                "$invocationId getToken() returns null (token provider is not (yet?) available)"
-              }
-              return null
+          is State.New -> {
+            logger.debug {
+              "$invocationId getToken() returns null (token provider is not (yet?) available)"
             }
+            return null
+          }
+          is State.Idle -> {
             newActiveState(invocationId, oldState.provider, oldState.forceTokenRefresh)
           }
           is State.Active -> {
@@ -352,6 +343,7 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
             removeTokenListener(newProvider)
             break
           }
+          is State.New -> State.Idle(newProvider, oldState.forceTokenRefresh)
           is State.Idle -> State.Idle(newProvider, oldState.forceTokenRefresh)
           is State.Active -> {
             val newProviderClassName = newProvider::class.qualifiedName
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectGrpcClient.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectGrpcClient.kt
index 26e9ce49c51..d0ebe3f84c3 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectGrpcClient.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectGrpcClient.kt
@@ -102,7 +102,7 @@ internal class DataConnectGrpcClient(
     )
   }
 
-  private suspend inline fun <T, R> T.retryOnGrpcUnauthenticatedError(
+  private inline fun <T, R> T.retryOnGrpcUnauthenticatedError(
     requestId: String,
     kotlinMethodName: String,
     block: T.() -> R

From ee8fbf3642dd980859aefe07ed0fd3ed2376fddf Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 07:21:51 +0000
Subject: [PATCH 03/10] DataConnectCredentialsTokenManager.kt:
 awaitTokenProvider() added

---
 .../DataConnectCredentialsTokenManager.kt     | 31 ++++++++++++++-----
 .../core/FirebaseDataConnectImpl.kt           | 29 ++---------------
 2 files changed, 26 insertions(+), 34 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index de9638d82c3..9458a678bff 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/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
   }
 
   /**
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
index d9bee50b89d..f49afda964f 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
+++ b/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 =

From 68f3adab8f247e6e3654501825b62ef648d56963 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 07:30:51 +0000
Subject: [PATCH 04/10] FirebaseDataConnectInternalExts.kt: add extension
 functions: FirebaseDataConnect.awaitAuthReady() and
 FirebaseDataConnect.awaitAppCheckReady()

---
 .../dataconnect/AuthIntegrationTest.kt        |  4 +--
 .../GrpcMetadataIntegrationTest.kt            | 23 ++++++++--------
 .../FirebaseDataConnectInternalExts.kt        | 26 +++++++++++++++++++
 3 files changed, 40 insertions(+), 13 deletions(-)
 create mode 100644 firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/testutil/FirebaseDataConnectInternalExts.kt

diff --git a/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/AuthIntegrationTest.kt b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/AuthIntegrationTest.kt
index ba314d0f9db..5f240960b7e 100644
--- a/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/AuthIntegrationTest.kt
+++ b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/AuthIntegrationTest.kt
@@ -17,10 +17,10 @@
 package com.google.firebase.dataconnect
 
 import com.google.firebase.auth.FirebaseAuth
-import com.google.firebase.dataconnect.core.FirebaseDataConnectInternal
 import com.google.firebase.dataconnect.testutil.DataConnectBackend
 import com.google.firebase.dataconnect.testutil.DataConnectIntegrationTestBase
 import com.google.firebase.dataconnect.testutil.InProcessDataConnectGrpcServer
+import com.google.firebase.dataconnect.testutil.awaitAuthReady
 import com.google.firebase.dataconnect.testutil.newInstance
 import com.google.firebase.dataconnect.testutil.property.arbitrary.dataConnect
 import com.google.firebase.dataconnect.testutil.schemas.PersonSchema
@@ -202,7 +202,7 @@ class AuthIntegrationTest : DataConnectIntegrationTestBase() {
   }
 
   private suspend fun signIn() {
-    (personSchema.dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    personSchema.dataConnect.awaitAuthReady()
     val authResult = auth.run { signInAnonymously().await() }
     withClue("authResult.user returned from signInAnonymously()") {
       authResult.user.shouldNotBeNull()
diff --git a/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/GrpcMetadataIntegrationTest.kt b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/GrpcMetadataIntegrationTest.kt
index 8ff7715ddb9..394c2195332 100644
--- a/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/GrpcMetadataIntegrationTest.kt
+++ b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/GrpcMetadataIntegrationTest.kt
@@ -23,7 +23,6 @@ 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
@@ -32,6 +31,8 @@ import com.google.firebase.dataconnect.testutil.DataConnectIntegrationTestBase
 import com.google.firebase.dataconnect.testutil.DataConnectTestAppCheckToken
 import com.google.firebase.dataconnect.testutil.FirebaseAuthBackend
 import com.google.firebase.dataconnect.testutil.InProcessDataConnectGrpcServer
+import com.google.firebase.dataconnect.testutil.awaitAppCheckReady
+import com.google.firebase.dataconnect.testutil.awaitAuthReady
 import com.google.firebase.dataconnect.testutil.getFirebaseAppIdFromStrings
 import com.google.firebase.dataconnect.testutil.newInstance
 import com.google.firebase.dataconnect.util.SuspendingLazy
@@ -138,7 +139,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldNotSendAuthMetadataWhenNotLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.awaitAuthReady()
     val queryRef = dataConnect.query("qryfyk7yfppfe", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
 
@@ -151,7 +152,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldNotSendAuthMetadataWhenNotLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutckjpte9v9j", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -165,7 +166,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldSendAuthMetadataWhenLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.awaitAuthReady()
     val queryRef = dataConnect.query("qryyarwrxe2fv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
     firebaseAuthSignIn(dataConnect)
@@ -179,7 +180,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldSendAuthMetadataWhenLoggedIn() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutayn7as5k7d", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -194,7 +195,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldNotSendAuthMetadataAfterLogout() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.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() }
@@ -212,7 +213,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldNotSendAuthMetadataAfterLogout() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
+    dataConnect.awaitAuthReady()
     val mutationRef =
       dataConnect.mutation("mutvw945ag3vv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob1 = async { grpcServer.metadatas.first() }
@@ -233,7 +234,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
     // appcheck token is sent at all.
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
+    dataConnect.awaitAppCheckReady()
     val queryRef = dataConnect.query("qrybbeekpkkck", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
 
@@ -248,7 +249,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
     // appcheck token is sent at all.
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
+    dataConnect.awaitAppCheckReady()
     val mutationRef =
       dataConnect.mutation("mutbs7hhxk39c", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
@@ -262,7 +263,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeQueryShouldSendAppCheckMetadataWhenAppCheckIsEnabled() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
+    dataConnect.awaitAppCheckReady()
     val queryRef = dataConnect.query("qryyarwrxe2fv", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
     val appCheck = FirebaseAppCheck.getInstance(dataConnect.app)
@@ -277,7 +278,7 @@ class GrpcMetadataIntegrationTest : DataConnectIntegrationTestBase() {
   fun executeMutationShouldSendAppCheckMetadataWhenAppCheckIsEnabled() = runTest {
     val grpcServer = inProcessDataConnectGrpcServer.newInstance()
     val dataConnect = dataConnectFactory.newInstance(grpcServer)
-    (dataConnect as FirebaseDataConnectInternal).awaitAppCheckReady()
+    dataConnect.awaitAppCheckReady()
     val mutationRef =
       dataConnect.mutation("mutz4hzqzpgb4", Unit, serializer<Unit>(), serializer<Unit>())
     val metadatasJob = async { grpcServer.metadatas.first() }
diff --git a/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/testutil/FirebaseDataConnectInternalExts.kt b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/testutil/FirebaseDataConnectInternalExts.kt
new file mode 100644
index 00000000000..f73f1d8b50d
--- /dev/null
+++ b/firebase-dataconnect/src/androidTest/kotlin/com/google/firebase/dataconnect/testutil/FirebaseDataConnectInternalExts.kt
@@ -0,0 +1,26 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.firebase.dataconnect.testutil
+
+import com.google.firebase.dataconnect.FirebaseDataConnect
+import com.google.firebase.dataconnect.core.FirebaseDataConnectInternal
+
+suspend fun FirebaseDataConnect.awaitAuthReady() =
+  (this as FirebaseDataConnectInternal).awaitAuthReady()
+
+suspend fun FirebaseDataConnect.awaitAppCheckReady() =
+  (this as FirebaseDataConnectInternal).awaitAppCheckReady()

From 6010e498d6460c2f31d917ce9931aac6413c9abc Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 12:42:23 -0400
Subject: [PATCH 05/10] empty commit to re-trigger github actions


From 415215c1fa6fe4d91f356b00bf346e6d346c0e58 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Fri, 4 Apr 2025 19:53:29 +0000
Subject: [PATCH 06/10] dataconnect: Use MutableStateFlow.update() instead of
 MutableStateFlow.compareAndSet() directly for improved readability and less
 potential for bugs

---
 .../DataConnectCredentialsTokenManager.kt     | 72 +++++++------------
 .../core/FirebaseDataConnectImpl.kt           | 46 ++++++------
 .../dataconnect/core/QuerySubscriptionImpl.kt | 24 +++----
 .../querymgr/RegisteredDataDeserialzer.kt     | 12 ++--
 .../testutil/SuspendingCountDownLatch.kt      | 11 ++-
 5 files changed, 71 insertions(+), 94 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index 9458a678bff..df99cca553a 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
@@ -47,6 +47,9 @@ import kotlinx.coroutines.ensureActive
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.getAndUpdate
+import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.flow.updateAndGet
 import kotlinx.coroutines.launch
 
 /** Base class that shares logic for managing the Auth token and AppCheck token. */
@@ -177,20 +180,9 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
 
   // This function must ONLY be called from close().
   private fun setClosedState() {
-    while (true) {
-      val oldState = state.value
-      val provider: T? =
-        when (oldState) {
-          is State.Closed -> return
-          is State.New -> null
-          is State.Idle -> oldState.provider
-          is State.Active -> oldState.provider
-        }
-
-      if (state.compareAndSet(oldState, State.Closed)) {
-        provider?.let { removeTokenListener(it) }
-        break
-      }
+    val oldState = state.getAndUpdate { State.Closed }
+    if (oldState is State.StateWithProvider<T>) {
+      removeTokenListener(oldState.provider)
     }
   }
 
@@ -201,9 +193,8 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
    */
   fun forceRefresh() {
     logger.debug { "forceRefresh()" }
-    while (true) {
-      val oldState = state.value
-      val newState: State.StateWithForceTokenRefresh<T> =
+    val newState =
+      state.updateAndGet { oldState ->
         when (oldState) {
           is State.Closed -> return
           is State.New -> oldState.copy(forceTokenRefresh = true)
@@ -214,13 +205,10 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
             State.Idle(oldState.provider, forceTokenRefresh = true)
           }
         }
-
-      check(newState.forceTokenRefresh) {
-        "newState.forceTokenRefresh should be true (error code gnvr2wx7nz)"
-      }
-      if (state.compareAndSet(oldState, newState)) {
-        break
       }
+
+    check(newState is State.StateWithForceTokenRefresh<T> && newState.forceTokenRefresh) {
+      "newState.forceTokenRefresh should be true: $newState (error code gnvr2wx7nz)"
     }
   }
 
@@ -350,30 +338,24 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     logger.debug { "onProviderAvailable(newProvider=$newProvider)" }
     addTokenListener(newProvider)
 
-    while (true) {
-      val oldState = state.value
-      val newState =
-        when (oldState) {
-          is State.Closed -> {
-            logger.debug {
-              "onProviderAvailable(newProvider=$newProvider)" +
-                " unregistering token listener that was just added"
-            }
-            removeTokenListener(newProvider)
-            break
-          }
-          is State.New -> State.Idle(newProvider, oldState.forceTokenRefresh)
-          is State.Idle -> State.Idle(newProvider, oldState.forceTokenRefresh)
-          is State.Active -> {
-            val newProviderClassName = newProvider::class.qualifiedName
-            val message = "a new provider $newProviderClassName is available (symhxtmazy)"
-            oldState.job.cancel(message, NewProvider(message))
-            State.Idle(newProvider, forceTokenRefresh = false)
+    state.update { oldState ->
+      when (oldState) {
+        is State.Closed -> {
+          logger.debug {
+            "onProviderAvailable(newProvider=$newProvider)" +
+              " unregistering token listener that was just added"
           }
+          removeTokenListener(newProvider)
+          oldState
+        }
+        is State.New -> State.Idle(newProvider, oldState.forceTokenRefresh)
+        is State.Idle -> State.Idle(newProvider, oldState.forceTokenRefresh)
+        is State.Active -> {
+          val newProviderClassName = newProvider::class.qualifiedName
+          val message = "a new provider $newProviderClassName is available (symhxtmazy)"
+          oldState.job.cancel(message, NewProvider(message))
+          State.Idle(newProvider, forceTokenRefresh = false)
         }
-
-      if (state.compareAndSet(oldState, newState)) {
-        break
       }
     }
   }
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
index f49afda964f..3a7569967de 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
@@ -54,6 +54,7 @@ import kotlinx.coroutines.async
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.collect
+import kotlinx.coroutines.flow.updateAndGet
 import kotlinx.coroutines.runBlocking
 import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.sync.withLock
@@ -407,34 +408,37 @@ internal class FirebaseDataConnectImpl(
     dataConnectAppCheck.close()
 
     // Start the job to asynchronously close the gRPC client.
-    while (true) {
-      val oldCloseJob = closeJob.value
-
-      oldCloseJob.ref?.let {
-        if (!it.isCancelled) {
-          return it
+    val newCloseJobRef =
+      closeJob.updateAndGet { oldCloseJob ->
+        oldCloseJob.ref?.let {
+          if (!it.isCancelled) {
+            return it
+          }
         }
-      }
 
-      @OptIn(DelicateCoroutinesApi::class)
-      val newCloseJob =
-        GlobalScope.async<Unit>(start = CoroutineStart.LAZY) {
-          lazyGrpcRPCs.initializedValueOrNull?.close()
-        }
+        @OptIn(DelicateCoroutinesApi::class)
+        val newCloseJob =
+          GlobalScope.async<Unit>(start = CoroutineStart.LAZY) {
+            lazyGrpcRPCs.initializedValueOrNull?.close()
+          }
 
-      newCloseJob.invokeOnCompletion { exception ->
-        if (exception === null) {
-          logger.debug { "close() completed successfully" }
-        } else {
-          logger.warn(exception) { "close() failed" }
+        newCloseJob.invokeOnCompletion { exception ->
+          if (exception === null) {
+            logger.debug { "close() completed successfully" }
+          } else {
+            logger.warn(exception) { "close() failed" }
+          }
         }
+
+        NullableReference(newCloseJob)
       }
 
-      if (closeJob.compareAndSet(oldCloseJob, NullableReference(newCloseJob))) {
-        newCloseJob.start()
-        return newCloseJob
+    val newCloseJob =
+      checkNotNull(newCloseJobRef.ref) {
+        "newCloseJobRef.ref should not be null (error code j3gbhd6e4j)"
       }
-    }
+    newCloseJob.start()
+    return newCloseJob
   }
 
   // The generated SDK relies on equals() and hashCode() using object identity.
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/QuerySubscriptionImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/QuerySubscriptionImpl.kt
index ceeb861cab8..2ca9aea6771 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/QuerySubscriptionImpl.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/QuerySubscriptionImpl.kt
@@ -27,6 +27,7 @@ import kotlinx.coroutines.cancelAndJoin
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.channelFlow
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 
 internal class QuerySubscriptionImpl<Data, Variables>(query: QueryRefImpl<Data, Variables>) :
@@ -80,22 +81,17 @@ internal class QuerySubscriptionImpl<Data, Variables>(query: QueryRefImpl<Data,
   }
 
   private fun updateLastResult(prospectiveLastResult: QuerySubscriptionResultImpl) {
-    // Update the last result in a compare-and-swap loop so that there is no possibility of
-    // clobbering a newer result with an older result, compared using their sequence numbers.
     // TODO: Fix this so that results from an old query do not clobber results from a new query,
     //  as set by a call to update()
-    while (true) {
-      val currentLastResult = _lastResult.value
-      if (currentLastResult.ref != null) {
-        val currentSequenceNumber = currentLastResult.ref.sequencedResult.sequenceNumber
-        val prospectiveSequenceNumber = prospectiveLastResult.sequencedResult.sequenceNumber
-        if (currentSequenceNumber >= prospectiveSequenceNumber) {
-          return
-        }
-      }
-
-      if (_lastResult.compareAndSet(currentLastResult, NullableReference(prospectiveLastResult))) {
-        return
+    _lastResult.update { currentLastResult ->
+      if (
+        currentLastResult.ref != null &&
+          currentLastResult.ref.sequencedResult.sequenceNumber >=
+            prospectiveLastResult.sequencedResult.sequenceNumber
+      ) {
+        currentLastResult
+      } else {
+        NullableReference(prospectiveLastResult)
       }
     }
   }
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/querymgr/RegisteredDataDeserialzer.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/querymgr/RegisteredDataDeserialzer.kt
index 3f94a7f95a0..1fa6d94eae4 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/querymgr/RegisteredDataDeserialzer.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/querymgr/RegisteredDataDeserialzer.kt
@@ -31,6 +31,7 @@ import kotlinx.coroutines.channels.BufferOverflow
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.onSubscription
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.withContext
 import kotlinx.serialization.DeserializationStrategy
 import kotlinx.serialization.modules.SerializersModule
@@ -84,17 +85,14 @@ internal class RegisteredDataDeserializer<T>(
         lazyDeserialize(requestId, sequencedResult)
       )
 
-    // Use a compare-and-swap ("CAS") loop to ensure that an old update never clobbers a newer one.
-    while (true) {
-      val currentUpdate = latestUpdate.value
+    latestUpdate.update { currentUpdate ->
       if (
         currentUpdate.ref !== null &&
           currentUpdate.ref.sequenceNumber > sequencedResult.sequenceNumber
       ) {
-        break // don't clobber a newer update with an older one
-      }
-      if (latestUpdate.compareAndSet(currentUpdate, NullableReference(newUpdate))) {
-        break
+        currentUpdate // don't clobber a newer update with an older one
+      } else {
+        NullableReference(newUpdate)
       }
     }
 
diff --git a/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt b/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
index 7098a390886..a4ff85d04f0 100644
--- a/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
+++ b/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
@@ -19,6 +19,7 @@ package com.google.firebase.dataconnect.testutil
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.update
 
 /**
  * An implementation of [java.util.concurrent.CountDownLatch] that suspends instead of blocking.
@@ -60,14 +61,10 @@ class SuspendingCountDownLatch(count: Int) {
    * @throws IllegalStateException if called when the count has already reached zero.
    */
   fun countDown(): SuspendingCountDownLatch {
-    while (true) {
-      val oldValue = _count.value
+    _count.update { oldValue ->
       check(oldValue > 0) { "countDown() called too many times (oldValue=$oldValue)" }
-
-      val newValue = oldValue - 1
-      if (_count.compareAndSet(oldValue, newValue)) {
-        return this
-      }
+      oldValue - 1
     }
+    return this
   }
 }

From c9d70a3e1be9c149dbe56c6dba2c90c5d26f6bdf Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Wed, 9 Apr 2025 19:52:33 +0000
Subject: [PATCH 07/10] FirebaseDataConnectImpl.kt: fix minor typo in comment

---
 .../google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
index 38cdcdcf12c..c96c990b7c0 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
@@ -432,7 +432,7 @@ internal class FirebaseDataConnectImpl(
         }
       }
 
-    // If the update "close job" was the one that we created, then start it!
+    // If the updated "close job" was the one that we created, then start it!
     if (updatedCloseJob.ref === newCloseJob) {
       newCloseJob.start()
     }

From 3aef735192af10bcb52490f9722e9c02e55fbb04 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Tue, 15 Apr 2025 14:06:08 -0400
Subject: [PATCH 08/10] minor cleanup

---
 .../core/FirebaseDataConnectImpl.kt           | 31 ++++++++++---------
 1 file changed, 16 insertions(+), 15 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
index c96c990b7c0..d05ca685db2 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
@@ -421,27 +421,28 @@ internal class FirebaseDataConnectImpl(
       }
     }
 
-    // Register the new "close job", unless there is a "close job" already in progress or one that
-    // completed successfully.
-    val updatedCloseJob =
-      closeJob.updateAndGet { oldCloseJob ->
-        if (oldCloseJob.ref !== null && !oldCloseJob.ref.isCancelled) {
-          oldCloseJob
+    // Register the new "close job". Do not overwrite a close job that is already in progress (to
+    // avoid having more than one close job in progress at a time) or a close job that completed
+    // successfully (since there is nothing to do if a previous close job was successful).
+    val updatedCloseJobRef =
+      closeJob.updateAndGet { oldCloseJobRef: NullableReference<Deferred<Unit>> ->
+        if (oldCloseJobRef.ref !== null && !oldCloseJobRef.ref.isCancelled) {
+          oldCloseJobRef
         } else {
           NullableReference(newCloseJob)
         }
       }
 
-    // If the updated "close job" was the one that we created, then start it!
-    if (updatedCloseJob.ref === newCloseJob) {
-      newCloseJob.start()
-    }
+    // Start the updated "close job" (if it was already started then start() is a no-op).
+    val updatedCloseJob =
+      checkNotNull(updatedCloseJobRef.ref) {
+        "internal error: closeJob.updateAndGet() returned a NullableReference whose 'ref' " +
+          "property was null; however it should NOT have been null (error code y5fk4ntdnd)"
+      }
+    updatedCloseJob.start()
 
-    // Return the job "close job" that is active or already completed so that the caller can await
-    // its result.
-    return checkNotNull(updatedCloseJob.ref) {
-      "updatedCloseJob.ref should not have been null (error code y5fk4ntdnd)"
-    }
+    // Return the "close job", which _may_ already be completed, so the caller can await it.
+    return updatedCloseJob
   }
 
   // The generated SDK relies on equals() and hashCode() using object identity.

From 9f3b7942b05a51a73132c495f9c4a203d57eb5a0 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Wed, 16 Apr 2025 13:53:41 -0400
Subject: [PATCH 09/10] Rename "old" state variables to "current" state, as
 suggested by Copilot:
 https://github.com/firebase/firebase-android-sdk/pull/6840#pullrequestreview-2773350309

---
 .../core/DataConnectCredentialsTokenManager.kt | 18 +++++++++---------
 .../core/FirebaseDataConnectImpl.kt            |  6 +++---
 .../testutil/SuspendingCountDownLatch.kt       |  6 +++---
 3 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index 10d9a1ce4f9..15b3a729821 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
@@ -193,12 +193,12 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
   fun forceRefresh() {
     logger.debug { "forceRefresh()" }
     val oldState =
-      state.getAndUpdate { oldState ->
-        when (oldState) {
+      state.getAndUpdate { currentState ->
+        when (currentState) {
           is State.Closed -> State.Closed
-          is State.New -> oldState.copy(forceTokenRefresh = true)
-          is State.Idle -> oldState.copy(forceTokenRefresh = true)
-          is State.Active -> State.Idle(oldState.provider, forceTokenRefresh = true)
+          is State.New -> currentState.copy(forceTokenRefresh = true)
+          is State.Idle -> currentState.copy(forceTokenRefresh = true)
+          is State.Active -> State.Idle(currentState.provider, forceTokenRefresh = true)
         }
       }
 
@@ -340,11 +340,11 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     addTokenListener(newProvider)
 
     val oldState =
-      state.getAndUpdate { oldState ->
-        when (oldState) {
+      state.getAndUpdate { currentState ->
+        when (currentState) {
           is State.Closed -> State.Closed
-          is State.New -> State.Idle(newProvider, oldState.forceTokenRefresh)
-          is State.Idle -> State.Idle(newProvider, oldState.forceTokenRefresh)
+          is State.New -> State.Idle(newProvider, currentState.forceTokenRefresh)
+          is State.Idle -> State.Idle(newProvider, currentState.forceTokenRefresh)
           is State.Active -> State.Idle(newProvider, forceTokenRefresh = false)
         }
       }
diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
index d05ca685db2..9ad64cc6054 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/FirebaseDataConnectImpl.kt
@@ -425,9 +425,9 @@ internal class FirebaseDataConnectImpl(
     // avoid having more than one close job in progress at a time) or a close job that completed
     // successfully (since there is nothing to do if a previous close job was successful).
     val updatedCloseJobRef =
-      closeJob.updateAndGet { oldCloseJobRef: NullableReference<Deferred<Unit>> ->
-        if (oldCloseJobRef.ref !== null && !oldCloseJobRef.ref.isCancelled) {
-          oldCloseJobRef
+      closeJob.updateAndGet { currentCloseJobRef: NullableReference<Deferred<Unit>> ->
+        if (currentCloseJobRef.ref !== null && !currentCloseJobRef.ref.isCancelled) {
+          currentCloseJobRef
         } else {
           NullableReference(newCloseJob)
         }
diff --git a/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt b/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
index a4ff85d04f0..b7828e0c36a 100644
--- a/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
+++ b/firebase-dataconnect/testutil/src/main/kotlin/com/google/firebase/dataconnect/testutil/SuspendingCountDownLatch.kt
@@ -61,9 +61,9 @@ class SuspendingCountDownLatch(count: Int) {
    * @throws IllegalStateException if called when the count has already reached zero.
    */
   fun countDown(): SuspendingCountDownLatch {
-    _count.update { oldValue ->
-      check(oldValue > 0) { "countDown() called too many times (oldValue=$oldValue)" }
-      oldValue - 1
+    _count.update { currentValue ->
+      check(currentValue > 0) { "countDown() called too many times (currentValue=$currentValue)" }
+      currentValue - 1
     }
     return this
   }

From b1f01b0530ea88f93a4c351975422270485d9435 Mon Sep 17 00:00:00 2001
From: Denver Coneybeare <dconeybe@google.com>
Date: Wed, 16 Apr 2025 14:48:35 -0400
Subject: [PATCH 10/10] DataConnectCredentialsTokenManager.kt: verify the new
 state in forceRefresh() to avoid unintentional internal state corruption.

This was suggested by Copilot: https://github.com/firebase/firebase-android-sdk/pull/6840#pullrequestreview-2773477393
---
 .../DataConnectCredentialsTokenManager.kt     | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
index 15b3a729821..d96e544d6af 100644
--- a/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
+++ b/firebase-dataconnect/src/main/kotlin/com/google/firebase/dataconnect/core/DataConnectCredentialsTokenManager.kt
@@ -194,12 +194,23 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     logger.debug { "forceRefresh()" }
     val oldState =
       state.getAndUpdate { currentState ->
-        when (currentState) {
-          is State.Closed -> State.Closed
-          is State.New -> currentState.copy(forceTokenRefresh = true)
-          is State.Idle -> currentState.copy(forceTokenRefresh = true)
-          is State.Active -> State.Idle(currentState.provider, forceTokenRefresh = true)
+        val newState =
+          when (currentState) {
+            is State.Closed -> State.Closed
+            is State.New -> currentState.copy(forceTokenRefresh = true)
+            is State.Idle -> currentState.copy(forceTokenRefresh = true)
+            is State.Active -> State.Idle(currentState.provider, forceTokenRefresh = true)
+          }
+
+        check(newState is State.Closed || newState is State.StateWithForceTokenRefresh<T>) {
+          "internal error gbazc7qr66: newState should have been Closed or " +
+            "StateWithForceTokenRefresh, but got: $newState"
+        }
+        check((newState as? State.StateWithForceTokenRefresh<T>)?.forceTokenRefresh !== false) {
+          "internal error fnzwyrsez2: newState.forceTokenRefresh should have been true"
         }
+
+        newState
       }
 
     when (oldState) {