TaskScope.swift

TaskScope.swift

import Synchronization

public final class TaskScope: Sendable {
    public typealias Job = @Sendable () async -> ()

    private let stream: AsyncStream<Job>
    private let continuation: AsyncStream<Job>.Continuation
    private let isRunning = Atomic<Bool>(false)

    public init() {
        (stream, continuation) = AsyncStream.makeStream(bufferingPolicy: .unbounded)
    }

    public func launch(
        @_inheritActorContext _ task: @escaping Job
    ) {
        assert(
            isRunning.load(ordering: .relaxed) == true,
            "TaskScope is not runnig. Did you forget `.compose(viewModel)` or `.task { await taskScope.run() }` ?"
        )
        continuation.yield(task)
    }

    public func run() async {
        let oldIsRunning = isRunning.exchange(true, ordering: .relaxed)
        precondition(oldIsRunning == false, "run() is being called twice simultaneously.")
        await withDiscardingTaskGroup { group in
            for await task in stream {
                group.addTask(operation: task)
            }
        }
        isRunning.store(false, ordering: .relaxed)
    }
}