completion_awaiter.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
 
#pragma once
#include <atomic>
#include <core/async/Awaitable.hpp>
#include <core/async/Cancellation.hpp>
#include <core/async/ExecutorContext.hpp>
#include <core/async/ParkedWork.hpp>
#include <core/async/StopToken.hpp>
#include <coroutine>
#include <cstdint>
#include <exception>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>
 
#include "../callback_scope.hpp"
#include "../executor.hpp"
 
/// @file
/// @brief The awaiter behind `Completion<T>::operator co_await`.
///
/// A fragment of `completion.hpp` rather than of `coroutine.hpp`, so a
/// `Completion<T>` is awaitable wherever it is visible. Specified in
/// `docs/spec/core/coroutines.md`.
 
namespace morph::async::detail {
 
template <typename T>
struct CompletionState;
 
/// @brief The awaiter `Completion<T>::operator co_await() &&` returns.
///
/// The await is one more `then`/`onError` pair on the completion, holding only
/// a heap `Shared` -- never the coroutine frame -- and gated by a
/// `CallbackToken`. With a stoppable token on the awaiting promise it also
/// registers a stop callback. Exactly one of *the completion's handler* and
/// *the stop callback* moves `Shared::outcome` off `Pending`, and only that one
/// resumes the coroutine. A stop that withdraws the await also stops the call
/// the completion reports on, when that call carries a stop source.
///
/// The pair is attached on the completion's executor, which owns its handler
/// lists: directly when the awaiting coroutine is running there, otherwise by
/// a task posted there. So any coroutine may await any completion.
///
/// `await_suspend` touches only local copies of the shared pointers once it has
/// published anything another thread could act on: from then on the coroutine
/// may already have been resumed elsewhere, and the awaiter with its frame
/// destroyed.
/// @tparam T The completion's value type; copied out of the settled state.
template <typename T>
class CompletionAwaiter {
    enum class Outcome : std::uint8_t { Pending, Settled, Cancelled };
 
    struct Shared;
 
    /// The stop callback's callable: forwards to `Shared::onStop`.
    struct OnStop {
        Shared* shared;
        void operator()() const noexcept { shared->onStop(); }
    };
 
    struct Shared {
        std::atomic<Outcome> outcome{Outcome::Pending};
        /// Set once `await_suspend` has finished registering; before that a
        /// stop only records itself and `await_suspend` answers it.
        std::atomic<bool> armed{false};
        /// Taken by whichever of `await_suspend` and `onStop` answers a stop.
        std::atomic<bool> claimed{false};
        std::optional<T> value;
        std::exception_ptr error;
        std::coroutine_handle<> continuation;
        /// Where the coroutine was running when it suspended: resumed there.
        /// Empty outside every executor's task.
        ::core::async::ResumeTarget context;
        ::morph::exec::IExecutor* fallback = nullptr;
        CallbackScope scope;
        std::optional<::core::async::StopCallback<OnStop>> stopCallback;
 
        // seq_cst, as `armed` is: whoever decides and then reads `armed` must
        // not miss the other side's store (a store-buffering pattern).
        bool decide(Outcome outcome_) noexcept {
            auto expected = Outcome::Pending;
            return outcome.compare_exchange_strong(expected, outcome_, std::memory_order_seq_cst);
        }
 
        /// A settled completion, on the completion's executor: resumes in the
        /// suspending context, or right here.
        void resumeSettled() {
            if (context) {
                context.submit(::core::async::ParkedWork{.resume = continuation});
            } else {
                continuation.resume();
            }
        }
 
        /// A stop, on whichever thread requested it: resumes through the
        /// suspending context, or on the completion's executor. Not inline
        /// here, because the requesting thread is not the coroutine's; the
        /// context itself may resume inline (a Task handler's resumer whose
        /// strands are closed does).
        void resumeCancelled() {
            if (context) {
                context.submit(::core::async::ParkedWork{.resume = continuation});
            } else {
                fallback->post([handle = continuation] { handle.resume(); });
            }
        }
 
        void onStop() noexcept {
            if (!decide(Outcome::Cancelled)) {
                return;
            }
            scope.requestStop();
            if (armed.load(std::memory_order_seq_cst) && !claimed.exchange(true, std::memory_order_acq_rel)) {
                try {
                    resumeCancelled();
                } catch (...) {  // NOLINT(bugprone-empty-catch): a stop callback must not throw
                    // A queue that cannot accept one more closure leaves nothing to
                    // resume through; the coroutine stays suspended until its
                    // owner destroys it.
                }
            }
        }
    };
 
public:
    /// @param state The completion's state; consumed. Null for an empty completion.
    explicit CompletionAwaiter(std::shared_ptr<CompletionState<T>> state)
        : _state{std::move(state)}, _shared{std::make_shared<Shared>()} {}
 
    /// Withdraws the await if the frame is destroyed while suspended, so neither
    /// a later settlement nor a later stop resumes a frame that is gone.
    ~CompletionAwaiter() {
        _shared->stopCallback.reset();
        _shared->scope.requestStop();
    }
 
    CompletionAwaiter(const CompletionAwaiter&) = delete;
    CompletionAwaiter& operator=(const CompletionAwaiter&) = delete;
    CompletionAwaiter(CompletionAwaiter&&) = delete;
    CompletionAwaiter& operator=(CompletionAwaiter&&) = delete;
 
    /// @return Always false: every decision is `await_suspend`'s, so this stays
    ///         a constant (see core-cpp's `awaitReadyIsConstantFalse`).
    [[nodiscard]] constexpr bool await_ready() const noexcept { return false; }
 
