Fetch the initial responses of watch() from the watcher interceptor

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 - [Improvement] Derive the keys a watcher matches cache changes against lazily, keeping the re-normalization of the whole response off the path that delivers the initial responses of `watch` (#378)
 - [Improvement] Reduce allocations when deriving a watcher's dependent keys and matching them against a cache change (#378)
+- [Improvement] `watch` now runs the operation through the interceptor chain once instead of twice. Its initial responses are fetched by the watcher interceptor, so the last of them no longer waits on a second execution of the chain to subscribe to the cache. Interceptors installed ahead of the cache run once per `watch` (#379)
+- [Fix] Cache headers set on the client are no longer discarded by a call that sets cache headers of its own. Whether the two were merged used to depend on the order the options were set in (#379)
 - Enable parallel sync for Tooling API clients (#380)
 
 PUT_CHANGELOG_HERE

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

@@ -10,13 +10,21 @@ import com.apollographql.apollo.exception.DefaultApolloException
 import com.apollographql.apollo.interceptor.ApolloInterceptor
 import com.apollographql.apollo.interceptor.ApolloInterceptorChain
 import com.apollographql.cache.normalized.CacheManager
+import com.apollographql.cache.normalized.DefaultFetchPolicyInterceptor
+import com.apollographql.cache.normalized.FetchPolicyContext
+import com.apollographql.cache.normalized.RefetchPolicyContext
 import com.apollographql.cache.normalized.api.CacheKey
 import com.apollographql.cache.normalized.api.dependentKeys
 import com.apollographql.cache.normalized.api.withErrors
+import com.apollographql.cache.normalized.options.noCache
+import com.apollographql.cache.normalized.options.onlyIfCached
+import com.apollographql.cache.normalized.options.refetchNoCache
+import com.apollographql.cache.normalized.options.refetchOnlyIfCached
 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.flow
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.onSubscription
 import kotlinx.coroutines.flow.transform
@@ -40,11 +48,11 @@ internal class WatcherInterceptor(val cacheManager: CacheManager) : ApolloInterc
      * 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.
+     * This matters because of when the work would otherwise happen: the last of the initial
+     * responses is withheld below 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, which is why those only record the
+     * data and leave the keys stale.
      */
     var dataToWatch: Operation.Data? = watchContext.data
     var errorsToWatch: List<Error>? = null
@@ -75,6 +83,36 @@ internal class WatcherInterceptor(val cacheManager: CacheManager) : ApolloInterc
       }
     }
 
+    /**
+     * The request used when the cache changes.
+     *
+     * `watch()` fetches its initial responses with the fetch policy and refetches with the refetch
+     * policy, so the latter has to be applied here rather than to the whole call. `watch(data)` has
+     * no initial fetch and keeps the fetch policy throughout.
+     */
+    val refetchRequest = if (watchContext.fetchInitialResponses) {
+      request.newBuilder()
+          .addExecutionContext(
+              FetchPolicyContext(request.executionContext[RefetchPolicyContext]?.interceptor ?: DefaultFetchPolicyInterceptor),
+          )
+          .noCache(request.refetchNoCache)
+          .onlyIfCached(request.refetchOnlyIfCached)
+          .build()
+    } else {
+      request
+    }
+
+    fun proceedRecordingData(request: ApolloRequest<D>): Flow<ApolloResponse<D>> {
+      return chain.proceed(request)
+          .onEach { response ->
+            if (response.data != null) {
+              dataToWatch = response.data
+              errorsToWatch = response.errors
+              watchedKeysAreStale = true
+            }
+          }
+    }
+
     fun isWatched(changedKeys: Set<*>): Boolean {
       if (changedKeys === CacheManager.ALL_KEYS) {
         // Matches regardless of the watched keys, so there is no need to compute them.
@@ -87,30 +125,58 @@ internal class WatcherInterceptor(val cacheManager: CacheManager) : ApolloInterc
       return (changedKeys as Set<String>).anyIntersection(watched)
     }
 
-    return (cacheManager.changedKeys as SharedFlow<Any>)
-        .onSubscription {
-          emit(Unit)
-        }
-        .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
+    return flow {
+      /**
+       * The last of the initial responses, withheld until the cache subscription is established so
+       * that callers can use it as a synchronisation point: modifying the store once it arrives is
+       * guaranteed to be observed. Subscribing first instead would make the watcher fire on the
+       * initial fetch's own write.
+       *
+       * See https://github.com/apollographql/apollo-kotlin/pull/3853
+       */
+      var lastResponse: ApolloResponse<D>? = null
+
+      if (watchContext.fetchInitialResponses) {
+        proceedRecordingData(request).collect { response ->
+          if (response.isLast) {
+            /**
+             * If we ever come here it means some interceptors built a new Flow and forgot to reset the isLast flag
+             * Better safe than sorry: emit them when we realize that. This will introduce a delay in the response.
+             */
+            lastResponse?.let { emit(it) }
+            lastResponse = response
+          } else {
+            emit(response)
           }
-          emitAll(
-              chain.proceed(request)
-                  .onEach { response ->
-                    if (response.data != null) {
-                      dataToWatch = response.data
-                      errorsToWatch = response.errors
-                      watchedKeysAreStale = true
-                    }
-                  }
-          )
         }
+      }
+
+      emitAll(
+          (cacheManager.changedKeys as SharedFlow<Any>)
+              .onSubscription {
+                emit(Unit)
+              }
+              .transform { event ->
+                if (event !is Set<*>) {
+                  // The marker emitted by `onSubscription`: subscribed, so the withheld response can
+                  // be released. Reaching this point costs a `SharedFlow` subscription and nothing
+                  // else, because the initial responses were fetched from this same execution of the
+                  // interceptor chain.
+                  val held = lastResponse
+                  if (held != null) {
+                    lastResponse = null
+                    emit(held)
+                  } else if (!watchContext.fetchInitialResponses) {
+                    emit(ApolloResponse.Builder(request.operation, request.requestUuid).exception(WatcherSentinel).build())
+                  }
+                  return@transform
+                }
+                if (!isWatched(event)) {
+                  return@transform
+                }
+                emitAll(proceedRecordingData(refetchRequest))
+              }
+      )
+    }
   }
 }

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/options/cacheHeaders.kt

@@ -12,29 +12,27 @@ import com.apollographql.cache.normalized.api.CacheHeaders
 import kotlin.time.Duration
 
 internal class CacheHeadersContext(val value: CacheHeaders) : ExecutionContext.Element {
-  override val key: ExecutionContext.Key<*>
-    get() = Key
-
   /**
-   * Merge the [CacheHeaders] instead of letting the right element replace the right element.
+   * Each [CacheHeadersContext] gets a unique key, so headers are merged as expected.
    */
-  override fun <R> fold(initial: R, operation: (R, ExecutionContext.Element) -> R): R {
-    val existing = (initial as? ExecutionContext)?.get(Key)
-    val element = if (existing != null) CacheHeadersContext(existing.value + value) else this
-    return operation(initial, element)
-  }
+  override val key: ExecutionContext.Key<*> = Key()
 
-  companion object Key : ExecutionContext.Key<CacheHeadersContext>
+  private class Key : ExecutionContext.Key<CacheHeadersContext>
 }
 
+private fun ExecutionContext.mergedCacheHeaders(): CacheHeaders =
+  fold(CacheHeaders.NONE) { acc, element ->
+    if (element is CacheHeadersContext) acc + element.value else acc
+  }
+
 internal val ExecutionOptions.cacheHeaders: CacheHeaders
-  get() = (executionContext[CacheHeadersContext]?.value ?: CacheHeaders.NONE)
+  get() = executionContext.mergedCacheHeaders()
 
 fun <D : Operation.Data> ApolloResponse.Builder<D>.cacheHeaders(cacheHeaders: CacheHeaders) =
   addExecutionContext(CacheHeadersContext(cacheHeaders))
 
 val <D : Operation.Data> ApolloResponse<D>.cacheHeaders
-  get() = executionContext[CacheHeadersContext]?.value ?: CacheHeaders.NONE
+  get() = executionContext.mergedCacheHeaders()
 
 
 /**
@@ -48,7 +46,7 @@ fun <T> MutableExecutionOptions<T>.cacheHeaders(cacheHeaders: CacheHeaders) = ad
  * Add a cache header to be passed to your [com.apollographql.cache.normalized.api.NormalizedCache]
  */
 fun <T> MutableExecutionOptions<T>.addCacheHeader(key: String, value: String) = cacheHeaders(
-    cacheHeaders.newBuilder().addHeader(key, value).build()
+    CacheHeaders.Builder().addHeader(key, value).build()
 )
 
 /**

normalized-cache/src/commonMain/kotlin/com/apollographql/cache/normalized/options/watch.kt

@@ -18,7 +18,18 @@ import kotlinx.coroutines.flow.filter
 import kotlinx.coroutines.flow.flow
 
 internal class WatchContext(
+    /**
+     * The data to derive the initially watched keys from, for the overload that does not fetch.
+     */
     val data: Query.Data?,
+
+    /**
+     * Whether to execute the request once, with the fetch policy, before observing the cache.
+     *
+     * Doing this from the interceptor rather than as a separate execution means the operation goes
+     * through the interceptor chain once instead of twice.
+     */
+    val fetchInitialResponses: Boolean,
 ) : ExecutionContext.Element {
   override val key: ExecutionContext.Key<*>
     get() = Key
@@ -33,8 +44,8 @@ internal val <D : Operation.Data> ApolloRequest<D>.watchContext: WatchContext?
 /**
  * Gets initial response(s) then observes the cache for any changes.
  *
- * The cache subscription is established before the initial fetch completes, so any external cache update made
- * after collecting the last initial response will be received.
+ * The cache subscription is established before the last initial response is collected, so any external cache update made
+ * after collecting it will be received.
  *
  * Note: when using [writeToCacheAsynchronously], the cache updates are postponed and behave as external cache updates. They may trigger emission. 
  *
@@ -46,53 +57,21 @@ internal val <D : Operation.Data> ApolloRequest<D>.watchContext: WatchContext?
  * @see refetchPolicy
  */
 fun <D : Query.Data> ApolloCall<D>.watch(): Flow<ApolloResponse<D>> {
-  return flow {
-    var lastResponse: ApolloResponse<D>? = null
-    var response: ApolloResponse<D>? = null
-
-    toFlow()
-        .collect {
-          response = it
-
-          if (it.isLast) {
-            if (lastResponse != null) {
-              /**
-               * If we ever come here it means some interceptors built a new Flow and forgot to reset the isLast flag
-               * Better safe than sorry: emit them when we realize that. This will introduce a delay in the response.
-               */
-              println("ApolloGraphQL: extra response received after the last one")
-              emit(lastResponse!!)
-            }
-            /**
-             * Remember the last response so that we can send it after we subscribe to the store
-             *
-             * This allows callers to use the last element as a synchronisation point to modify the store and still have the watcher
-             * receive subsequent updates
-             *
-             * See https://github.com/apollographql/apollo-kotlin/pull/3853
-             */
-            lastResponse = it
-          } else {
-            emit(it)
-          }
-        }
-
-
-    copy().fetchPolicyInterceptor(refetchPolicyInterceptor)
-        .noCache(refetchNoCache)
-        .onlyIfCached(refetchOnlyIfCached)
-        .watchInternal(response?.data)
-        .collect {
-          if (it.exception === WatcherSentinel) {
-            if (lastResponse != null) {
-              emit(lastResponse!!)
-              lastResponse = null
-            }
-          } else {
-            emit(it)
-          }
-        }
-  }
+  /**
+   * The initial responses are fetched by the interceptor, which subscribes to the cache right after
+   * and only then releases the last of them.
+   *
+   * Executing them here instead would mean a second trip through the interceptor chain to subscribe,
+   * and that last response would be withheld until the trip finished - delaying it by the teardown
+   * of the first flow, a dispatch, and a re-run of every interceptor ahead of the cache. Done from
+   * the interceptor, the wait is a [kotlinx.coroutines.flow.SharedFlow] subscription and nothing
+   * else, while the synchronisation point callers rely on is unchanged.
+   *
+   * See https://github.com/apollographql/apollo-kotlin/pull/3853
+   */
+  return copy()
+      .addExecutionContext(WatchContext(data = null, fetchInitialResponses = true))
+      .toFlow()
 }
 
 /**
@@ -108,5 +87,5 @@ fun <D : Query.Data> ApolloCall<D>.watch(data: D?): Flow<ApolloResponse<D>> {
  * The fetch policy set by [fetchPolicy] will be used.
  */
 internal fun <D : Query.Data> ApolloCall<D>.watchInternal(data: D?): Flow<ApolloResponse<D>> {
-  return copy().addExecutionContext(WatchContext(data)).toFlow()
+  return copy().addExecutionContext(WatchContext(data, fetchInitialResponses = false)).toFlow()
 }

tests/cache-options/src/commonTest/kotlin/test/CacheOptionsTest.kt

@@ -1073,6 +1073,61 @@ class CacheOptionsTest {
           }
         }
   }
