`withTaskCancellationHandler` seemingly introduces suspension point

Using Swift
1

I have a couple assumptions about what I am trying to achieve here (which I think are confirmed by either documentation or forum posts):

  1. withUnsafeContinuation synchronously executes its body, without introducing any suspension points before running it. From the docs:

The body of the closure executes synchronously on the calling task, and once it returns the calling task is suspended. It is possible to immediately resume the task, or escape the continuation in order to complete it afterwards, which will then resume the suspended task.

  1. withTaskCancellationHandler doesn't suspend or change isolation unless the closure contained does. From the docs:

The operation closure executes on the calling execution context, and doesn’t suspend or change execution context unless code contained within the closure does so. In other words, the potential suspension point of the withTaskCancellationHandler(operation:onCancel:) never suspends by itself before executing the operation.

  1. When there is no switch in isolation, there is no suspension point. Was confirmed in this post: Suspension guarantees when isolation does not change
  2. @isolated(any) and therefore TaskGroup.addTask(_:) + isolated closure guarantees that work is scheduled in order of calling addTask. Kind of mentioned in the proposal: swift-evolution/proposals/0431-isolated-any-functions.md at main · swiftlang/swift-evolution · GitHub

However, I recently wrote something similar to this:

@Test
func testSomething() async {
    
    let lock = OSAllocatedUnfairLock(initialState: 0)
    
    func shouldNotSuspend(
        value: Int,
        isolation: isolated (any Actor)? = #isolation
    ) async {
        await withTaskCancellationHandler {
            await withUnsafeContinuation { continuation in
                lock.withLock { $0 = value }
                continuation.resume()
            }
        } onCancel: { }
    }
    
    await withTaskGroup(of: Void.self) { taskGroup in
        taskGroup.addTask { @Sendable @MainActor in
            await shouldNotSuspend(value: 1)
        }
        taskGroup.addTask { @Sendable @MainActor in
            await shouldNotSuspend(value: 2)
        }
        await taskGroup.waitForAll()
    }
    
    lock.withLock { value in
        #expect(value == 2)
    }
}

If you run this test a couple of times (I mostly just do 100 runs), it will fail at some point, saying that value is 1 instead of 2. Now, the title mentions that the problem seems to be withTaskCancellationHandler. This is because, if you'd remove withTaskCancellationHandler and just keep withUnsafeContinuation, you can run the test 1000 times and it will not fail.

Would someone mind explaining, what I could be doing wrong? Thank you.

2

Doing even more tests, I suspect that withTaskCancellationHandlers operation does not inherit isolation as I expected it to. The closure is called from the global executor pool and therefore introduces a suspension point.

Nested async closures are hard.

3

You will need to forward the isolation to the inner continuation:

await withTaskCancellationHandler {
   await withUnsafeContinuation(isolation: isolation) {
   ...
  }

I do this within IdentifableContinuation and it works well.

4 Likes
4

Ah, forgot about that. I had:

await withTaskCancellationHandler {
  _ = isolation
 await withUnsafeContinuation { }
}

... as a workaround, but passing it along looks cleaner.

1 Like
5

I haven’t looked at this closely enough, but it seems like perhaps the compiler should be able to reason that the static isolation does not change here?

6

This might be even spicier than I initially thought. This seems to break isolation assumptions:

func baz(
    isolation: isolated (any Actor)? = #isolation,
    closure: (UnsafeContinuation<Void, Never>) -> Void
) async {
    await withTaskCancellationHandler {
        await withUnsafeContinuation { continuation in
            closure(continuation)
        }
    } onCancel: {
        
    }
}

@MainActor
class Foo {
    
    var isolated = "isolated"
    
    func bar() async {
        await baz { continuation in
            isolated = "mutated"
            continuation.resume()
        }
    }
}

Task {
    let foo = Foo()
    await foo.bar()
}

It crashes because isolated will be mutated on a non-MainActor isolation. At this point I think that this is actually a bug, somehow.

It works if I do this instead:

func baz(
    isolation: isolated (any Actor)? = #isolation,
    closure: (UnsafeContinuation<Void, Never>) -> Void
) async {
    await withTaskCancellationHandler {
        await withUnsafeContinuation(isolation: isolation) { continuation in
            closure(continuation)
        }
    } onCancel: {
        
    }
}

Any opinions/objections?

EDIT: It works on newest nightly.