dataconnect: Use MutableStateFlow.update() instead of MutableStateFlow.compareAndSet() directly for improved readability and less potential for bugs

Author: dconeybe

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
       }
     }
   }

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.

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)
       }
     }
   }

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)
       }
     }
 

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
   }
 }