cancel.h

Go to the documentation of this file.

 
#pragma once
 
#include <neograph/api.h>
 
#include <asio/any_io_executor.hpp>
#include <asio/cancellation_signal.hpp>
#include <asio/post.hpp>
#include <asio/strand.hpp>
 
#include <algorithm>
#include <condition_variable>
#include <stop_token>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <utility>
#include <vector>
 
namespace neograph::graph {
 
class NEOGRAPH_API CancelledException : public std::runtime_error {
public:
    CancelledException()
        : std::runtime_error("neograph: run cancelled") {}
    explicit CancelledException(const std::string& detail)
        : std::runtime_error("neograph: run cancelled — " + detail) {}
};
 
class NEOGRAPH_API CancelToken {
public:
    CancelToken() = default;
 
    CancelToken(const CancelToken&) = delete;
    CancelToken& operator=(const CancelToken&) = delete;
 
    void cancel() noexcept {
        if (stop_source_.stop_requested()) return;
        // Retain operation children before a native callback can release its
        // awaiting frame. No callback runs while either token mutex is held.
        auto keep_alive = self_keep_alive_for_post();
        if (!stop_source_.request_stop()) return;
 
        // Serialize executor submission with unbind_executor(), so completed
        // operations cannot post into a context that has begun teardown.
        {
            std::lock_guard<std::mutex> lk(mu_);
            if (ex_) {
                asio::post(ex_,
                           [this, binding = binding_generation_,
                            keep_alive = std::move(keep_alive)]() {
                               emit_if_bound(binding);
                           });
            }
        }
 
        // v1.0 (9c): the legacy ``add_cancel_hook`` / hooks_ iteration
        // is gone — ``fork()`` is the canonical primitive for nested
        // cancel scopes, and every internal ``run_sync`` cascades via
        // a forked child.
 
        // Cascade to live children produced by ``fork()``.
        // Snapshot live shared_ptrs under the lock, then call
        // ``cancel()`` on each outside the lock — the recursive call
        // would otherwise re-enter ``children_mu_`` if a grandchild
        // exists, and lock_guard isn't reentrant.
        std::vector<std::shared_ptr<CancelToken>> live_children;
        {
            std::lock_guard<std::mutex> lk(children_mu_);
            live_children.reserve(children_.size());
            for (auto& w : children_) {
                if (auto sp = w.lock(); sp && sp.get() != this) {
                    live_children.push_back(std::move(sp));
                }
            }
            // Don't bother pruning expired entries here — fork() does
            // it on the next call, and after this cancel() the parent
            // is "done" anyway.
        }
        for (auto& child : live_children) {
            child->cancel();
        }
    }
 
    [[nodiscard]] bool is_cancelled() const noexcept {
        return stop_source_.stop_requested();
    }
 
    [[nodiscard]] std::stop_token stop_token() const noexcept {
        return stop_source_.get_token();
    }
 
    [[nodiscard]] asio::any_io_executor bind_executor(asio::any_io_executor ex) {
        asio::any_io_executor bound;
        bool                   fire_immediately = false;
        std::uint64_t          binding = 0;
        {
            std::lock_guard<std::mutex> lk(mu_);
            ex_ = ex ? asio::any_io_executor(asio::make_strand(std::move(ex)))
                     : asio::any_io_executor{};
            binding = ++binding_generation_;
            bound = ex_;
            fire_immediately = stop_source_.stop_requested();
        }
        emit_finished_.notify_all();
        if (fire_immediately && bound) {
            auto keep_alive = self_keep_alive_for_post();
            asio::post(bound,
                       [this, binding, keep_alive = std::move(keep_alive)]() {
                           emit_if_bound(binding);
                       });
        }
        return bound;
    }
 
    void unbind_executor() noexcept {
        asio::any_io_executor released;
        {
            std::unique_lock<std::mutex> lk(mu_);
            released = std::move(ex_);
            ++binding_generation_;
            emit_finished_.wait(lk, [this] { return !emitting_; });
        }
        // No emit can now touch the signal. Clear outside mu_: destruction of
        // a user's completed handler may itself request cancellation.
        sig_.slot().clear();
        // Keep the former executor alive through this method's end, while its
        // execution context is still known to be alive to the caller.
        static_cast<void>(released);
    }
 
    asio::cancellation_slot slot() noexcept {
        return sig_.slot();
    }
 