    /// @tparam Promise The awaiting coroutine's promise type.
    /// @param awaiting The coroutine performing the `co_await`.
    /// @return False to continue at once (an unusable completion, or a stop
    ///         already requested); true once the await is attached.
    template <typename Promise>
    // NOLINTNEXTLINE(readability-function-cognitive-complexity): the whole attach decision lives here by design (see await_ready); each branch is one way the await ends or suspends.
    bool await_suspend(std::coroutine_handle<Promise> awaiting) {
        auto shared = _shared;
        auto state = _state;
        if (state == nullptr || state->cbExec == nullptr) {
            shared->error = std::make_exception_ptr(std::logic_error{
                state == nullptr ? "co_await on an empty morph::async::Completion (default-constructed or moved-from)"
                                 : "co_await on a morph::async::Completion with no callback executor to resume on"});
            shared->outcome.store(Outcome::Settled);
            return false;
        }
        shared->continuation = awaiting;
        shared->context = ::core::async::ResumeTarget::current();
        shared->fallback = state->cbExec;
 
        if constexpr (::core::async::HasStopToken<Promise>) {
            ::core::async::StopToken const token = awaiting.promise().stopToken();
            if (token.stop_requested()) {
                shared->outcome.store(Outcome::Cancelled);
                stopUnattached(state, token);
                return false;
            }
            if (token.stop_possible()) {
                shared->stopCallback.emplace(token, OnStop{shared.get()});
                shared->armed.store(true, std::memory_order_seq_cst);
                if (shared->outcome.load(std::memory_order_seq_cst) == Outcome::Cancelled) {
                    // A stop landed while registering. Whoever claims answers
                    // it; if onStop claimed first it is resuming us already.
                    stopUnattached(state, token);
                    return shared->claimed.exchange(true, std::memory_order_acq_rel);
                }
            }
        }
 
        try {
            if (::morph::exec::runningOn(*state->cbExec)) {
                attach(*state, shared);
            } else {
                // Attaching is the owner's: post it there. The executor's
                // queue then orders it with the settle's delivery, and the
                // handlers it attaches resume the coroutine as above.
                state->cbExec->post([state, shared] {
                    try {
                        attach(*state, shared);
                    } catch (...) {
                        if (shared->decide(Outcome::Settled)) {
                            shared->error = std::current_exception();
                            shared->resumeSettled();
                        }
                    }
                });
            }
        } catch (...) {
            // The throw resumes the coroutine, so nothing else may. If a stop
            // has claimed the await it is resuming the coroutine already, and
            // the frame is no longer this call's to throw into: that stop's
            // OperationCancelled is what the coroutine sees.
            if (shared->claimed.exchange(true, std::memory_order_acq_rel)) {
                shared->scope.requestStop();
                return true;
            }
            // Likewise if the handler attached before the throw has already
            // settled the await: it resumed the coroutine.
            if (!shared->decide(Outcome::Settled) && shared->outcome.load() == Outcome::Settled) {
                shared->scope.requestStop();
                return true;
            }
            // Still suspended, and decided: withdraw the stop callback (waiting
            // out one that is running) and the handler attached before the throw.
            shared->stopCallback.reset();
            shared->scope.requestStop();
            throw;
        }
        return true;
    }
 
    /// @return The settled value.
    /// @throws core::async::OperationCancelled if a stop withdrew the await.
    /// @throws Whatever the completion was rejected with.
    /// @throws std::logic_error for an empty completion or one with no executor.
    T await_resume() {
        _shared->stopCallback.reset();
        if (_shared->outcome.load() == Outcome::Cancelled) {
            throw ::core::async::OperationCancelled{};
        }
        if (_shared->error) {
            std::rethrow_exception(_shared->error);
        }
        return std::move(*_shared->value);
    }
 
private:
    /// Attaches the await's handler pair to @p state; runs on its executor.
    ///
    /// The await owns the call it waits on, as a scope-gated continuation
    /// does: the scope's stop -- a stop that withdraws the await, or the frame
    /// destroyed while suspended -- is linked to the call's stop source, so
    /// withdrawing the await also asks the call to stop. Linked after the
    /// scope has stopped, it stops the call at once.
    static void attach(CompletionState<T>& state, const std::shared_ptr<Shared>& shared) {
        auto const token = shared->scope.token();
        state.linkStop(token.stopToken());
        state.attachThen([token, shared](const T& settled) {
            if (token.active() && shared->decide(Outcome::Settled)) {
                shared->value.emplace(settled);
                shared->resumeSettled();
            }
        });
        state.attachOnError([token, shared](std::exception_ptr error) {
            if (token.active() && shared->decide(Outcome::Settled)) {
                shared->error = std::move(error);
                shared->resumeSettled();
            }
        });
    }
 
    /// Stops the call @p state reports on, for an await a stop withdrew before
    /// it attached, so no link carries that stop to the call. @p stopped,
    /// already stopped, is linked on the completion's executor, which stops the
    /// call there. Best effort: an executor that refuses the task leaves the
    /// call running; the coroutine resumes with the cancellation either way.
    static void stopUnattached(const std::shared_ptr<CompletionState<T>>& state,
                               const ::core::async::StopToken& stopped) noexcept {
        try {
            if (::morph::exec::runningOn(*state->cbExec)) {
                state->linkStop(stopped);
                return;
            }
            state->cbExec->post([state, stopped] { state->linkStop(stopped); });
        } catch (...) {  // NOLINT(bugprone-empty-catch): the await is answered already
        }
    }
 
    std::shared_ptr<CompletionState<T>> _state;
    std::shared_ptr<Shared> _shared;
};
 
}  // namespace morph::async::detail