How to use withTaskCancellationHandler properly?

Using Swift
1

This is something that was asked already for a Combine Publisher type here. But is such a common pattern that I'm still curious what's the best approach.

The question is, when you have a type that takes care of cancelling, but that must be created inside the async operation closure, how you keep it around safely?

 public func asyncFunc() async throws -> FeatureQueryResult {
        var cancellation: TypeThatCancels? // Publisher subscription, network operation...
        return try await withTaskCancellationHandler {
            cancellation?.cancel()
        } operation: {
            try await withCheckedThrowingContinuation { continuation in
                cancellation = StartTheWork() {
                    continuation.resume(with: $0)
                }
            }
        }
    }

The compiler will complain that we're referencing var cancellation from concurrent code, which is of course true.

Looking online the only solutions I've seen is to wrap the cancellation in a class, but is that the correct and intended solution? It feels too fragile. Seems like it just shuts up the compiler but probably doesn't fix the real issues.

Another approach would be to wrap it in an actor but doesn't that feel too much? Also at that point you need to care a lot of timings and check for cancellation inside the operation closure too.

I'm surprised to not find any more official or well explained solutions for such a common scenario. Any advice will be appreciated. Thanks!

Sometimes the workaround is easy, if the object we're using can give us a cancellation token separetly from the closure that gives us the results then we can just move the assignment outside.

 public func asyncFunc() async throws -> FeatureQueryResult {
        var request = someObject.createRequest() // gives us something that can start and also cancel the request
        return try await withTaskCancellationHandler {
            request.cancel()
        } operation: {
            try await withCheckedThrowingContinuation { continuation in
                request.start {
                    continuation.resume(with: $0)
                }
            }
        }
    }

But that is rarely the case.

2 Likes
2

Closures that have @Sendable function type can only use by-value captures. Captures of immutable values introduced by let are implicitly by-value; any other capture must be specified via a capture list:

public func asyncFunc() async throws -> FeatureQueryResult {
        var cancellation: TypeThatCancels? // Publisher subscription, network operation...
        return try await withTaskCancellationHandler { [cancellation] in
            cancellation?.cancel()
        } operation: {
            try await withCheckedThrowingContinuation { continuation in
                cancellation = StartTheWork() {
                    continuation.resume(with: $0)
                }
            }
        }
    }
3

Thanks for the link ^^

Sadly the proposed solution doesn't seem valid. Since you capture by-value it means what is being captured is nil, and thus when the cancellation closure runs it doesn't point to the correct instance set by the start work.

I thought that Sendable would help, but even making the cancellation token sendable doesn't change anything since what the compiler complains about is the var binding itself.

This brings us back to the suggestions I found, wrap the cancellation token in another class so you can have a let and still keep a reference to the token. Seems unfortunate with such a common scenario.

So my original doubts still stand: is wrapping it in a class the best solution? is it safe enough or should we use an actor? Or is there some other trick?

Thanks

4

Hope this example I just coded helps:

func requestPlacemark() async -> CLPlacemark? {
    let geocoder = CLGeocoder()
    return await withTaskCancellationHandler {
        geocoder.cancelGeocode()
    } operation: {
        await withCheckedContinuation { continuation in
            geocoder.reverseGeocodeLocation(location) { placemarks, error in
                continuation.resume(returning: placemarks?.first)
            }
        }
    }
}

Next I plan on handling the returned error, likely converting it to a throw.

5

Thanks but this example is exactly what my last paragraph mentions.

So that's the easy case ^^ I still wonder about the more tricky and common one :slight_smile:

6

I found a workaround for the compilation error here you could try, e.g.

        let onCancel = { dataTask?.cancel() }
        return try await withTaskCancellationHandler {
            onCancel()
        } operation: {
           ...
        }
    }
5 Likes
7

That's actually an interesting one!

The closure performs the same trick as the class: hiding the capture in a reference type that can itself be captured. But with less boilerplate.

Thanks! :clap:

2 Likes
8

Glad to help!

