Compute watcher dependent keys lazily instead of before subscribing

WatcherInterceptor derived the watched dependent keys eagerly, before
subscribing to CacheManager.changedKeys. That work re-serializes the whole
response through MapJsonWriter (withErrors) and normalizes it again
(normalize + dependentKeys), duplicating what ApolloCacheInterceptor already
does when it writes the response to the cache.

It also lands on a latency-sensitive path. watch() withholds the last of its
initial responses until this interceptor has subscribed, so everything done
before subscribing delays that response reaching the caller. For large
operations the extra normalization is substantial: on a big query in a
production Android app it accounted for most of a ~100ms gap between the
network response being ready and the collector observing it, compared to the
same query run through query().

The keys are only ever read to filter an incoming cache change, so compute
them on the first such change instead. The subscription marker and ALL_KEYS
events are answered without needing them at all, and a watcher that is
cancelled before any change never pays the cost. Refetch responses now only
record their data and mark the keys stale, keeping the same work off their
delivery path too.

Filtering is unchanged: the same keys are derived from the same data, just
later.

Three related cleanups while here:

- Add Set.anyIntersection(), which short-circuits on the first shared element
  and looks up against the larger set, replacing an intersect() call that
  materialized a whole set only to test it for emptiness.

- Build the dependent keys in a single pass. Going through Record.fieldKeys()
  allocated a list and a set per record, plus a list holding every key, on top
  of the set actually returned.

- Replace filter + map + flattenConcatPolyfill with a single transform, which
  has the same sequential semantics without a flow per event, and drops the
  re-test of a condition the filter had already decided.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: AlexanderGH

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/api/Record.kt