    void throw_if_cancelled(const std::string& detail = {}) const {
        if (is_cancelled()) {
            throw CancelledException(detail);
        }
    }
 
    [[nodiscard]] std::shared_ptr<CancelToken> fork() {
        // Construct child outside any lock — its ctor takes no lock.
        // shared_ptr ctor allocates the control block; cheap relative
        // to a real cancel scope (one io_context spin-up downstream).
        auto child = std::shared_ptr<CancelToken>(new CancelToken());
 
        // The weak self-entry retains operation children through posted emits.
        // cancel() excludes this entry from parent-to-child propagation.
        {
            std::lock_guard<std::mutex> lk(child->children_mu_);
            child->children_.push_back(child);
        }
 
        {
            std::lock_guard<std::mutex> lk(children_mu_);
            // Opportunistic prune: every fork() is a natural place to
            // drop expired weak_ptrs from previous children that have
            // already been released by their owners. Bounds the
            // children_ vector at the live-fan-out width even on
            // long-lived parents (e.g. an engine with 1000 LLM calls
            // per run, each forking once).
            children_.erase(
                std::remove_if(
                    children_.begin(), children_.end(),
                    [](const std::weak_ptr<CancelToken>& w) {
                        return w.expired();
                    }),
                children_.end());
            children_.push_back(child);
        }
 
        // Eager propagation: parent already cancelled → child sees
        // the polling flag immediately. ``bind_executor`` on the
        // child will then fire its signal at the next co_await.
        if (stop_source_.stop_requested()) {
            child->cancel();
        }
        return child;
    }
 
private:
    void emit_if_bound(std::uint64_t binding) {
        {
            std::unique_lock<std::mutex> lk(mu_);
            emit_finished_.wait(lk, [this, binding] {
                return !ex_ || binding != binding_generation_ || !emitting_;
            });
            if (!ex_ || binding != binding_generation_) return;
            emitting_ = true;
        }
        struct EmitCompletion {
            CancelToken& token;
            ~EmitCompletion() {
                {
                    std::lock_guard<std::mutex> lk(token.mu_);
                    token.emitting_ = false;
                }
                token.emit_finished_.notify_all();
            }
        } complete{*this};
        // Signal handlers may reenter bind_executor() or cancel(). Admission
        // is fenced above, but no token mutex may be held through a callback.
        sig_.emit(asio::cancellation_type::all);
    }
 
    std::shared_ptr<CancelToken> self_keep_alive_for_post() {
        std::lock_guard<std::mutex> lk(children_mu_);
        for (auto& candidate : children_) {
            if (auto self = candidate.lock(); self.get() == this) {
                return self;
            }
        }
        return {};
    }
 
    std::stop_source        stop_source_;
    mutable std::mutex       mu_;        // guards executor and emit admission
    std::condition_variable emit_finished_;
    std::uint64_t           binding_generation_ = 0;
    bool                    emitting_ = false;
    asio::any_io_executor    ex_;        // bound by engine before HTTP I/O
    asio::cancellation_signal sig_;      // for asio operation cancel
 
    // Hierarchical cascade list. Forked children retain a weak self-entry for
    // cancellation callbacks and posted emits; cancel() skips it when cascading.
    mutable std::mutex children_mu_;
    std::vector<std::weak_ptr<CancelToken>> children_;
};
 
 
// Keeps a token's executor binding scoped to the operation that drained it.
// Construct after the token is bound; destruction must occur before the
// corresponding io_context or thread_pool is torn down.
class CancelExecutorLease final {
public:
    explicit CancelExecutorLease(std::shared_ptr<CancelToken> token) noexcept
        : token_(std::move(token)) {}
 
    CancelExecutorLease(const CancelExecutorLease&) = delete;
    CancelExecutorLease& operator=(const CancelExecutorLease&) = delete;
 
    ~CancelExecutorLease() noexcept {
        if (token_) token_->unbind_executor();
    }
 
private:
    std::shared_ptr<CancelToken> token_;
};
 
 
// v1.0 (9d): the `current_cancel_token()` thread_local +
// `CurrentCancelTokenScope` RAII smuggling channel is gone.
// `RunContext::cancel_token` (on `NodeInput::ctx`) is the only cancel
// channel. Engine threads it through every dispatch.
 
}  // namespace neograph::graph