Task Cancellation Handlers

In the first post on task cancellation, I showcased a few simple examples of manual cancellation handling using Task.isCancelled and Task.checkCancellation. Using these functions is passive (or reactive) cancellation support: If your code happens to pass over such a check, it cancels. But if not, there might be a long delay before it notices it’s been cancelled. In some cases, a task might never pass over such a check because it’s suspended waiting for something.

This is where cancellation handlers come into play. As the name suggests, you can register a handler closure that gets called on task cancellation. Sounds easy? It looks easy:

await withTaskCancellationHandler {
  // Primary operation
} onCancel: {
  // Handle cancellation
}

But is it easy? Not really. The signature of the cancellation handler is @Sendable () -> Void—a synchronous sendable non-throwing non-isolated closure with no arguments and no return value. Which means:

• You cannot “return” or “throw” out of the main operation from the handler; you can only influence its course of execution.
• To achieve that, you must use only sendable (thread-safe) constructs.
• And these constructs must be synchronously accessible.

Depending on your experience with structured concurrency and strict sendability checking, you might already know this isn’t trivial. Let’s make sure we’re on the same page:

Anything we use to communicate between the handler and the operation needs to be available outside the call. The handler cannot, for example, access state created in the operation. A closure that’s supposed to influence the course of execution of another closure has no access to anything created in there. We’ll see an example and a solution shortly.

Second, the typical tools for attaining a sendable entity are severely limited. Whatever you share between the two closures cannot be an actor or isolated on a global actor (like @MainActor). The cancellation handler is synchronous and non-isolated. We can’t make asynchronous calls with await, and we can’t know to be on the right isolation to access these types synchronously.

The only standard library construct that fulfills all these requirements out of the box is Task itself. We’ll see an example shortly. But for anything else, you have no choice but to go to lower-level mechanisms…like locks. Yes, the thing you thought you’d never need again with structured concurrency—well,

at least I did—it’s still around.

Luckily, since Swift 6, there’s the new Synchronization framework. It contains helper types that are primitive and low-level but work great with strict sendability checking and our scenario here—they’re synchronous and sendable. The type you’ll most likely use here is Mutex, a container that provides thread-safe mutually exclusive read-write access to some value.

A lot to take in? 😅 Let’s look at some examples.

The simple example

Let’s start with Task to see why it’s such a good fit for cancellation handling.

What follows is a bit contrived. You usually won’t have a Task created and passed around just for a single client. It’s not wrong, but it has a smell to it, as there are simpler tools available. But that’s a story for another post. For now, it’s illustrative because it’s simple.

Say you have a task doing some work. To keep things simple, it supports cancellation by returning nil:

let workTask: Task<Value?, Never> = …

At some point, you’re awaiting the workTask’s result to update the interface:

let result = await workTask.value

The task was created just for this one client. When the interface ends before the task finishes, the task result is no longer needed and the task can be cancelled.

Using the task cancellation handler in this case is straightforward:

let result = await withTaskCancellationHandler {
  await workTask.value
} onCancel: {
  workTask.cancel()
}

if let result {
  // Do the update
}

Let’s check the requirements: workTask is created outside the cancellation block, so it can be available inside both closures. It’s a sendable reference type, so it can be used in both closures (a reference to it is copied to each closure). Task’s cancel() function is synchronous and non-isolated, so it can be used in the cancellation handler. A perfect match!

The less simple example

I’ve been struggling to come up with a useful example that’s both short and not contrived. I have one in production, but it’s not short and requires some scene setting. We’ll save that for later.

For now, we’ll build an academic sample with little (or no) practical use: a function that suspends until cancelled. untilCancelled() has no return value and doesn’t throw. It just hangs in there and blocks the task until it’s cancelled. It might have some utility as a test helper, but I’m not going to argue that.

What makes it interesting is that we’ll use a CheckedContinuation for the implementation. It’s mainly a language primitive to bridge between synchronous and asynchronous code, but it’s helpful in other cases as well. A continuation suspends a task on one side and provides a handle to end that suspension on the other. Here’s a minimal example that doesn’t actually suspend because it immediately resumes:

await withCheckedContinuation { continuation in
  continuation.resume()
}

Continuations can have values and throw errors, but we don’t need these here. What makes it compelling is that a continuation is a sendable, copyable struct that suspends an asynchronous context and can be unlocked synchronously. You might recognize these trigger words—it sounds like a perfect match for what we require. And it is! With one challenge: The continuation handle is provided as an argument to the closure only.

Let’s sketch out the untilCancelled function to see the challenge. We need to support cancellation and want to use a checked continuation. The closure in withCheckedContinuation is always synchronous, so the only way these two fit together is like this:

func untilCancelled() async {
  // 1
  
  await withTaskCancellationHandler {
    await withCheckedContinuation { continuation in
      // 2
    }
  } onCancel: {
    // 3
  }
}

