Using TaskLocal in Interceptor Chain

[PatrickDanino]

Question

General question to try to clarify with more specificity how @TaskLocal is supposed to help us track context across interceptors. I'll caveat my questions by admitting that my experience with this is limited.

The examples that are in the docs seem somewhat constrained. More specifically, they all seem to assume that you have the data available at the start of the chain, not somewhere further down the chain. My understand of a limitation of @TaskLocal is that it only goes from Parent to child, not from child back to parent.

Said another way, unless you define your @TaskLocal at the very start of the entire interceptor chain, it's unlikely to be useful somewhere else because you otherwise have no way of being sure that async task has the context already defined.

Separately, it's not even clear to me that if you succeed in defining the @TaskLocal, it'll actually make its way down. This is in part due to code like this throughout the interceptors:

extension AsyncThrowingStream where Failure == any Swift.Error {
  static func executingInAsyncTask(
    bufferingPolicy: AsyncThrowingStream<Element, Failure>.Continuation.BufferingPolicy = .unbounded,
    _ block: @escaping @Sendable (Continuation) async throws -> Void
  ) -> Self {
    return AsyncThrowingStream(bufferingPolicy: bufferingPolicy) { continuation in
      let task = Task {
        do {
          try await block(continuation)
          continuation.finish()

        } catch {
          continuation.finish(throwing: error)
        }
      }

      continuation.onTermination = { _ in task.cancel() }
    }
  }
}

My (limited) understanding of this snippet is that the code being executed in the try await block(continuation) will not inherit the parent context due to the explicit Task { ... } declaration from within a continuation block. Given AsyncThrowingStream.executingInAsyncTask is being called in numerous places, all these blocks have their own task/context that isn't shared.

Am I misunderstanding or missing something? I'll admit it's still unclear to me how @TaskLocal can help expose the HTTP headers available in the GQL response given it's made available so late in the interceptor chain.

[PatrickDanino]

Keeping the above original question, but since then I had the chance to test out the parent/child @TaskLocal and that does appear to work even though the Task is created within a continuation.

There's a somewhat vague reference to AsyncStream continuations being "special" in the Apple docs, but it's not clear (to me) as to how this can be guaranteed.

For the time being, I've tried to work around this by creating a Sendable class that is thread-safe. This allows me to essentially create a context per operation and pass it through all the way down, then update it later. It did however require having to wrap various ApolloClient operations, but that's not a huge deal for us.

[AnthonyMDev]

Hi @PatrickDanino.

My (limited) understanding of this snippet is that the code being executed in the try await block(continuation) will not inherit the parent context due to the explicit Task { ... } declaration from within a continuation block

As you found, this is incorrect. The TaskLocal value will be inherited by child Tasks. The only time it is not is when you use Task.detached.

For the time being, I've tried to work around this by creating a Sendable class that is thread-safe. This allows me to essentially create a context per operation and pass it through all the way down, then update it later. It did however require having to wrap various ApolloClient operations, but that's not a huge deal for us.

I'm not sure what other issue you are trying to work around here. Could you give me some more information on what you're trying to achieve?

[PatrickDanino]

There are two claims made in the docs. You're correct that using @TaskLocal appears to work in child tasks, but the docs also claim the following:

You may define a @TaskLocal value anywhere you would like. In this example, we define a @TaskLocal value in a TracingInterceptor.

I don't understand how that claim can be true. You must define the task local in a parent task and prior to other interceptors, otherwise it won't be set.

As a simple example, imagine you have an HTTPInterceptor and want to measure latency metrics for the network operation. For all intents and purposes, you can't really define a @TaskLocal there because it will start very late and finish very early, so any "parent" GraphQLInterceptor interceptor (for example, the TracingInterceptor from your docs) won't have access to that value.

So in practice, the only way you can truly ensure any interceptor has reliable access to a @TaskLocal value and make that value available up the chain is if it's defined very early, which most likely means having to wrap all the ApolloClient operations like fetch(...).

Hope that clarifies?

[AnthonyMDev]

You could wrap your fetch request, but you could also start setting the TaskLocal value in a GraphQLInterceptor and just make sure that is the first interceptor in your provider. It will then be called immediately upon kicking off the request chain.

If you really need to have this value before kicking off the request chain, then yes you can wrap the ApolloClient and do it there.

it will start very late and finish very early, so any "parent" GraphQLInterceptor interceptor (for example, the TracingInterceptor from your docs) won't have access to that value.

Because the request chain goes down the interceptor chain and then back up, the TaskLocal values set in a later interceptor will be available while going back up through the previous interceptors still. This would be available from within the map function of the result stream.

[github-actions]

Do you have any feedback for the maintainers? Please tell us by taking a one-minute survey. Your responses will help us understand Apollo iOS usage and allow us to serve you better.