scope.hpp

Go to the documentation of this file.

#pragma once
 
// CoroutineScope — structured concurrency primitive, modelled after
// Kotlin's `CoroutineScope` / `viewModelScope` and Swift's `TaskGroup`.
//
// Contract (the things a "global-class C++ framework" must guarantee):
//
//   * Every coroutine launched into a scope is *owned* by the scope.
//   * `cancel()` is non-blocking: it requests cancellation, no more.
//   * `cancel_and_join()` (sync) and `co_await join()` (async) wait until
//     every in-flight coroutine has exited. Tests, ViewModel teardown,
//     app shutdown — all use one of these to drain.
//   * The destructor MUST NOT let a launched coroutine outlive the scope.
//     It therefore performs `cancel_and_join()` with a bounded wait
//     (5 s by default). If the wait times out (a coroutine is stuck on
//     a non-cancellable await), the leak is reported through the async
//     error sink — we never block process exit indefinitely.
//   * Unhandled exceptions on the detached path do NOT vanish: any
//     non-`OperationCancelled` exception is forwarded to the async
//     error sink (see <aria/async/async_error_sink.hpp>).
//   * Scopes nest: a child scope constructed with a parent token is
//     automatically cancelled when the parent cancels.
//
// Backward-compatible API: existing `scope.launch(factory)` and
// `scope.launch_simple(task)` calls keep working unchanged.
//
//   CoroutineScope scope;
//   scope.launch([](CancellationToken t) -> Task<void> {
//       while (!t.is_cancelled()) {
//           co_await schedule_on(pool);
//           do_work();
//       }
//   });
//   co_await scope.join();   // wait for everyone to drain (cooperatively)
//
//   // Or, in a synchronous teardown path:
//   scope.cancel_and_join();
 
#include "aria/async/async_error_sink.hpp"
#include "aria/async/cancellation.hpp"
#include "aria/async/task.hpp"
 
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <coroutine>
#include <cstddef>
#include <exception>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <utility>
#include <vector>
 
namespace aria::async {
 
namespace detail {

struct ScopeState {
    std::atomic<std::size_t> inflight{0};
    std::mutex mu;
    std::condition_variable cv;
    // Awaiters waiting for `inflight == 0` (registered by `join()`).
    std::vector<std::function<void()>> drain_waiters;
};
 
}  // namespace detail
 
class CoroutineScope : public std::enable_shared_from_this<CoroutineScope> {
public:
    CoroutineScope() : state_(std::make_shared<detail::ScopeState>()) {}

    explicit CoroutineScope(CancellationToken parent)
        : state_(std::make_shared<detail::ScopeState>()) {
        register_parent_link_(std::move(parent));
    }
 
    ~CoroutineScope() {
        // Structured-concurrency invariant: no coroutine outlives the scope.
        //
        // `cancel_and_join` is itself `noexcept`, but the diagnostic path
        // it routes through (`report_async_error` -> user-installed sink)
        // can in principle throw `bad_alloc` from within the sink's
        // implementation. A destructor is implicitly `noexcept`, so any
        // escaping exception during stack unwinding would call
        // `std::terminate`. Wrap defensively.
        try {
            cancel_and_join();
        } catch (...) {
            // Last line of defence — an exception from the async error
            // sink is itself a (non-fatal) async error; we can't report
            // it back through the same channel without recursing, so
            // we silently drop. The leak (if any) was already accounted
            // for via `inflight_` reads.
        }
    }
 
    CoroutineScope(const CoroutineScope&) = delete;
    CoroutineScope& operator=(const CoroutineScope&) = delete;
    CoroutineScope(CoroutineScope&&) = delete;
    CoroutineScope& operator=(CoroutineScope&&) = delete;
 
    // ── Inspection ────────────────────────────────────────────────────
 
    [[nodiscard]] CancellationToken token() const noexcept { return src_->token(); }
    [[nodiscard]] bool is_cancelled() const noexcept { return src_->is_cancelled(); }