The challenge: the continuation is available only at (2), but we need to resume it at (3). Remember, anything that has to be used in both closures must be defined outside. This means we require something defined at (1) to share state between these two places.

Luckily, continuation handles can escape from the closure, they can be stored somewhere. They’re also sendable, which makes them safe to use from any thread. The code structure requires us to set up the state first and only then get access to and store the handle. We need something that’s mutable and optional so it can store a continuation once known. The simplest idea might be something like this:

var cont: CheckedContinuation<Void, Never>?

Unfortunately, variables aren’t thread-safe and thus not sendable. You could store the continuation in a variable, but the compiler won’t allow access to it from the cancellation handler. But we can wrap the variable in a mutex to work around that:

let mutex: Mutex<CheckedContinuation<Void, Never>?>

If you haven’t used it before, Mutex has one primary function withLock that provides a passed closure with exclusive access to the contained value:

mutex.withLock { value in
  // protected read/write access to `value`
}

With all pieces in place, let’s implement our function:

func untilCancelled() async {
  let mutex: Mutex<CheckedContinuation<Void, Never>?> = .init(nil) // 1
  
  await withTaskCancellationHandler { // 2
    await withCheckedContinuation { continuation in // 3
      mutex.withLock { $0 = continuation } // 4
    }
  } onCancel: {
    let continuation = mutex.withLock { $0.take() } // 5
    continuation?.resume() // 6
  }
}

(1) First, we set up the state to hold the continuation. Initially, the continuation isn’t known, so it starts with nil. In (2) we begin cancellation handling and in (3) we create the continuation that the task is suspended on. In (4), we access the state’s mutex and store the continuation for later retrieval. The task will get suspended waiting on the continuation. When cancellation happens, in (5) we retrieve the continuation and in (6) resume it to end the function.

As I said, it’s an academic example with little use in practice. But if you want to try it, here’s how you’d set up a task that hangs forever—_unless_ canceled:

let hangsForever = Task {
  await untilCancelled()
}
…
hangsForever.cancel()

The task cancellation handler trap

If you try this code, it works fine if you let the caller sleep for a moment (during the …) before cancellation. However, if you call cancel immediately, the Task will never end. What’s up with that?

There’s one more aspect to withTaskCancellationHandler that we’re missing. The documentation says (emphasis mine):

The operation closure is always invoked, even when […] called from a task that was already canceled.
When […] used in a task that has already been canceled, the cancellation handler will be executed immediately before the operation closure gets to execute.

So what happens if the function is called from an already cancelled task? The state is set up as usual in (1) and then cancellation handling set up in (2). But instead of continuing into the operation, the cancellation handler (5) will be called first. The state is still empty, the retrieved continuation will be nil, and so the resume call (6) will be skipped. Only then, the continuation is set up (3) and stored (4). And there it will hang…forever. Like, forever-forever. Because the task is already cancelled and the handler has already been invoked, and cancellation handlers are only invoked at most once, nothing can now reach the continuation.

There are two ways to fix this: We could have some additional state to denote early cancellation. Or—which I find simpler—we could manually check for cancellation in the operation. Which brings us to this final implementation:

func untilCancelled() async {
  let mutex: Mutex<CheckedContinuation<Void, Never>?> = .init(nil)
  
  await withTaskCancellationHandler {
    // 1
    await withCheckedContinuation { continuation in
      // 2
      mutex.withLock {
        // 3
        if Task.isCancelled {
          continuation.resume()
          return
        }
        $0 = continuation
      }
    }
  } onCancel: {
    let continuation = mutex.withLock { $0.take() }
    continuation?.resume()
  }
}

Note how we do the cancellation check only inside the mutex access (3). This avoids race conditions where cancellation might happen concurrently with the continuation setup. When checking before withCheckedContinuation (1) or before withLock (2), in rare cases cancellation might happen right between the check and the mutex access, exposing the same lost cancellation behavior as before. It’s rare, but not impossible.

At that point (3), we also can’t just return from the closure. The continuation handle has already been provided, so we need to resume it. We also avoid storing it. While not strictly necessary here, continuations must be resumed exactly once. Not keeping a reference after it’s been resumed is one effective way to ensure this consistency.

The real world

So that is a simple function that suspends until cancelled. 😅 Only that, in my opinion, it’s not simple at all—it’s really complex with numerous pitfalls small and large. Should this keep you from implementing cancellable tasks? It shouldn’t, but you should sharpen your senses.

I found building this dummy function and playing around with it quite insightful. My initial post written for our team internally, for example, missed the aspect of early cancellation. I encourage you to do the same and try it out. I’ve posted a slightly extended version on GitHub that’s ready to run.

Ideally, frameworks would come with more cancellation support built in. Until they do, I recommend wrapping these complex intricacies in helper types that are less complex to use. Above, I mentioned a real-world example, which happens to be such a helper type. It needs a bit more scene setting, so I’ll leave this to a separate post.

Thank you for reading, and let me know what you think.