Propagating TaskLocal storage to escaping closures

Using Swift
1

A problem I frequently run into: TaskLocal values do not propagate into escaping closures (except unstructured tasks).

@TaskLocal var requestID = 0

$requestID.withValue(1) {
  print(requestID) // 1
  someMethodWithEscapingClosure {
    print(requestID) // 0
  }
}

I'd love for there to be some way to propagate all TaskLocal values into an escaping closure. I'm not sure what the API would look like, but one idea:

$requestID.withValue(1) { withTaskLocals in
  print(requestID) // 1
  someMethodWithEscapingClosure {
    print(requestID) // 0
    withTaskLocals {
      print(requestID) // 1
    }
  }
}

So, making an overload of withValue that provides a closure with which the current TaskLocal value. However, if it's on the individual withValue, then it obviously wouldn't compose well with other TaskLocals. Another idea...

@TaskLocal var taskLocalA = 0
@TaskLocal var taskLocalB = 1

$taskLocalA.withValue(37) {
  $taskLocalB.withValue(42) {
    print(taskLocalA, taskLocalB) // 37, 42
    captureTaskLocals { withTaskLocals in
      print(taskLocalA, taskLocalB) // 37, 42
      methodWithEscapingClosure {
        print(taskLocalA, taskLocalB) // 0, 1
        withTaskLocals {
          print(taskLocalA, taskLocalB) // 37, 42
        }
      }
    }
  }
}

This is getting pretty nuts with the levels of nesting though. Ideally there'd be some way to propagate it without introducing further nesting. Maybe using ~Escapable values?

struct TaskLocalScopeFactory {
  func callAsFunction() -> TaskLocalScope
}
struct TaskLocalScope: ~Escapable, ~Copyable {}

@TaskLocal var taskLocalA = 0
@TaskLocal var taskLocalB = 1

$taskLocalA.withValue(37) {
  $taskLocalB.withValue(42) {
    print(taskLocalA, taskLocalB) // 37, 42
    captureTaskLocals { makeScope in
      print(taskLocalA, taskLocalB) // 37, 42
      methodWithEscapingClosure {
        print(taskLocalA, taskLocalB) // 0, 1
        let scope = makeScope()
        print(taskLocalA, taskLocalB) // 37, 42
        // automatically closes scope
      }
    }
  }
}

Is this feasible to implement? Is it advisable? Is there a better way to do this?

2

I'm pretty sure that swift-dependencies package uses that kind of pattern here:

Unless I am misunderstanding (usually am), I think what they're doing here is passing the task-local DependencyValues through to child task. Maybe try it and see if that works for you?

3

Yes, that's what I was kinda inspired by; I guess I'm thinking about a more universal version that copies all TaskLocal values, not just ones explicitly referred to.

1 Like
4

Couldn't you just wrap all those values in a single enum case associated value or a single struct and just pass that?

5

Fundamentally this is something task locals were designed to not do. Let me explain:

Task locals, same as structured concurrency in general, enforce the structure and lifetime of tasks (and task locals). So you know that once you've exited the with {} block, the binding is gone. So that's why task local propagation should not just work in situations that you described. The reason we value this property is issues with thread-locals where you can set one, and forget to ever unset it and you'd have "leaked" the value which may affect some other execution on the same task/thread etc. So this is a property we actively want to maintain.

I do understand the use-case though, and sometimes it's not that you want to keep the binding alive, but you just want to bind the same thing, in a new context -- this "works" because it's a new binding scope.

Currently, you'd have to spell it as a new binding:

$requestID.withValue(1) {
  print(requestID) // 1
  someMethodWithEscapingClosure { [requestID] in  // breaks out of 'structure'
    $requestID.withValue(requestID) { // restore structure
    }
  }
}

On one hand, this i annoying of course; on another though, you're breaking the structured nature of structured concurrency here so there's nothing the runtime can automatically guarantee for you here. We have no idea how long that escaping closure would live.

Or rather; this is not at all about the closure at all; this is about specifically where the closure is invoked:

someMethodWithEscapingClosure(theClosure: () -> ()) {
  theClosure() // cool cool...  same context...

// or
self.theClosure = theClosure
// ...
Task.detached { self.theClosure() } // ofc not structured, no values.

This is what GitHub - apple/swift-service-context: Minimal type-safe context propagation container · GitHub was designed for;

You don't set many individual locals, but set a context and propagate that context propagation container with all the values. This is what swift-distributed-tracing uses under the hood, and you're encouraged to just add more values into it if you'd like to do so.

You just have to keep passing the single ServiceContext, without exploding into tens of TL values.

So in most situations IMHO, the ServiceContext is sufficient when you control all the values.

I guess I'm thinking about a more universal version that copies all TaskLocal values, not just ones explicitly referred to.

Right, so that operation doesn't exist yet; but it could exist; because this is what the runtime does when we say "copy all bindings" to a new unstructuerd Task() which DOES copy value bindings, unlike the Task.detached.

There is another question about the ability to propagate task local values you don't even have the keys for and make new bindings for them... That would have to be a runtime offering because user-land just doesn't know the keys to propagate.

This has come up from time to time, but every time someone did ask for it they had changed their mind a week later. So we've not yet offered this as an API. I'd love to explore the actual use-cases and real world scenarios here, so we can think about how to best solve this.

The general form of this would have to be something like...

let taskLocalBindings = captureAllTaskLocals()
kappa { 
  taskLocalBindings.with { 
  ...
}
}

where it's the same as the existing pattern for a single value but saying "for all".

Perhaps we can make this a bit nicer with a macro or something, but fundamentally it needs to be before the closure creation and not inside it -- because inside it's "too late" to capture the context.

Open to ideas, but would love to hear specific real world scenarios -- because for many the ServiceContext is a pretty good answer already.

2 Likes
6

In the precise example given, it feels more like the problem is that the completion-handler-based API should itself be async (the task-local isn't being explicitly defaulted; it's just being run in a different task). If you convert it,

func someMethodAsync() async {
    let (stream, continuation) = makeAsyncStream(of: Void.self)
    someMethodWithEscapingClosure {
        continuation.finish()
    }
    for await _ in stream {}
    // or ideally handle task cancellation correctly too
}

...

@TaskLocal var requestID = 0

$requestID.withValue(1) {
  print(requestID) // 1
  await someMethodAsync()
  print(requestID) // 1
}
7

Yeah; basically the nature of structured concurrency, and therefore task locals, pushes you to structure your programs this way and doesn't compose well with "executes in some random place and time" escaping callbacks. This is indeed pushing you to correct your design towards more structured concurrency and in a way that's a good thing -- that's how we want modern Swift code to look like.

I do understatnd though that not always you're able to change code, so perhaps a pattern here might be warranted so I'm keeping an open mind still; at the same time agreeing that this friction isn't necessarily entirely bad...