+
+  /**
+   * Setting a cache header on a call must not discard the ones set on the client: both end up in the
+   * same [com.apollographql.apollo.api.ExecutionContext] and have to be merged.
+   */
+  @Test
+  fun clientCacheHeadersSurviveCallCacheHeaders() = runTest(before = { setUp() }, after = { tearDown() }) {
+    mockServer.enqueueString(
+        // language=JSON
+        """
+          {
+            "data": {
+              "car": {
+                "__typename": "Car",
+                "id": "1",
+                "color": "Red",
+                "doors": null
+              }
+            },
+            "errors": [
+              {
+                "message": "Doors does not exist",
+                "path": ["car", "doors"]
+              }
+            ]
+          }
+          """,
+    )
+
+    ApolloClient.Builder()
+        .serverUrl(mockServer.url())
+        .serverErrorsAsException(false)
+        .cacheManager(memoryThenSqlCacheManager)
+        .build()
+        .use { apolloClient ->
+          apolloClient.query(GetCarQuery(carId = "1"))
+              .fetchPolicy(FetchPolicy.NetworkOnly)
+              .execute()
+
+          val cacheResponse = apolloClient.query(GetCarQuery(carId = "1"))
+              .fetchPolicy(FetchPolicy.CacheOnly)
+              // A cache header of the call's own, which the client's must be merged with rather than
+              // replaced by.
+              .memoryCacheOnly(true)
+              .execute()
+
+          // serverErrorsAsException is false, so the stored error is in errors and not thrown
+          assertNull(cacheResponse.exception)
+          assertEquals("Red", cacheResponse.data?.car?.color)
+          assertErrorsEquals(
+              listOf(Error.Builder("Doors does not exist").path(listOf("car", "doors")).build()),
+              cacheResponse.errors,
+          )
+        }
+  }
 }
 
 private fun <D : Operation.Data> ApolloCall<D>.executeCacheAndNetwork(): Flow<ApolloResponse<D>> {