For-await during cancelling a Task

Using Swift
,
1

I've noticed, that when a Task is cancelled while the work is suspended (on some for-await), awaiting is "interrupted" and work is continued, event if the stream is constantly sending values (not in a example for simplicity).
Here is a simple code that allow to recreate this behaviour.

import SwiftUI
import ConcurrencyExtras

struct QuestionTaskView: View {
    @State var visible: Bool = false
    var body: some View {
        Button(action: {
            visible.toggle()
        }) {
            Text("Toggle")
        }
        if visible {
            Text("Task")
                .task {
                    print("Task started")
                    let stream = AsyncStream<Int> { _ in } // stream that never ends
                    for await _ in stream {
                    }
                    print("Task ended")
                }
        }
    }
}

#Preview {
    QuestionTaskView()
}

After tapping twice on the Toggle button (to start and cancel the Task) this code prints in the console:

Task started
Task ended

I did not find any documentation about how cancelling a Task is dealing with suspension points. It is logical that a Task should allow to continue in order to clean up/finish, without this interruption Task would wait indefinitely. Where I can find more about which awaits are handled this way?

2

As I understand it, only 2 things happen when a task gets cancelled:

  • The isCancelled flag inside the task is set. From this point, any future call to Task.isCancelled will return true, and any future call to Task.checkCancellation() will throw.

  • The concurrency runtime calls all cancellation handlers that have been registered with the task (using withTaskCancellationHandler).

That's it. So the fact that the task is currently suspended inside an await is not handled directly.

In your example, AsyncStream's iterator type handles cancellation by wrapping the body of its next() function in a cancellation handler (the relevant source code in the stdlib). This means the runtime will notify the iterator when the iterating task gets cancelled, and the iterator reacts to it by stopping and returning nil, which ends the loop.

Note that this behavior is implemented by AsyncStream. Other AsyncSequences might handle cancellation differently. (Although all Swift Concurrency code is expected to react to cancellation by finishing as soon as possible.)

2 Likes
3

Thank you, now I understand!

1 Like
4

It is expected that all AsyncSequence conforming types will terminate eventually on cancellation. Some may not do it immediately (i.e. via checking Task.isCancelled and not withTaskCancellationHandler) but contractually it is expected to terminate at some point in time because else wise it is a leak.

2 Likes
5

The documentation for AsyncInteratorProtocol says:

Cancellation

Types conforming to AsyncIteratorProtocol should use the cancellation primitives provided by Swift’s Task API. The iterator can choose how to handle and respond to cancellation, including:

  • Checking the isCancelled value of the current Task inside next() and returning nil to terminate the sequence.
  • Calling checkCancellation() on the Task, which throws a CancellationError.
  • Implementing next() with a withTaskCancellationHandler(handler:operation:)invocation to immediately react to cancellation.

If the iterator needs to clean up on cancellation, it can do so after checking for cancellation as described above, or in deinit if it’s a reference type.

So, when writing your own AsyncSequence you need to make sure your iterator responds to cancelation.

But in this case, you are using the concrete AsyncStream, which correctly handles cancelation for you. When a task that is iterating through the sequence is canceled, the AsyncStream is automatically canceled, too. And you are using the .task view modifier, for which “SwiftUI will automatically cancel the task at some point after the view disappears before the action completes.”

So, the task associated with your .task view modifier will be canceled when the Text("Task") disappears as visible is toggled. And AsyncStream responds to this cancelation, thereby stopping the sequence.

Now, the question is how one terminates an AsyncStream that is “constantly sending values”. The key observation is that you need an onTermination clause that will stop the work that is sending the values. You did not give an example of your stream, but consider this, which sends integer values, with one second delay between each:

let sequence = AsyncStream<Int> { continuation in
    let task = Task {
        do {
            for i in 0 ..< 1_000 {
                try await Task.sleep(for: .seconds(1))
                continuation.yield(i)
            }
        } catch { }
        continuation.finish()
    }
    
    continuation.onTermination = { state in
        if case .cancelled = state {
            task.cancel()
        }
    }
}

So, the onTermination will determine if the sequence was canceled (in which case we have to cancel our task yielding all these values) or not (i.e., it finished in the course of its natural lifecycle, in the above example, if it successfully yielded the 1,000 values and finished).

Now, there are lots of variations on the theme, but the idea is that when the stream is terminated (which happens when the parent task is canceled), you need to stop your process that is sending all these values.

2 Likes