Why `AsyncStream` breaks structured concurrency?

Using Swift
17

Okay, I see this somehow got to the 16th post with no one giving an actual example of how this would be done in true structured concurrency. So here goes:

func withStream<T>(_ body: @Sendable (AsyncStream<Int>) async throws -> T) async rethrows -> T {
    let (stream, continuation) = AsyncStream.makeStream(of: Int.self)

    async let _: () = {
        var n = 0
        while !Task.isCancelled {
            print("Sending", n)
            continuation.yield(n)
            try? await Task.sleep(nanoseconds: 1_000_000_000)
            n += 1
        }
        print("Cancelled!")
     }()

    return try await body(stream)
}

Usage is:

async let outerTask = withStream() { stream in
    for await n in stream {
        print("Received", n)
    }
}

try? await Task.sleep(nanoseconds: 3_500_000_000)
outerTask.cancel()
try? await Task.sleep(nanoseconds: 5_000_000_000)

The use of the with* pattern is less than pretty, but is here to stay until ~Copyable evolves to the point of replacing it. (At a bare minimum ~SyncConsumable is needed in order to be able to properly clean up task)

An alternative is to take a task group as an input:

func makeStream(_ group: TaskGroup<Void>) -> AsyncStream<Int> {
    return AsyncStream { continuation in
        group.addTask {
            var n = 0
            while !Task.isCancelled {
                print("Sending", n)
                continuation.yield(n)
                try? await Task.sleep(nanoseconds: 1_000_000_000)
                n += 1
            }
            print("Cancelled!")
        }
    }
}

Then the usage becomes:

async let outerTask = withTaskGroup(of: Int.Self) { group in
    defer { group.cancelAll() }
    for await n in makeStream(group) {
        print("Received", n)
    }
}

try? await Task.sleep(nanoseconds: 3_500_000_000)
outerTask.cancel()
try? await Task.sleep(nanoseconds: 5_000_000_000)

The general rules are:

  1. Any use of Task { violates structured concurrency. If you want to spawn a child task use either async let or GroupTask.addTask. Task { is basically banned in true structured concurrency.
  2. If you want to return a value that depends on a task running in the background, you need to either pass in a TaskGroup to manage the lifetime of that task, or use a with* structure to ensure the lifetime of the task is limited to a closure.

The key part to structured concurrency isn't the delivery of task cancel. It's the strict definition of task lifetimes. It prevents tasks from being leaked, by requiring tasks adhere to bounded scopes, from which they cannot escape.

The "cheating the structure" solution using continuation.onTermination actually still allows the task to leak: If you call makeStream() but don't perform for await on it, and don't cancel outerTask, the stream's background task will leak (and retain a strong reference to the AsyncStream, via continuation, which should be collected at this point). Structured concurrency makes sure the stream's background task cannot outlive outerTask, even if for await is not used.

Structured concurrency is basically RAII for tasks. Which is why it should one day benefit from true RAII via ~Copyable.

3 Likes