AsyncStream as a Task queue

Using Swift
1

I am trying to use AsyncStream as a queue to use to serialize the calling of certain async methods, is this implementation valid ?

import Foundation

final class AsyncTaskQueue: Sendable {
    typealias TaskWork = @Sendable () async -> Void
    private let continuation: AsyncStream<TaskWork>.Continuation

    init() {
        let (stream, continuation) = AsyncStream<TaskWork>.makeStream()
        self.continuation = continuation

        Task { [stream] in
            for await work in stream {
                await work()
            }
        }
    }

    func enqueue(_ work: @escaping TaskWork) {
        continuation.yield(work)
    }

    deinit {
        continuation.finish()
    }
}
2

Generally yes. I think you also don't strictly need the deinit, since the continuation will finish itself when deallocated.

One thing of note is that you're not handling cancellation, which is fine if that's what you want, but in case you'd like to have cancellation eventually by cancelling the one task created in the init, you have to be aware that it will affect all the future work submitted to the queue, and not one single item.

1 Like
3

Is there a way to only cancel the one task that is canceled without affecting the rest ?

4

Yes, it is possible. You can have a look at the implementation of GitHub - groue/AsyncQueues: Utilities for serializing asynchronous operations · GitHub, which can individually cancel submitted jobs.

There are other queue implementations in the wild, you can search for them. The most recent ones deal with task priority escalation (not the case of the above link).

Gven the complexity of a robust and well-behaved implementation, there should exist a ready-made solution.

1 Like
5

You mean TaskGate ?

6

There's a different technique to serialize tasks, which is to "link up" a previous task with a new one:

func enqueue(work: sending @isolated(any) @escaping () async -> ()) {
    let oldTask = self.task
    self.task = Task {
        let _ = await oldTask.value
        await work()
    }
}

— then you can return the task or store it somewhere else and cancel one-by-one; notably it will not "poison" the whole queue with the cancellation flag (which cannot be unset).

One thing to keep in mind in this case still though is that now, if you have a chain of tasks A -> B -> C, then canceling B will not automatically cancel A. But it can be done by using a withTaskCancellationHandler, which will "recursively" bubble up.

7

Why this technique is different ? if await oldTask.value is cancel how it will not affect the main task ?

And also at this point we will not use AsyncStream no need for it, Right ?

8

Yes, it's different in that it doesn't use an AsyncStream.

Cancelling one unstructured task does not affect any other unstructured tasks, even if they await each other — this is the point of unstructured tasks, they do not propagate several properties such as cancellation or task-local values.

The distinction is subtle, but important: using an AsyncStream and canceling the task that hosts the for await loop will also cancel or outright drop any future work; using a chain like I provided above can allow you to cancel previous work if it's no longer needed.

But again, these concerns only matter if you care about cancellation in the first place. Arguably, it's a good thing to support, but may not matter in your case.

1 Like
9

Thanks man so much for this detailed explanation :heart:

But in genral there is no problem with "link up" a previous task with a new one, it safe and will not cause any issues right ?

10

Depends on what you want to consider an "issue".

Semantically, everything is sound: await oldTask.value waits for the old task to finish before executing work, so this is the exact serialization behavior you want. If oldTask already finished, this is a no-op, and the new task will start running work right away.

Both the stream and the linking techniques technically allow for unbounded latency and memory growth: if you start spawning tasks faster than they manage to complete (and you choose to not cancel them), then you'll run into such kinds of problems, but this can happen with virtually every serial scheduler unless a concrete decision is made on how to prevent that (backpressure? drop old tasks? drop new tasks?).

1 Like
11

So with this implemntation what I need to focus on is task handlation basically to prevent unbounded latency and memory growth.