fix(ai): Fix concurrency issues (WIP)

firebase-ai/src/main/kotlin/com/google/firebase/ai/common/util/android.kt

@@ -19,7 +19,9 @@ package com.google.firebase.ai.common.util
 import android.media.AudioRecord
 import kotlin.time.Duration.Companion.milliseconds
 import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.flow.flow
+import kotlinx.coroutines.isActive
 import kotlinx.coroutines.yield
 
 /**
@@ -35,20 +37,18 @@ internal val AudioRecord.minBufferSize: Int
  *
  * Will yield when this instance is not recording.
  */
-internal fun AudioRecord.readAsFlow() = flow {
+internal fun AudioRecord.readAsFlow() = callbackFlow {
  val buffer = ByteArray(minBufferSize)
 
- while (true) {
+ while (isActive) {
  if (recordingState != AudioRecord.RECORDSTATE_RECORDING) {
- // TODO(vguthal): Investigate if both yield and delay are required.
- delay(10.milliseconds)
- yield()
+ delay(0)
  continue
  }
  val bytesRead = read(buffer, 0, buffer.size)
  if (bytesRead > 0) {
- emit(buffer.copyOf(bytesRead))
+ send(buffer.copyOf(bytesRead))
  }
- yield()
+ delay(0)
  }
 }

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

@@ -33,13 +33,15 @@ import io.ktor.client.plugins.websocket.DefaultClientWebSocketSession
 import io.ktor.websocket.Frame
 import io.ktor.websocket.close
 import io.ktor.websocket.readBytes
+import kotlinx.coroutines.CoroutineName
 import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.coroutines.CoroutineContext
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.cancel
 import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
+import kotlinx.coroutines.delay
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.buffer
 import kotlinx.coroutines.flow.catch
@@ -49,7 +51,6 @@ import kotlinx.coroutines.flow.onCompletion
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.isActive
 import kotlinx.coroutines.launch
-import kotlinx.coroutines.yield
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.encodeToString
@@ -120,7 +121,6 @@ internal constructor(
  functionCallHandler: ((FunctionCallPart) -> FunctionResponsePart)? = null,
  enableInterruptions: Boolean = false,
  ) {
-
  val context = firebaseApp.applicationContext
  if (
  ContextCompat.checkSelfPermission(context, RECORD_AUDIO) != PackageManager.PERMISSION_GRANTED
@@ -137,8 +137,8 @@ internal constructor(
  )
  return@catchAsync
  }
-
- scope = CoroutineScope(blockingDispatcher + childJob())
+ // TODO: maybe it should be THREAD_PRIORITY_AUDIO anyways for playback and recording (not network though)
+ scope = CoroutineScope(blockingDispatcher + childJob() + CoroutineName("LiveSession Scope"))
  audioHelper = AudioHelper.build()
 
  recordUserAudio()
@@ -201,7 +201,7 @@ internal constructor(
  )
  }
  ?.let { emit(it.toPublic()) }
- yield()
+ delay(0)
  }
  }
  .onCompletion { stopAudioConversation() }
@@ -326,7 +326,10 @@ internal constructor(
  ?.listenToRecording()
  ?.buffer(UNLIMITED)
  ?.accumulateUntil(MIN_BUFFER_SIZE)
- ?.onEach { sendMediaStream(listOf(MediaData(it, "audio/pcm"))) }
+ ?.onEach {
+ sendMediaStream(listOf(MediaData(it, "audio/pcm")))
+ delay(0)
+ }
  ?.catch { throw FirebaseAIException.from(it) }
  ?.launchIn(scope)
  }
@@ -374,7 +377,7 @@ internal constructor(
  if (it.interrupted) {
  playBackQueue.clear()
  } else {
- println("Sending audio parts")
+ println("Queuing audio parts from model")
  val audioParts = it.content?.parts?.filterIsInstance<InlineDataPart>().orEmpty()
  for (part in audioParts) {
  playBackQueue.add(part.inlineData)
@@ -390,7 +393,7 @@ internal constructor(
  }
  }
  }
- .launchIn(CoroutineScope(Dispatchers.IO))
+ .launchIn(scope)
  }
 
  /**
@@ -401,7 +404,7 @@ internal constructor(
  * Launched asynchronously on [scope].
  */
  private fun listenForModelPlayback(enableInterruptions: Boolean = false) {
- CoroutineScope(Dispatchers.IO).launch {
+ scope.launch {
  while (isActive) {
  val playbackData = playBackQueue.poll()
  if (playbackData == null) {
@@ -410,8 +413,9 @@ internal constructor(
  if (!enableInterruptions) {
  audioHelper?.resumeRecording()
  }
- yield()
+ delay(0)
  } else {
+ println("Playing audio data")
  /**
  * We pause the recording while the model is speaking to avoid interrupting it because of
  * no echo cancellation

gradle/libs.versions.toml

@@ -16,7 +16,7 @@ benchmarkMacro = "1.3.4"
 browser = "1.3.0"
 cardview = "1.0.0"
 constraintlayout = "2.1.4"
-coroutines = "1.9.0"
+coroutines = "1.10.2"
 dagger = "2.51" # Don't bump above 2.51 as it causes a bug in AppDistro FeedbackSender JPEG code
 datastore = "1.1.7"
 dexmaker = "2.28.1"