@@ -142,8 +142,22 @@ fun Record.expirationDate(field: String) = metadata[field]?.get(ApolloCacheHeade
  */
 typealias RecordValue = Any?
 
+/**
+ * Returns the set of all field keys of all the given records.
+ * A field key incorporates any GraphQL arguments in addition to the field name.
+ */
 fun Collection<Record>?.dependentKeys(): Set<String> {
-  return this?.flatMap {
-    it.fieldKeys()
-  }?.toSet() ?: emptySet()
+  if (this == null) {
+    return emptySet()
+  }
+  // Built in one pass: going through `fieldKeys()` would allocate a list and a set per record, plus
+  // a list holding every key, on top of the set actually returned. For a large operation that is
+  // thousands of throwaway collections.
+  return buildSet {
+    for (record in this@dependentKeys) {
+      for (fieldName in record.fields.keys) {
+        add(record.key.fieldKey(fieldName))
+      }
+    }
+  }
 }

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/internal/WatcherInterceptor.kt

@@ -3,6 +3,7 @@ package com.apollographql.cache.normalized.internal
 import com.apollographql.apollo.api.ApolloRequest
 import com.apollographql.apollo.api.ApolloResponse
 import com.apollographql.apollo.api.CustomScalarAdapters
+import com.apollographql.apollo.api.Error
 import com.apollographql.apollo.api.Operation
 import com.apollographql.apollo.api.Query
 import com.apollographql.apollo.exception.DefaultApolloException
@@ -16,12 +17,9 @@ import com.apollographql.cache.normalized.watchContext
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.emitAll
-import kotlinx.coroutines.flow.filter
-import kotlinx.coroutines.flow.flow
-import kotlinx.coroutines.flow.flowOf
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onSubscription
+import kotlinx.coroutines.flow.transform
 
 internal val WatcherSentinel = DefaultApolloException("The watcher has started")
 
@@ -35,12 +33,37 @@ internal class WatcherInterceptor(val cacheManager: CacheManager) : ApolloInterc
 
     val customScalarAdapters = request.executionContext[CustomScalarAdapters]!!
 
+    /**
+     * The data whose dependent keys are watched, and the keys themselves.
+     *
+     * Computing the keys normalizes the whole data again, which is significant work for large
+     * operations. The keys are only ever read to filter an incoming cache change, so they are
+     * computed on the first such change rather than up front.
+     *
+     * This matters because of when the work would otherwise happen:
+     * [com.apollographql.cache.normalized.watch] withholds the last of its initial responses until
+     * this interceptor has subscribed to [CacheManager.changedKeys], so anything done before
+     * subscribing delays that response reaching the caller. The same applies to the responses of a
+     * refetch below, which is why those only record the data and leave the keys stale.
+     */
+    var dataToWatch: Operation.Data? = watchContext.data
+    var errorsToWatch: List<Error>? = null
+    var watchedKeys: Set<String>? = null
+    var watchedKeysAreStale = true
+
     @Suppress("UNCHECKED_CAST")
-    var watchedKeys: Set<String>? =
-      watchContext.data?.let { data ->
+    fun computeWatchedKeysIfStale() {
+      if (!watchedKeysAreStale) {
+        return
+      }
+      watchedKeysAreStale = false
+      val data = dataToWatch
+      watchedKeys = if (data == null) {
+        null
+      } else {
         val dataWithErrors = (data as D).withErrors(
             executable = request.operation,
-            errors = null,
+            errors = errorsToWatch,
             customScalarAdapters = customScalarAdapters,
         )
         cacheManager.normalize(
@@ -50,47 +73,44 @@ internal class WatcherInterceptor(val cacheManager: CacheManager) : ApolloInterc
             customScalarAdapters = customScalarAdapters,
         ).values.dependentKeys()
       }
+    }
+
+    fun isWatched(changedKeys: Set<*>): Boolean {
+      if (changedKeys === CacheManager.ALL_KEYS) {
+        // Matches regardless of the watched keys, so there is no need to compute them.
+        return true
+      }
+      computeWatchedKeysIfStale()
+      val watched = watchedKeys ?: return true
+
+      @Suppress("UNCHECKED_CAST")
+      return (changedKeys as Set<String>).anyIntersection(watched)
+    }
 
     return (cacheManager.changedKeys as SharedFlow<Any>)
         .onSubscription {
           emit(Unit)
         }
-        .filter { changedKeys ->
-          changedKeys !is Set<*> ||
-              changedKeys === CacheManager.ALL_KEYS ||
-              watchedKeys == null ||
-              changedKeys.intersect(watchedKeys!!).isNotEmpty()
-        }.map {
-          if (it == Unit) {
-            flowOf(ApolloResponse.Builder(request.operation, request.requestUuid).exception(WatcherSentinel).build())
-          } else {
-            chain.proceed(request)
-                .onEach { response ->
-                  if (response.data != null) {
-                    val dataWithErrors = response.data!!.withErrors(
-                        executable = request.operation,
-                        errors = response.errors,
-                        customScalarAdapters = customScalarAdapters,
-                    )
-                    watchedKeys = cacheManager.normalize(
-                        executable = request.operation,
-                        dataWithErrors = dataWithErrors,
-                        rootKey = CacheKey.QUERY_ROOT,
-                        customScalarAdapters = customScalarAdapters,
-                    ).values.dependentKeys()
-                  }
-                }
+        .transform { event ->
+          if (event !is Set<*>) {
+            // The marker emitted by `onSubscription`. Answered without computing the watched keys,
+            // so that the initial response of `watch` is not held back by normalization.
+            emit(ApolloResponse.Builder(request.operation, request.requestUuid).exception(WatcherSentinel).build())
+            return@transform
+          }
+          if (!isWatched(event)) {
+            return@transform
           }
+          emitAll(
+              chain.proceed(request)
+                  .onEach { response ->
+                    if (response.data != null) {
+                      dataToWatch = response.data
+                      errorsToWatch = response.errors
+                      watchedKeysAreStale = true
+                    }
+                  }
+          )
         }
-        .flattenConcatPolyfill()
   }
 }
-
-/**
- * A copy/paste of the kotlinx.coroutines version until it becomes stable
- *
- * This is taken from 1.5.2 and replacing `unsafeFlow {}` with `flow {}`
- */
-private fun <T> Flow<Flow<T>>.flattenConcatPolyfill(): Flow<T> = flow {
-  collect { value -> emitAll(value) }
-}

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/internal/anyIntersection.kt

@@ -0,0 +1,16 @@
+package com.apollographql.cache.normalized.internal
+
+/**
+ * Returns whether this set and [other] have at least one element in common.
+ *
+ * Unlike `intersect`, this allocates nothing and stops at the first match. Elements are looked up in
+ * the larger of the two sets, so the number of (constant time) lookups is bounded by the size of the
+ * smaller one.
+ */
+internal fun <T> Set<T>.anyIntersection(other: Set<T>): Boolean {
+  return if (size < other.size) {
+    any { it in other }
+  } else {
+    other.any { it in this }
+  }
+}

tests/normalized-cache/src/commonTest/kotlin/WatcherTest.kt

@@ -62,6 +62,12 @@ class WatcherTest {
   private val episodeHeroNameChangedTwoData = EpisodeHeroNameQuery.Data(EpisodeHeroNameQuery.Hero("Droid", "ArTwo"))
 
   private val episodeHeroNameWithIdData = EpisodeHeroNameWithIdQuery.Data(EpisodeHeroNameWithIdQuery.Hero("Droid", "2001", "R2-D2"))
+  private val episodeHeroNameWithIdChangedData =
+    EpisodeHeroNameWithIdQuery.Data(EpisodeHeroNameWithIdQuery.Hero("Droid", "2001", "ArTwo"))
+
+  private val starshipByIdData = StarshipByIdQuery.Data(
+      StarshipByIdQuery.Starship("Starship", "Starship1", "SuperRocket", listOf(listOf(900.0, 800.0)))
+  )
 
 
   private val heroAndFriendsNamesWithIDsData = HeroAndFriendsNamesWithIDsQuery.Data(
@@ -316,6 +322,57 @@ class WatcherTest {
     job.cancel()
   }
 
+  /**
+   * The keys a watcher matches cache changes against come from the latest response it has seen, and
+   * they are only computed when a change actually needs them. Both of those have to keep holding
+   * across a refetch: a watcher that loses its keys would match every subsequent change.
+   */
+  @Test
+  fun watchedKeysStillFilterAfterARefetch() = runTest(before = { setUp() }) {
+    val channel = Channel<EpisodeHeroNameWithIdQuery.Data?>()
+
+    // The first query should get a "R2-D2" name
+    val episodeHeroNameWithIdQuery = EpisodeHeroNameWithIdQuery(Episode.EMPIRE)
+    apolloClient.enqueueTestResponse(episodeHeroNameWithIdQuery, episodeHeroNameWithIdData)
+    val job = launch {
+      apolloClient.query(episodeHeroNameWithIdQuery).watch().collect {
+        channel.send(it.data)
+      }
+    }
+
+    // Cache miss is emitted first (null data)
+    assertNull(channel.awaitElement())
+    assertEquals(channel.awaitElement()?.hero?.name, "R2-D2")
+
+    // An overlapping query triggers a refetch, which is where the watched keys are refreshed
+    val heroAndFriendsNamesWithIDsQuery = HeroAndFriendsNamesWithIDsQuery(Episode.NEWHOPE)
+    apolloClient.enqueueTestResponse(heroAndFriendsNamesWithIDsQuery, heroAndFriendsNamesWithIDsNameChangedData)
+    apolloClient.query(heroAndFriendsNamesWithIDsQuery)
+        .fetchPolicy(FetchPolicy.NetworkOnly)
+        .execute()
+
+    assertEquals(channel.awaitElement()?.hero?.name, "Artoo")
+
+    // A query sharing no key with the watched one is still filtered out after that refetch
+    val starshipByIdQuery = StarshipByIdQuery("Starship1")
+    apolloClient.enqueueTestResponse(starshipByIdQuery, starshipByIdData)
+    apolloClient.query(starshipByIdQuery)
+        .fetchPolicy(FetchPolicy.NetworkOnly)
+        .execute()
+
+    channel.assertEmpty()
+
+    // ...while an overlapping one still reaches the watcher
+    apolloClient.enqueueTestResponse(episodeHeroNameWithIdQuery, episodeHeroNameWithIdChangedData)
+    apolloClient.query(episodeHeroNameWithIdQuery)
+        .fetchPolicy(FetchPolicy.NetworkOnly)
+        .execute()
+
+    assertEquals(channel.awaitElement()?.hero?.name, "ArTwo")
+
+    job.cancel()
+  }
+
   /**
    * A test to test refetching with a NetworkOnly refetchPolicy. On every change, the watcher should get new information
    * from the network