Something very important I also learned the hard way is that the operation is executed even if cancel is called so you have to check if the task is cancelled inside operation. You might not see this in some code examples that use Foundation's async methods because apparently check for cancellation and thus throw exceptions inside them. When calling other completion block-based APIs we do need to check for cancellation inside the operation block, e.g.

    func requestPlacemark(location: CLLocation) async throws -> CLPlacemark? {
        var geocoder: CLGeocoder?
        let onCancel = { geocoder?.cancelGeocode() }
        return try await withTaskCancellationHandler {
            onCancel()
        } operation: {
            try Task.checkCancellation() // need to check if the task was cancelled before it started.
            return try await withCheckedThrowingContinuation { continuation in
            // should we also check cancellation here? I don't know, it seems to be on the same thread called inline with earlier line so maybe not needed.

I had been calling my func from SwiftUI inside .task(id:location) and found that when location changed a few times in quick succession on app resume (from CLLocationManagerDelegate) the tasks were not being cancelled thus many happening simultaneously and resulted in the wrong placemark being shown.

1 Like
9

Yes sure! Thankfully I already knew that from the early days of the proposals so I didn't catch me by surprise :joy: but I admit is yet another trap with the cancellation API.

I usually write two cancellation unit test for code that uses continuations:

  1. Check that cancellation happens when the task is immediately cancelled. This usually uncovers the issue that you mention.
  2. Check that cancellation happens when the task is cancelled with a delay.

I sometimes have this question too but I found that is not needed, at least with the tests I've been doing.

1 Like
10

This now produces a warning with the Xcode 13.3 beta:

Capture of 'onCancel' with non-sendable type '() -> Void?' in a `@Sendable` closure

The best workaround I have come up with so far is to use an unchecked Sendable class to hold the cancellable object:

private class URLSessionDataTaskHolder: @unchecked Sendable {
    var task: URLSessionDataTask?
    func cancel() { task?.cancel() }
}

func data(for request: URLRequest) async throws -> (Data, URLResponse) {
    let dataTaskHolder = URLSessionDataTaskHolder()

    return try await withTaskCancellationHandler(
        handler: {
            dataTaskHolder.cancel()
        },
        operation: {
            try await withCheckedThrowingContinuation { continuation in
                let dataTask = self.dataTask(with: request) { data, response, error in
                    guard let data = data, let response = response else {
                        let error = error ?? URLError(.badServerResponse)
                        return continuation.resume(throwing: error)
                    }

                    continuation.resume(returning: (data, response))
                }

                dataTask.resume()
                dataTaskHolder.task = dataTask
            }
        }
    )
}

Are there better solutions here that I may be missing?

1 Like
11

I also ended up with this approach of using a class wrapper and making it Sendable.

But I don't think your implementation is safe, because you're telling the compiler your URLSessionDataTaskHolder class is Sendable when in fact it isn't. You have a data race because you're accessing dataTaskHolder.task from two different concurrency contexts. I ended up protecting access to to the task property with a lock/mutex.

Additionally, I think you have another race condition: if the cancellation handler runs before the line dataTaskHolder.task = dataTask is executed, your code won't cancel the URLSessionTask.

12

Thanks for calling those points out @ole. Here's an updated (and slightly more generalized) version which should address them:

public extension URLSession {
    func data(for request: URLRequest) async throws -> (Data, URLResponse) {
        let dataTaskHolder = CancellableHolder<URLSessionDataTask>()

        return try await withTaskCancellationHandler(
            handler: {
                dataTaskHolder.cancel()
            },
            operation: {
                try await withCheckedThrowingContinuation { continuation in
                    dataTaskHolder.value = self.dataTask(with: request) { data, response, error in
                        guard let data = data, let response = response else {
                            let error = error ?? URLError(.badServerResponse)
                            return continuation.resume(throwing: error)
                        }

                        continuation.resume(returning: (data, response))
                    }

                    dataTaskHolder.value?.resume()
                }
            }
        )
    }
}

private class CancellableHolder<T: Cancellable>: @unchecked Sendable {
    private var lock = NSRecursiveLock()
    private var innerCancellable: T?

    private func synced<Result>(_ action: () throws -> Result) rethrows -> Result {
        lock.lock()
        defer { lock.unlock() }
        return try action()
    }

    var value: T? {
        get { synced { innerCancellable } }
        set { synced { innerCancellable = newValue } }
    }

    func cancel() {
        synced { innerCancellable?.cancel() }
    }
}

protocol Cancellable {
    func cancel()
}

extension URLSessionDataTask: Cancellable {}

What do you think?

3 Likes
13

We use an actor in AsyncHTTPClient to make it thread safe. A bit more complicated because we also internally schedule a timeout/deadline but it is solving the same problem.

I'm also wondering if there is a better solution. The code example from the Structured concurrency Proposal doesn't actually compile with Xcode 13.2.1 (Swift 5.5.2):

func download(url: URL) async throws -> Data? {
  var urlSessionTask: URLSessionTask?

  return try withTaskCancellationHandler {
    return try await withUnsafeThrowingContinuation { continuation in
      urlSessionTask = URLSession.shared.dataTask(with: url) { data, _, error in
        if let error = error {
          // Ideally translate NSURLErrorCancelled to CancellationError here
          continuation.resume(throwing: error)
        } else {
          continuation.resume(returning: data)
        }
      }
      urlSessionTask?.resume()
    }
  } onCancel: {
    urlSessionTask?.cancel() // runs immediately when cancelled
  }
}

After adding the missing await in front of withTaskCancellationHandler, the compiler complains with the error message: Reference to captured var 'urlSessionTask' in concurrently-executing code

cc: @ktoso

1 Like
14

This looks right but please remember that the dataTaskHolder.cancel() must be implemented in a thread-safe manner, because it might execute from any thread, also concurrently with the store of the dataTaskHolder.value store. So be careful but this, or similar shapes, is how you could do these things today.

There were some discussions to see if we could simplify execution semantics of the handler (make it async) but it ends up in very difficult tradeoffs we're not sure we're willing to take... I offloaded the discussion from my memory but remember it being a world of pain to do any other way :thinking: Maybe @rokhinip remembers the tradeoff this would cost us...

3 Likes
15

I also considered using an actor for synchronization. I ultimately rejected it because it felt wrong when the two contexts from which you call into the actor are synchronous, so you have to wrap the actor calls with Task { … }. Using a lock felt like less overhead to me, but I have no intuition whether that's correct or not.

Yeah, the code example in the proposal is definitely wrong.

3 Likes
16

There were some thoughts around making the cancellation handler async such that it runs synchronous to the task. This way, clients will not have to maintain their own synchronization between the cancellation handler and the task that is being cancelled. However, if we were to do this, the earliest point we could run the cancellation handler, is the next potential suspension point in the task. So if we had something like this:

Task {
     <Synchronous code CPU intensive work>
     await <some async function>
}

If the task was cancelled while it was beginning to execute the CPU intensive operation, the cancellation handler would not be called on the task until it is done with the CPU intensive work. It becomes a question of whether it is acceptable to delay the cancellation handler's execution for a potentially arbitrary period of time.

It is worth noting that the Task will still be able to observe that it has been cancelled using Task.isCancelled and "Synchronous CPU intensive work" ought to periodically check that value and break out of it. This plays further into the fact that cancellation is cooperative with client code.

8 Likes
17

Thanks @rokhinip for the reminder! So much stuff going on that I forgot this… That’d be quite a painful tradeoff and change in execution semantics…

18

That sounds like a fairly intuitive semantic to me actually - if running long bound cpu intensive tasks periodically checking for cancellation is quite reasonable and needed as the alternatives aren’t nice (coming from a pthread_cancel perspective :joy:). (And cancellation at potential await suspension points is perfectly ok)

19

The problem is that some use cases rely on the cancellation handler running to interrupt (set some atomic flag on a shared object that the synchronous work is checking sometimes; or asynchronously cancel some underlying resource (connection)) some work that does NOT suspend anywhere. It would force sprinkling around yields in synchronous code defensively just because “someone might need to wrap this call with a cancellation handler”.

So we’d need to need to become very explicit about sprinkling around ‘await Task.yield()’ if one ever wanted cancellation handlers to run during such code… and that would only be possible in asynchronous code to begin with :worried:

20

The distinction is urgency at which the cancellation work propagates to other pieces of work, like cancelling that connection, etc.

It seems like the two code blocks below are equivalent:

Task {
    while (condition) {
        if (Task.isCancelled) { 
           /* cancel other things like connection objects */ 
           /* short-circuit */
        }
       // Do work
   }
}

vs

withTaskCancellationHandler {
  while (condition) {
     // Do work
  }
} onCancel: {
    /* cancel other things like connection objects */ 
}

but the former runs synchronous to the task, the latter asynchronous to the task. As a result, the immediacy with which "cancel other things" executes changes.

This is intentional which SE-0304 points out:

/// This differs from the operation cooperatively checking for cancellation
/// and reacting to it in that the cancellation handler is _always_ and
/// _immediately_ invoked when the task is cancelled. For example, even if the
/// operation is running code which never checks for cancellation, a cancellation
/// handler still would run and give us a chance to run some cleanup code.
2 Likes