    [[nodiscard]] std::size_t inflight_count() const noexcept {
        return state_->inflight.load(std::memory_order_acquire);
    }
 
    // ── Cancellation / join ───────────────────────────────────────────

    void cancel() noexcept {
        auto source = src_;
        try { source->cancel(); } catch (...) {}
    }

    bool cancel_and_join(std::chrono::milliseconds timeout =
                             std::chrono::milliseconds{5000}) noexcept {
        cancel();
        std::unique_lock lk(state_->mu);
        const bool drained = state_->cv.wait_for(lk, timeout, [this] {
            return state_->inflight.load(std::memory_order_acquire) == 0;
        });
        if (!drained) {
            const auto leaked =
                state_->inflight.load(std::memory_order_acquire);
            // A bounded synchronous wait must not complete asynchronous
            // joiners early. Their shared state remains alive until the last
            // task exits and resumes them through decrement_inflight_.
            lk.unlock();
            report_async_error(
                std::string("CoroutineScope: dtor leaked ") +
                std::to_string(leaked) +
                " task(s) (cancellation observed but coroutines did not "
                "exit within the timeout)");
        }
        return drained;
    }

    struct JoinAwaiter {
        std::shared_ptr<detail::ScopeState> st;
 
        bool await_ready() const noexcept {
            return st->inflight.load(std::memory_order_acquire) == 0;
        }

        bool await_suspend(std::coroutine_handle<> h) {
            std::unique_lock lk(st->mu);
            if (st->inflight.load(std::memory_order_acquire) == 0) {
                return false;// resume in place, on the caller's frame
            }
            st->drain_waiters.emplace_back([h]() mutable { h.resume(); });
            return true;
        }
 
        void await_resume() const noexcept {}
    };

    [[nodiscard]] JoinAwaiter join() noexcept {
        cancel();
        return JoinAwaiter{state_};
    }

    [[nodiscard]] JoinAwaiter join_existing() noexcept {
        return JoinAwaiter{state_};
    }
 
    // ── Launch ────────────────────────────────────────────────────────

    template<typename Fn>
    void launch(Fn&& factory) {
        using FnDecay = std::decay_t<Fn>;
        spawn_tracked_(
            launch_owner_coro_<FnDecay>(std::forward<Fn>(factory), token()));
    }

    void launch_simple(Task<void> task) {
        spawn_tracked_(std::move(task));
    }
 
private:
    // ── Tracked spawn ─────────────────────────────────────────────────

    static void decrement_inflight_(
        const std::shared_ptr<detail::ScopeState>& st) noexcept {
        if (!st) return;
        if (st->inflight.fetch_sub(1, std::memory_order_acq_rel) == 1) {
            // Last one out: fire drain waiters and notify cv.
            std::vector<std::function<void()>> waiters;
            {
                std::lock_guard lk(st->mu);
                waiters.swap(st->drain_waiters);
            }
            st->cv.notify_all();
            for (auto& w : waiters) {
                try { w(); } catch (...) {}
            }
        } else {
            // Not the last; still wake any timed waiters that may
            // want to observe progress (cheap).
            st->cv.notify_all();
        }
    }

    template<typename Fn>
    static Task<void> launch_owner_coro_(Fn factory, CancellationToken tok) {
        co_await factory(std::move(tok));
    }

    void spawn_tracked_(Task<void> task) {
        // Snapshot the shared accounting block — wrapper holds it by
        // shared_ptr so even if the scope object is destroyed first,
        // the bookkeeping completes safely.
        auto st = state_;
        st->inflight.fetch_add(1, std::memory_order_acq_rel);
        // Detached driver coroutine — captures `st` by value. If
        // building or starting the driver throws (e.g. bad_alloc when
        // the coroutine frame can't be allocated), we must roll back
        // the inflight increment so the scope can still be drained.
        try {
            spawn_driver_(std::move(task), std::move(st)).start_detached_();
        } catch (...) {
            decrement_inflight_(state_);
            throw;
        }
    }
 
