dataconnect(chore): refactor cache database management logic out of network layer (#8211)

Author: dconeybe

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

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2026 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.core
+
+import com.google.firebase.dataconnect.core.LoggerGlobals.Logger
+import com.google.firebase.dataconnect.core.LoggerGlobals.debug
+import com.google.firebase.dataconnect.sqlite.DataConnectCacheDatabase
+import com.google.firebase.dataconnect.util.ObjectLifecycleManager
+import java.io.File
+import kotlinx.coroutines.CoroutineDispatcher
+
+internal class DataConnectCache(
+  private val dbFile: File?,
+  val maxAge: kotlin.time.Duration,
+  private val cpuDispatcher: CoroutineDispatcher,
+  private val logger: Logger,
+) : ObjectLifecycleManager<DataConnectCacheDatabase>(cpuDispatcher, logger) {
+
+  val maxAgeProto: com.google.protobuf.Duration =
+    maxAge.toComponents { seconds, nanos ->
+      com.google.protobuf.Duration.newBuilder().setSeconds(seconds).setNanos(nanos).build()
+    }
+
+  override fun create() =
+    DataConnectCacheDatabase(
+      dbFile,
+      cpuDispatcher,
+      Logger("DataConnectCacheDatabase").apply { debug { "created by ${logger.nameWithId}" } }
+    )
+
+  override suspend fun initialize(instance: DataConnectCacheDatabase) {
+    instance.initialize()
+  }
+
+  override suspend fun destroy(instance: DataConnectCacheDatabase) {
+    instance.close()
+  }
+
+  override fun toString() = "DataConnectCache(dbFile=$dbFile, maxAge=$maxAge)"
+}

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

@@ -35,14 +35,12 @@ import com.google.firebase.dataconnect.util.CoroutineUtils
 import com.google.firebase.dataconnect.util.GrpcBidiFlow
 import com.google.firebase.dataconnect.util.GrpcBidiFlowListenerMessageFormatter
 import com.google.firebase.dataconnect.util.IdStringGenerator
-import com.google.firebase.dataconnect.util.NullableReference
 import com.google.firebase.dataconnect.util.ProtoUtil.buildStructProto
 import com.google.firebase.dataconnect.util.ProtoUtil.toCompactString
 import com.google.firebase.dataconnect.util.ProtoUtil.toDataConnectPath
 import com.google.firebase.dataconnect.util.ProtoUtil.toStructProto
 import com.google.firebase.dataconnect.util.SuspendingLazy
 import com.google.firebase.dataconnect.util.copy
-import com.google.protobuf.Duration as DurationProto
 import com.google.protobuf.Struct
 import google.firebase.dataconnect.proto.ConnectorServiceGrpc
 import google.firebase.dataconnect.proto.ConnectorServiceGrpcKt
@@ -69,11 +67,9 @@ import io.grpc.Status
 import io.grpc.StatusException
 import io.grpc.StatusRuntimeException
 import io.grpc.android.AndroidChannelBuilder
-import java.io.File
 import java.lang.System.currentTimeMillis
 import java.util.concurrent.Executor
 import java.util.concurrent.TimeUnit
-import kotlin.time.Duration
 import kotlinx.coroutines.CancellationException
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.asExecutor
@@ -95,7 +91,7 @@ internal class DataConnectGrpcRPCs(
   private val nonBlockingCoroutineDispatcher: CoroutineDispatcher,
   private val blockingCoroutineDispatcher: CoroutineDispatcher,
   private val grpcMetadata: DataConnectGrpcMetadata,
-  private val cacheSettings: CacheSettings?,
+  private val cache: DataConnectCache?,
   parentLogger: Logger,
 ) {
   private val logger =
@@ -104,7 +100,7 @@ internal class DataConnectGrpcRPCs(
         "created by ${parentLogger.nameWithId} with" +
           " host=$host" +
           " sslEnabled=$sslEnabled" +
-          " cacheSettings=$cacheSettings" +
+          " cache=$cache" +
           " grpcMetadata=${grpcMetadata.instanceId}"
       }
     }
@@ -119,38 +115,6 @@ internal class DataConnectGrpcRPCs(
   private val mutex = Mutex()
   private var closed = false
 
-  data class CacheSettings(val dbFile: File?, val maxAge: Duration)
-
-  private data class CacheDbSettingsPair(
-    val db: DataConnectCacheDatabase,
-    val maxAge: DurationProto,
-  )
-
-  // Use the non-main-thread CoroutineDispatcher to avoid blocking operations on the main thread.
-  private val lazyCacheDb =
-    SuspendingLazy(mutex = mutex, coroutineContext = blockingCoroutineDispatcher) {
-      check(!closed) { "DataConnectGrpcRPCs ${logger.nameWithId} instance has been closed" }
-      if (cacheSettings === null) {
-        NullableReference()
-      } else {
-        logger.debug { "Creating GRPC ManagedChannel for host=$host sslEnabled=$sslEnabled" }
-
-        val maxAge =
-          cacheSettings.maxAge.toComponents { seconds, nanos ->
-            DurationProto.newBuilder().setSeconds(seconds).setNanos(nanos).build()
-          }
-
-        val dbFile = cacheSettings.dbFile
-        val cacheLogger = Logger("DataConnectCacheDatabase")
-        cacheLogger.debug {
-          "created by ${logger.nameWithId} with dbFile=$dbFile maxAge=${cacheSettings.maxAge}"
-        }
-        val cacheDb = DataConnectCacheDatabase(dbFile, cacheLogger)
-        cacheDb.initialize()
-        NullableReference(CacheDbSettingsPair(cacheDb, maxAge))
-      }
-    }
-
   // Use the non-main-thread CoroutineDispatcher to avoid blocking operations on the main thread.
   private val lazyGrpcChannel =
     SuspendingLazy(mutex = mutex, coroutineContext = blockingCoroutineDispatcher) {
@@ -238,28 +202,20 @@ internal class DataConnectGrpcRPCs(
   }
 
   private class QueryCacheInfo(
-    val cacheDb: DataConnectCacheDatabase,
+    val cache: DataConnectCache,
     val authUid: AuthUid?,
     val queryId: QueryId,
-    val maxAge: DurationProto,
   )
 
-  private suspend fun queryCacheInfo(
+  private suspend fun DataConnectCache.queryCacheInfo(
     authToken: DataConnectAuth.GetAuthTokenResult?,
     request: ExecuteQueryRequest,
-  ): QueryCacheInfo? {
+  ): QueryCacheInfo {
     val queryId =
       withContext(nonBlockingCoroutineDispatcher) {
         calculateQueryId(request.operationName, request.variables)
       }
-    return lazyCacheDb.get().ref?.let { (cacheDb, maxAge) ->
-      QueryCacheInfo(
-        cacheDb,
-        authUid = authToken?.authUid,
-        queryId = queryId,
-        maxAge = maxAge,
-      )
-    }
+    return QueryCacheInfo(this, authToken?.authUid, queryId)
   }
 
   suspend fun executeQuery(
@@ -282,7 +238,7 @@ internal class DataConnectGrpcRPCs(
         it.build()
       }
 
-    val cacheInfo = queryCacheInfo(authToken, request)
+    val cacheInfo = cache?.queryCacheInfo(authToken, request)
     if (cacheInfo == null && fetchPolicy == FetchPolicy.CACHE_ONLY) {
       throw CachedDataNotFoundException(
         "FetchPolicy.CACHE_ONLY cannot be used because local caching is not configured. " +
@@ -335,14 +291,16 @@ internal class DataConnectGrpcRPCs(
       )
 
       cacheInfo?.run {
-        cacheDb.insertQueryResult(
-          authUid,
-          queryId,
-          queryData = response.data,
-          maxAge = maxAge,
-          currentTimeMillis = currentTimeMillis(),
-          getEntityIdForPath = response.getEntityIdForPathFunction(),
-        )
+        cache
+          .open()
+          .insertQueryResult(
+            authUid,
+            queryId,
+            queryData = response.data,
+            maxAge = cache.maxAgeProto,
+            currentTimeMillis = currentTimeMillis(),
+            getEntityIdForPath = response.getEntityIdForPathFunction(),
+          )
       }
     }
 
@@ -368,7 +326,8 @@ internal class DataConnectGrpcRPCs(
         else -> DataConnectCacheDatabase.GetQueryResultResult.Stale::class
       }
 
-    val cachedResult = cacheDb.getQueryResult(authUid, queryId, currentTimeMillis(), staleResult)
+    val cachedResult =
+      cache.open().getQueryResult(authUid, queryId, currentTimeMillis(), staleResult)
 
     val cachedData: Struct? =
       when (cachedResult) {
@@ -607,37 +566,14 @@ internal class DataConnectGrpcRPCs(
   suspend fun close() {
     logger.debug { "close()" }
     mutex.withLock { closed = true }
-    connectCoroutineScope.cancel("DataConnectGrpcRPCs.close() called [xn8dqn8dzm]")
 
-    val grpcChannel = lazyGrpcChannel.initializedValueOrNull
-    val cacheDb = lazyCacheDb.initializedValueOrNull?.ref
-
-    if (grpcChannel === null && cacheDb === null) {
-      connectCoroutineScope.coroutineContext.job.join()
-      return
-    }
+    connectCoroutineScope.cancel("DataConnectGrpcRPCs.close() called [xn8dqn8dzm]")
 
-    // Avoid blocking the calling thread by running potentially-blocking code on the dispatcher
-    // given to the constructor, which should have similar semantics to [Dispatchers.IO].
-    val grpcChannelShutdownResult: Result<*>
-    val cacheDbCloseResult: Result<*>
-    withContext(blockingCoroutineDispatcher) {
-      grpcChannelShutdownResult = runCatching {
-        grpcChannel?.shutdownNow()
-        grpcChannel?.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS)
+    lazyGrpcChannel.initializedValueOrNull?.let { grpcChannel ->
+      withContext(blockingCoroutineDispatcher) {
+        grpcChannel.shutdownNow()
+        grpcChannel.awaitTermination(Long.MAX_VALUE, TimeUnit.SECONDS)
       }
-      cacheDbCloseResult = runCatching { cacheDb?.db?.close() }
-    }
-
-    // Bundle together any exceptions that were thrown.
-    val exceptions =
-      listOf(
-          grpcChannelShutdownResult,
-          cacheDbCloseResult,
-        )
-        .mapNotNull { it.exceptionOrNull() }
-    if (exceptions.isNotEmpty()) {
-      throw exceptions.first().apply { exceptions.drop(1).forEach { addSuppressed(it) } }
     }
 
     connectCoroutineScope.coroutineContext.job.join()