dataconnect(change): Add 15 second grace before disconnecting a streaming connection with the backend (#8481)

Author: dconeybe

firebase-dataconnect/CHANGELOG.md

@@ -10,6 +10,10 @@
 - [changed] Add grpc request headers for platform name and sdk version
   to enable metrics collection in cloud monitoring.
   ([#8486](https://github.com/firebase/firebase-android-sdk/pull/8486))
+- [changed] Wait for 15 seconds before closing realtime streaming connection
+  with backend after last subscriber unsubscribes (instead of closing the
+  connection immediately).
+  ([#8481](https://github.com/firebase/firebase-android-sdk/pull/8481))
 
 # 17.3.2
 

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

@@ -233,7 +233,11 @@ internal class DataConnectBidiConnectStream(
         .buffer(capacity = 64) // Use a finite buffer to activate gRPC flow control, when needed
         .shareIn(
           coroutineScope,
-          started = SharingStarted.WhileSubscribed(replayExpirationMillis = 0),
+          started =
+            SharingStarted.WhileSubscribed(
+              stopTimeoutMillis = 15_000,
+              replayExpirationMillis = 0,
+            ),
           replay = 0,
         )
 

firebase-dataconnect/src/test/kotlin/com/google/firebase/dataconnect/core/QuerySubscriptionImplUnitTest.kt

@@ -19,13 +19,15 @@ import android.content.Context.CONNECTIVITY_SERVICE
 import android.net.ConnectivityManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import app.cash.turbine.ReceiveTurbine
+import app.cash.turbine.TurbineContext
 import app.cash.turbine.test
 import app.cash.turbine.turbineScope
 import com.google.firebase.appcheck.interop.InteropAppCheckTokenProvider
 import com.google.firebase.auth.internal.InternalAuthProvider
 import com.google.firebase.dataconnect.DataConnectSettings
 import com.google.firebase.dataconnect.FirebaseDataConnect.CallerSdkType
 import com.google.firebase.dataconnect.QueryRef
+import com.google.firebase.dataconnect.QuerySubscriptionResult
 import com.google.firebase.dataconnect.core.DataConnectAuth.GetAuthTokenResult
 import com.google.firebase.dataconnect.core.DataConnectBidiConnectStream.Companion.setReconnectPendingAuthTokenForTesting
 import com.google.firebase.dataconnect.core.DataConnectBidiConnectStream.Companion.unsetReconnectPendingAuthTokenForTesting
@@ -91,9 +93,11 @@ import io.kotest.assertions.print.print
 import io.kotest.assertions.withClue
 import io.kotest.common.DelicateKotest
 import io.kotest.common.ExperimentalKotest
+import io.kotest.matchers.booleans.shouldBeTrue
 import io.kotest.matchers.collections.shouldBeIn
 import io.kotest.matchers.collections.shouldContainExactly
 import io.kotest.matchers.maps.shouldBeEmpty
+import io.kotest.matchers.nulls.shouldBeNull
 import io.kotest.matchers.result.shouldBeSuccess
 import io.kotest.matchers.shouldBe
 import io.kotest.matchers.shouldNotBe
@@ -111,6 +115,7 @@ import io.kotest.property.arbitrary.az
 import io.kotest.property.arbitrary.distinct
 import io.kotest.property.arbitrary.enum
 import io.kotest.property.arbitrary.int
+import io.kotest.property.arbitrary.long
 import io.kotest.property.arbitrary.map
 import io.kotest.property.arbitrary.next
 import io.kotest.property.arbitrary.of
@@ -2036,6 +2041,133 @@ class QuerySubscriptionImplUnitTest {
     }
   }
 
+  @Test
+  fun `connection is kept alive for exactly the grace period after last subscriber unsubscribes`() =
+    testConnectionGracePeriod(Arb.long(0L until CONNECTION_GRACE_PERIOD_MS)) { context ->
+      context.clientCollector1.cancelAndIgnoreRemainingEvents()
+
+      // Wait a random duration less than CONNECTION_GRACE_PERIOD_MS
+      delay(context.delayMillis.milliseconds)
+
+      // Verify that the connection is still open (no close event received on serverCollector).
+      // Based on the timing, we _may_ receive the "cancel" event, which is expected.
+      val event = context.serverCollector.asChannel().tryReceive().getOrNull()
+      if (event != null) {
+        val streamRequest = event.shouldBeInstanceOf<StreamRequestReceived>().streamRequest
+        streamRequest.hasCancel().shouldBeTrue()
+        context.serverCollector.asChannel().tryReceive().getOrNull().shouldBeNull()
+      }
+
+      // Measure the remaining time until the client closes the connection
+      val time1 = @OptIn(ExperimentalCoroutinesApi::class) context.testScheduler.currentTime
+      context.serverCollector.awaitUntilClientClosesConnection()
+      val time2 = @OptIn(ExperimentalCoroutinesApi::class) context.testScheduler.currentTime
+
+      (time2 - time1) shouldBe (CONNECTION_GRACE_PERIOD_MS - context.delayMillis)
+    }
+
+  @Test
+  fun `re-subscribing within grace period keeps the connection alive and reuses it`() =
+    testConnectionGracePeriod(Arb.long(100L until CONNECTION_GRACE_PERIOD_MS)) { context ->
+      context.clientCollector1.cancelAndIgnoreRemainingEvents()
+
+      // Wait a random duration less than CONNECTION_GRACE_PERIOD_MS
+      delay(context.delayMillis.milliseconds)
+
+      // Re-subscribe clientCollector2
+      val clientCollector2 =
+        context.subscription.flow.testIn(context.backgroundScope, name = "clientCollector2")
+
+      // Verify that the server receives a subscribe request for the connection, but the connection
+      // ID remains the same (proving reuse)
+      val subscribeRequest = context.serverCollector.awaitUntilSubscribeStreamRequest()
+      subscribeRequest.connectionId shouldBe context.initialConnectionId
+
+      clientCollector2.cancelAndIgnoreRemainingEvents()
+      // Wait for the new grace period to expire so we don't leak connection closure errors/events
+      context.serverCollector.awaitUntilClientClosesConnection()
+    }
+
+  @Test
+  fun `re-subscribing after grace period establishes a new connection`() =
+    testConnectionGracePeriod(
+      Arb.long(CONNECTION_GRACE_PERIOD_MS..(CONNECTION_GRACE_PERIOD_MS * 4))
+    ) { context ->
+      context.clientCollector1.cancelAndIgnoreRemainingEvents()
+
+      // Wait a random duration of at least CONNECTION_GRACE_PERIOD_MS
+      delay(context.delayMillis.milliseconds)
+
+      // Verify that the connection is closed
+      context.serverCollector.awaitUntilClientClosesConnection()
+
+      // Re-subscribe clientCollector2
+      val clientCollector2 =
+        context.subscription.flow.testIn(context.backgroundScope, name = "clientCollector2")
+
+      // Verify that a new connection is established
+      val connection2 = context.serverCollector.awaitConnectRpcStarted()
+      connection2.connectionId shouldNotBe context.initialConnectionId
+      context.serverCollector.awaitUntilInitStreamRequest()
+      context.serverCollector.awaitUntilSubscribeStreamRequest()
+
+      clientCollector2.cancelAndIgnoreRemainingEvents()
+      // Wait for connection to close to keep clean state
+      context.serverCollector.awaitUntilClientClosesConnection()
+    }
+
+  @OptIn(ExperimentalCoroutinesApi::class)
+  private fun testConnectionGracePeriod(
+    delayMillisArb: Arb<Long>,
+    block: suspend TurbineContext.(TestConnectionGracePeriodContext) -> Unit
+  ) = runTest {
+    val server = runningInProcessDataConnectServer()
+    checkAll(propTestConfig, delayMillisArb, Arb.dataConnect.operationName(), testVariablesArb()) {
+      delayMillis,
+      operationName,
+      variables ->
+      runWithDataConnect(server) { dataConnect ->
+        val subscription = querySubscription(dataConnect, operationName, variables)
+
+        turbineScope {
+          val serverCollector = server.events.testIn(backgroundScope, name = "serverCollector")
+          val clientCollector1 =
+            subscription.flow.testIn(backgroundScope, name = "clientCollector1")
+
+          // Wait for initial connection and subscribe
+          val connection = serverCollector.awaitConnectRpcStarted()
+          serverCollector.awaitUntilInitStreamRequest()
+          serverCollector.awaitUntilSubscribeStreamRequest()
+
+          val context =
+            TestConnectionGracePeriodContext(
+              delayMillis = delayMillis,
+              serverCollector = serverCollector,
+              clientCollector1 = clientCollector1,
+              initialConnectionId = connection.connectionId,
+              subscription = subscription,
+              backgroundScope = backgroundScope,
+              testScheduler = testScheduler,
+            )
+
+          block(context)
+
+          serverCollector.cancelAndIgnoreRemainingEvents()
+        }
+      }
+    }
+  }
+
+  private data class TestConnectionGracePeriodContext(
+    val delayMillis: Long,
+    val serverCollector: ReceiveTurbine<InProcessDataConnectGrpcStreamingServer.Event>,
+    val clientCollector1: ReceiveTurbine<QuerySubscriptionResult<TestData, TestVariables>>,
+    val initialConnectionId: InProcessDataConnectGrpcStreamingServer.ConnectionId,
+    val subscription: QuerySubscriptionImpl<TestData, TestVariables>,
+    val backgroundScope: kotlinx.coroutines.CoroutineScope,
+    val testScheduler: TestCoroutineScheduler,
+  )
+
   private fun runningInProcessDataConnectServer(): InProcessDataConnectGrpcStreamingServer {
     val server = InProcessDataConnectGrpcStreamingServer()
     cleanups.register(server)
@@ -2273,3 +2405,14 @@ private class SequenceRandom(jitters: Sequence<Double>) : Random() {
     return lock.withLock { iterator.next() + 0.5 }
   }
 }
+
+/**
+ * The amount of time, in milliseconds, that [DataConnectBidiConnectStream] keeps the physical
+ * connection with the backend alive after the last subscriber unsubscribes. By keeping the
+ * connection alive for a short amount of time rather than closing it immediately it improves the
+ * latency and reduces the backend load if a new subscriber were to subscribe within this grace
+ * period. This rapid unsubscription and resubscription could happen, for example, between activity
+ * or fragment transitions in an application where the old activity/fragment unsubscribes in its
+ * onDestory() and the new activity/fragment subscribes in its onCreate().
+ */
+const val CONNECTION_GRACE_PERIOD_MS = 15_000L