Merge remote-tracking branch 'origin/main' into MutableStateFlowUseUpdateInsteadOfCompareAndSet

Author: dconeybe

encoders/protoc-gen-firebase-encoders/src/main/kotlin/com/google/firebase/encoders/proto/codegen/Types.kt

@@ -190,7 +190,7 @@ data class ProtoField(
 
   val lowerCamelCaseName: String
     get() {
-      return SNAKE_CASE_REGEX.replace(name) { it.value.replace("_", "").toUpperCase() }
+      return SNAKE_CASE_REGEX.replace(name) { it.value.replace("_", "").uppercase() }
     }
 
   val camelCaseName: String

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

@@ -17,6 +17,7 @@
 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
@@ -127,6 +128,7 @@ class AuthIntegrationTest : DataConnectIntegrationTestBase() {
       grpcServer.metadatas.map { it.get(firebaseAuthTokenHeader) }.toCollection(authTokens)
     }
     val dataConnect = dataConnectFactory.newInstance(auth.app, grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val operationName = Arb.dataConnect.operationName().next(rs)
     val queryRef =
       dataConnect.query(operationName, Unit, serializer<TestData>(), serializer<Unit>())
@@ -155,6 +157,7 @@ class AuthIntegrationTest : DataConnectIntegrationTestBase() {
       grpcServer.metadatas.map { it.get(firebaseAuthTokenHeader) }.toCollection(authTokens)
     }
     val dataConnect = dataConnectFactory.newInstance(auth.app, grpcServer)
+    (dataConnect as FirebaseDataConnectInternal).awaitAuthReady()
     val operationName = Arb.dataConnect.operationName().next(rs)
     val mutationRef =
       dataConnect.mutation(operationName, Unit, serializer<TestData>(), serializer<Unit>())

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

@@ -48,8 +48,6 @@ 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. */
@@ -151,9 +149,18 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
    */
   fun close() {
     logger.debug { "close()" }
+
     weakThis.clear()
     coroutineScope.cancel()
-    setClosedState()
+
+    val oldState = state.getAndUpdate { State.Closed }
+    when (oldState) {
+      is State.Closed -> {}
+      is State.New -> {}
+      is State.StateWithProvider -> {
+        removeTokenListener(oldState.provider)
+      }
+    }
   }
 
   /**
@@ -178,37 +185,31 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     logger.debug { "awaitTokenProvider() done: currentState=$currentState" }
   }
 
-  // This function must ONLY be called from close().
-  private fun setClosedState() {
-    val oldState = state.getAndUpdate { State.Closed }
-    if (oldState is State.StateWithProvider<T>) {
-      removeTokenListener(oldState.provider)
-    }
-  }
-
   /**
    * Sets a flag to force-refresh the token upon the next call to [getToken].
    *
    * If [close] has been called, this method does nothing.
    */
   fun forceRefresh() {
     logger.debug { "forceRefresh()" }
-    val newState =
-      state.updateAndGet { oldState ->
+    val oldState =
+      state.getAndUpdate { oldState ->
         when (oldState) {
-          is State.Closed -> return
+          is State.Closed -> State.Closed
           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))
-            State.Idle(oldState.provider, forceTokenRefresh = true)
-          }
+          is State.Active -> State.Idle(oldState.provider, forceTokenRefresh = true)
         }
       }
 
-    check(newState is State.StateWithForceTokenRefresh<T> && newState.forceTokenRefresh) {
-      "newState.forceTokenRefresh should be true: $newState (error code gnvr2wx7nz)"
+    when (oldState) {
+      is State.Closed -> {}
+      is State.New -> {}
+      is State.Idle -> {}
+      is State.Active -> {
+        val message = "needs token refresh (wgrwbrvjxt)"
+        oldState.job.cancel(message, ForceRefresh(message))
+      }
     }
   }
 
@@ -338,24 +339,30 @@ internal sealed class DataConnectCredentialsTokenManager<T : Any>(
     logger.debug { "onProviderAvailable(newProvider=$newProvider)" }
     addTokenListener(newProvider)
 
-    state.update { oldState ->
-      when (oldState) {
-        is State.Closed -> {
-          logger.debug {
-            "onProviderAvailable(newProvider=$newProvider)" +
-              " unregistering token listener that was just added"
-          }
-          removeTokenListener(newProvider)
-          oldState
+    val oldState =
+      state.getAndUpdate { oldState ->
+        when (oldState) {
+          is State.Closed -> State.Closed
+          is State.New -> State.Idle(newProvider, oldState.forceTokenRefresh)
+          is State.Idle -> State.Idle(newProvider, oldState.forceTokenRefresh)
+          is State.Active -> State.Idle(newProvider, forceTokenRefresh = false)
         }
-        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)
+      }
+
+    when (oldState) {
+      is State.Closed -> {
+        logger.debug {
+          "onProviderAvailable(newProvider=$newProvider)" +
+            " unregistering token listener that was just added"
         }
+        removeTokenListener(newProvider)
+      }
+      is State.New -> {}
+      is State.Idle -> {}
+      is State.Active -> {
+        val newProviderClassName = newProvider::class.qualifiedName
+        val message = "a new provider $newProviderClassName is available (symhxtmazy)"
+        oldState.job.cancel(message, NewProvider(message))
       }
     }
   }

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

@@ -407,38 +407,41 @@ internal class FirebaseDataConnectImpl(
     dataConnectAuth.close()
     dataConnectAppCheck.close()
 
-    // Start the job to asynchronously close the gRPC client.
-    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()
-          }
+    // Create the "close job" to asynchronously close the gRPC client.
+    @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" }
-          }
+    // 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
+        } else {
+          NullableReference(newCloseJob)
         }
-
-        NullableReference(newCloseJob)
       }
 
-    val newCloseJob =
-      checkNotNull(newCloseJobRef.ref) {
-        "newCloseJobRef.ref should not be null (error code j3gbhd6e4j)"
-      }
-    newCloseJob.start()
-    return newCloseJob
+    // If the update "close job" was the one that we created, then start it!
+    if (updatedCloseJob.ref === newCloseJob) {
+      newCloseJob.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)"
+    }
   }
 
   // The generated SDK relies on equals() and hashCode() using object identity.

firebase-vertexai/src/main/kotlin/com/google/firebase/vertexai/type/LiveSession.kt

@@ -24,6 +24,7 @@ import io.ktor.client.plugins.websocket.ClientWebSocketSession
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
+import java.io.ByteArrayOutputStream
 import java.util.concurrent.ConcurrentLinkedQueue
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
@@ -141,16 +142,14 @@ internal constructor(
   }
 
   private suspend fun sendAudioDataToServer() {
-    var offset = 0
-    val audioBuffer = ByteArray(MIN_BUFFER_SIZE * 2)
+
+    val audioBufferStream = ByteArrayOutputStream()
     while (isRecording) {
       val receivedAudio = audioQueue.poll() ?: continue
-      receivedAudio.copyInto(audioBuffer, offset)
-      offset += receivedAudio.size
-      if (offset >= MIN_BUFFER_SIZE) {
-        sendMediaStream(listOf(MediaData(audioBuffer, "audio/pcm")))
-        audioBuffer.fill(0)
-        offset = 0
+      audioBufferStream.write(receivedAudio)
+      if (audioBufferStream.size() >= MIN_BUFFER_SIZE) {
+        sendMediaStream(listOf(MediaData(audioBufferStream.toByteArray(), "audio/pcm")))
+        audioBufferStream.reset()
       }
     }
   }