    static Task<void> spawn_driver_(Task<void> body,
                                    std::shared_ptr<detail::ScopeState> st) {
        // Local guard: even if the body throws synchronously before its
        // first co_await (which a properly written Task should not, but
        // we don't want to rely on it), the inflight counter still
        // decrements via the destructor.
        struct InflightGuard {
            std::shared_ptr<detail::ScopeState> st;
            ~InflightGuard() { CoroutineScope::decrement_inflight_(st); }
        };
        InflightGuard guard{std::move(st)};
        try {
            co_await std::move(body);
        } catch (const OperationCancelled&) {
            // Expected on scope.cancel(); do nothing.
        } catch (const std::exception& e) {
            report_async_error(
                std::string("CoroutineScope: unhandled exception in launched task: ") +
                e.what());
        } catch (...) {
            report_async_error(
                "CoroutineScope: unhandled non-std exception in launched task");
        }
        co_return;
    }
 
    // ── Parent linkage ────────────────────────────────────────────────

    void register_parent_link_(CancellationToken parent) {
        // Keep the source alive for the entire broadcast, even if a child
        // cancellation callback destroys the scope. No proxy lock may be held
        // while invoking callbacks supplied by the application.
        parent.on_cancel([weak = std::weak_ptr<CancellationSource>(src_)] {
            if (auto source = weak.lock()) source->cancel();
        });
    }
 
    std::shared_ptr<CancellationSource> src_ =
        std::make_shared<CancellationSource>();
    std::shared_ptr<detail::ScopeState> state_;
};

inline auto operator co_await(CancellationToken tok) {
    struct Latch {
        enum State : int {
            preparing = 0,
            suspended = 1,
            cancellation_before_suspend = 2,
            resumed = 3,
        };
 
        std::atomic<int> state{preparing};
    };
    struct Awaiter {
        CancellationToken tok;
        std::shared_ptr<Latch> latch;
 
        bool await_ready() const noexcept { return tok.is_cancelled(); }
 
        bool await_suspend(std::coroutine_handle<> h) {
            auto l = latch;
            // Register the cancellation callback. It may fire
            // synchronously right here if `tok` is already cancelled,
            // or asynchronously later from another thread. Synchronous
            // firing while we are still preparing records cancellation
            // without resuming on this stack.
            tok.on_cancel([h, l]() mutable {
                int observed = l->state.load(std::memory_order_acquire);
                for (;;) {
                    if (observed == Latch::preparing) {
                        if (l->state.compare_exchange_weak(
                                observed,
                                Latch::cancellation_before_suspend,
                                std::memory_order_acq_rel,
                                std::memory_order_acquire)) {
                            return;
                        }
                        continue;
                    }
                    if (observed == Latch::suspended) {
                        if (l->state.compare_exchange_weak(
                                observed,
                                Latch::resumed,
                                std::memory_order_acq_rel,
                                std::memory_order_acquire)) {
                            // Defer the resume: see contract (1)/(2)
                            // above. The topmost cancel() on this
                            // thread will drain the pending list once
                            // all callbacks return.
                            detail::schedule_deferred_resume(h);
                            return;
                        }
                        continue;
                    }
                    return;
                }
            });
 
            // Commit to suspension only if no cancellation callback won
            // while we were registering it. If the callback already moved
            // the state to `cancellation_before_suspend`, returning false
            // lets the compiler resume safely on the caller's stack.
            int expected = Latch::preparing;
            return l->state.compare_exchange_strong(
                expected,
                Latch::suspended,
                std::memory_order_acq_rel,
                std::memory_order_acquire);
        }
 
        // Intentionally non-throwing. The body must follow up with an
        // explicit `tok.throw_if_cancelled()` probe (see the class
        // comment for why we don't throw here).
        void await_resume() noexcept {}
    };
    return Awaiter{std::move(tok), std::make_shared<Latch>()};
}
 
}  // namespace aria::async