spsc_live_handoff_queue.hpp

1// SpscLiveHandoffQueue combines a resource pool with an SPSC producer-consumer queue.2// Producers borrow and populate resources; completed consumption returns resources to the pool.3// Producer: borrow from the pool (try_acquire) -> write/publish Message -> hand off (send / send_and_acquire).4// Consumer: claim each completed record with try_claim -> finish on success -> continue to the next record.5//                                                    -> otherwise cancel_claim and retain the record at the head for retry.6// Differences from a conventional SPSC queue:7// 1. Production needs a resource pool, which may be replaced by a similar pool or an upstream completion queue.8// 2. After earlier completed records, the consumer may read the resource still being produced. Each {finished, next} record keeps current-resource reads aligned with the completion queue.9// 3. claim/cancel replaces pop.10// Compared with pop, claim/cancel: a. simplifies consumers by keeping pending resources in the queue;11//                               b. avoids losing a removed resource when a consumer exits unexpectedly;1213#pragma once1415#include "TwoBufQueue/src/base_components/warm_pool/warm_pool.hpp"16#include "TwoBufQueue/src/utils/spsc_queue.hpp"17#include <cassert>1819namespace TwoBufQueue {2021enum class LiveAcquireStatus { Success, ProducerBusy, QueueFull, PoolLimit, ResourceError, Closed };22struct LiveHandoffStats {23    std::uint64_t acquired = 0, queue_full = 0, pool_limit = 0, resource_error = 0;24};2526// Ordered SPSC handoff. Resource publication cursors and read logic do not belong to this queue.27// Except is_closed/pending_count/stats, each interface is used only by its owning side; stop both sides before destruction.28template<class TPool>29class SpscLiveHandoffQueue {30public:31    using Pool = TPool;32    using Resource = typename TPool::ResourceType;33    using Loan = typename TPool::Loan;34    static_assert(TPool::concurrent, "SpscLiveHandoffQueue requires a concurrent pool");35    template<bool TConsumer> class Ticket {36        friend class SpscLiveHandoffQueue;37    public:38        Ticket() noexcept = default;39        // A consumer ticket represents a dequeued record; begin records are valid too, with empty get().40        explicit operator bool() const noexcept {41            if constexpr (TConsumer) return owner_ != nullptr;42            else return bool(loan_);43        }44        Resource* get() const noexcept { return loan_.get(); }45        Resource* operator->() const noexcept { return get(); }46        Resource& operator*() const noexcept { return *get(); }47    private:48        const SpscLiveHandoffQueue* owner_ = nullptr;49        std::uint64_t id_ = 0;50        Loan loan_;51    };52    using ProducerLoan = Ticket<false>;53    using Claim = Ticket<true>;54    class Record : public Ticket<true> {55        friend class SpscLiveHandoffQueue;56    public:57        // Borrow the next producer resource only; read published data after earlier records finish, and never return it yourself.58        [[nodiscard]] Loan next() const noexcept { return next_; }59    private:60        Loan next_;61    };62    struct AcquireResult {63        LiveAcquireStatus status;64        ProducerLoan loan;65        explicit operator bool() const noexcept { return status == LiveAcquireStatus::Success; }66    };67    struct Callbacks {68        void (*on_send)(Resource&) noexcept = nullptr;69        // Preconditions for finish: retain an unfinished Record and resume it on the next attempt.70        bool (*can_finish)(const Resource&) noexcept = nullptr;71        // An optional callback performs the final return; otherwise the queue returns the resource directly to the external pool.72        // The callback may immediately destroy the resource; the queue does not access it afterward.73        void (*on_finish)(Resource&) noexcept = nullptr;74    };75    using NotifyCallback = void (*)(void*) noexcept;7677    SpscLiveHandoffQueue(TPool& pool, std::size_t capacity, Callbacks callbacks = {}, bool enable_stats = false)78        : pool_(pool), entries_(checked_capacity(capacity)), callbacks_(callbacks), enable_stats_(enable_stats) {}79    SpscLiveHandoffQueue(const SpscLiveHandoffQueue&) = delete;80    SpscLiveHandoffQueue& operator=(const SpscLiveHandoffQueue&) = delete;81    ~SpscLiveHandoffQueue() noexcept {82        if (producer_ || claim_ || popped_id_ != finished_id_ || current_ || !entries_.empty()) {83            std::fputs("destroying SpscLiveHandoffQueue before draining handoffs\n", stderr);84            std::terminate();85        }86    }8788    [[nodiscard]] TPool& pool() const noexcept { return pool_; }89    [[nodiscard]] const Callbacks& callbacks() const noexcept { return callbacks_; }9091#pragma region Producer side (writer thread)92public:93    // Producer: duplicate acquisition is an API error; never implicitly obtain another block or return the existing one.94    [[nodiscard]] AcquireResult try_acquire() noexcept {95        if (producer_) return {LiveAcquireStatus::ProducerBusy, {}};96        auto result = acquire();97        if (result) enqueue({}, producer_.loan_); // Begin record: exposes the new current block to the consumer.98        return result;99    }100    [[nodiscard]] ProducerLoan current_producer() const noexcept { return producer_; }101    bool send() noexcept {102        if (!producer_) return false;103        const auto old = producer_.loan_;104        if (callbacks_.on_send) callbacks_.on_send(*old);105        producer_ = {};106        enqueue(old, {});107        notify_read_available();108        return true;109    }110    // Like the original get_next_mem, a successful switch publishes exactly one {old block, new block} record.111    // Failure still hands off the old block; the caller may subsequently wait/retry try_acquire.112    [[nodiscard]] AcquireResult send_and_acquire() noexcept {113        const auto finished = producer_;114        auto finished_loan = finished ? finished.loan_ : Loan{};115        if (finished) {116            if (callbacks_.on_send) callbacks_.on_send(*finished);117            producer_ = {};118        }119        auto result = acquire();120        auto next_loan = producer_ ? producer_.loan_ : Loan{};121        if (finished || producer_) {122            enqueue(finished_loan, next_loan); // {old block, new block}; acquisition failure yields {old block, empty}.123            notify_read_available();124        }125        return result;126    }127    // Reject new acquisitions without implicitly sending current; an existing producer ticket may still be sent.128    void close() noexcept { closed_.store(true, std::memory_order_release); }129    // External wakeup only; do not reenter the producer interface.130    [[nodiscard]] bool has_read_available_callback() const noexcept { return notify_ != nullptr; }131132private:133    // Borrow from the pool and set producer_ without enqueueing; try_acquire / send_and_acquire publish the handoff record.134    AcquireResult acquire() noexcept {135        auto status = LiveAcquireStatus::Closed;136        Loan next;137        if (!is_closed()) {138            // Reserve this record's slot plus one extra so the current producer block can always be handed off.139            if (entries_.capacity() - entries_.size() < 2) {140                status = LiveAcquireStatus::QueueFull;141                if (enable_stats_) queue_full_.fetch_add(1, std::memory_order_relaxed);142            } else {143                PoolGetStatus pool_status;144                next = pool_.get(pool_status);145                if (next) {146                    status = LiveAcquireStatus::Success;147                    producer_ = make_ticket<false>(next, next_id(producer_id_));148                    if (enable_stats_) acquired_.fetch_add(1, std::memory_order_relaxed);149                } else if (pool_status == PoolGetStatus::Limit) {150                    status = LiveAcquireStatus::PoolLimit;151                    if (enable_stats_) pool_limit_.fetch_add(1, std::memory_order_relaxed);152                } else {153                    status = LiveAcquireStatus::ResourceError;154                    if (enable_stats_) resource_error_.fetch_add(1, std::memory_order_relaxed);155                }156            }157        }158        return {status, producer_};159    }160    void enqueue(Loan finished, Loan next) noexcept {161        if (!entries_.try_push(Entry{finished, next})) std::terminate(); // Reserved-capacity invariant.162    }163164#pragma endregion165166#pragma region Consumer side (reader thread)167public:168    // Claim the head without removing it; its queue slot stays occupied. Only one active Claim is allowed.169    [[nodiscard]] Claim try_claim() noexcept {170        assert(!claim_ && popped_id_ == finished_id_);171        if (claim_ || popped_id_ != finished_id_) return {};172        const auto* entry = entries_.front();173        if (!entry) return {};174        if (claim_id_ == std::numeric_limits<std::uint64_t>::max()) std::terminate();175        claim_ = make_ticket<true>(entry->finished, ++claim_id_);176        return claim_;177    }178179    // Cancel without dequeueing or returning the resource; the next try_claim yields the same head record.180    void cancel_claim(Claim& claim) noexcept {181        const bool valid = matches(claim, claim_);182        assert(valid);183        if (!valid) return;184        claim_ = {};185        claim = {};186    }187188    // Finish a Claim: dequeue, advance current and return any completed resource to the pool.189    void finish(Claim& claim) noexcept {190        const bool valid = matches(claim, claim_) && popped_id_ == finished_id_;191        assert(valid);192        if (!valid) return;193        const auto* entry = entries_.front();194        assert(entry);195        if (!entry) return;196        const auto loan = claim_.loan_;197        const auto next = entry->next;198        if (loan && callbacks_.can_finish && !callbacks_.can_finish(*loan)) {199            assert(false && "finish requires a fully processed resource");200            std::terminate();201        }202        claim_ = {};203        claim = {};204        entries_.pop();205        current_ = next;206        popped_id_ = next_id(popped_id_);207        finished_id_ = popped_id_;208        recycle_finished(loan);209    }210211    // Remove one record and immediately release its queue slot, without returning its resource.212    // {empty, next} is a valid Record with empty get(); it still requires finish acknowledgement.213    // Consecutive pop calls are allowed. The caller retains pending records and handles/finishes them in dequeue order.214    [[nodiscard]] Record try_pop() noexcept {215        assert(!claim_);216        if (claim_) return {};217        const auto* entry = entries_.front();218        if (!entry) return {};219        Record record;220        static_cast<Ticket<true>&>(record) = make_ticket<true>(entry->finished, next_id(popped_id_));221        record.next_ = entry->next;222        current_ = entry->next;223        entries_.pop();224        return record;225    }226227    // Consumer-side current is the last popped next; query only, without dequeueing or reading producer_.228    // When records remain queued, current may lag behind the tail's next; snapshot_for_read fixes this round's boundary first.229    // Finish all previously popped pending records before reading the resource; never return current directly.230    [[nodiscard]] Loan current() const noexcept { return current_; }231232    // Fix this round's queued record count: use the last record's next, or current() if the queue is empty.233    // Excludes already popped pending records; process those and this round's records before reading the returned block.234    [[nodiscard]] Loan snapshot_for_read(std::size_t& count) const noexcept {235        count = entries_.size();236        if (!count) return current();237#ifndef NDEBUG238        // Compare tickets only; do not dereference resources or access the producer's producer_.239        const auto& head = entries_.front()->finished;240        assert(head.get() == current_.get() && head.generation() == current_.generation());241#endif242        return entries_.front(count - 1)->next;243    }244245    // Consumer-side query for transferring pending work to another stopped reader; it does not extend resource lifetime.246    [[nodiscard]] bool is_pending(const Record& record) const noexcept {247        return record.owner_ == this && record.id_ > finished_id_ && record.id_ <= popped_id_;248    }249    // Acknowledge a dequeued record and return its resource without touching queue slots. Call finish in dequeue order.250    void finish(Record& record) noexcept {251        const bool valid = is_pending(record) && record.id_ - finished_id_ == 1;252        assert(valid);253        if (!valid) return;254        const auto loan = record.loan_;255        if (loan && callbacks_.can_finish && !callbacks_.can_finish(*loan)) {256            assert(false && "finish requires a fully processed resource");257            std::terminate();258        }259        finished_id_ = record.id_;260        record = {};261        if (!loan) return; // Begin records have no old resource; invoke neither the resource callback nor recycle.262        recycle_finished(loan);263    }264    [[nodiscard]] bool empty() const noexcept { return entries_.empty(); }265    [[nodiscard]] bool is_drained() const noexcept {266        return is_closed() && popped_id_ == finished_id_ && !current_ && entries_.empty();267    }268269#pragma endregion270271#pragma region Queries / cross-side operations272public:273    [[nodiscard]] bool is_closed() const noexcept { return closed_.load(std::memory_order_acquire); }274    [[nodiscard]] std::size_t pending_count() const noexcept { return entries_.size(); }275    [[nodiscard]] std::size_t capacity() const noexcept { return entries_.capacity(); }276    [[nodiscard]] bool owns(const Loan& loan) const noexcept { return pool_.owns(loan); }277    [[nodiscard]] LiveHandoffStats stats() const noexcept {278        return {acquired_.load(std::memory_order_relaxed), queue_full_.load(std::memory_order_relaxed),279                pool_limit_.load(std::memory_order_relaxed), resource_error_.load(std::memory_order_relaxed)};280    }281    // Install/remove while the producer is quiescent; external wakeup only, never reenter production.282    void set_read_available_callback(NotifyCallback callback, void* context) noexcept {283        notify_ = callback;284        notify_context_ = context;285    }286    // Notify the consumer that new published data may be read; the producer calls this after send, send_and_acquire or flush.287    void notify_read_available() noexcept {288        if (notify_) notify_(notify_context_);289    }290291#pragma endregion292293private:294    void recycle_finished(const Loan& loan) noexcept {295        if (!loan) return;296        if (callbacks_.on_finish) callbacks_.on_finish(*loan);297        else if (!pool_.recycle(loan)) std::terminate();298    }299    struct Entry { Loan finished, next; };300    static std::size_t checked_capacity(std::size_t capacity) {301        if (capacity < 2) throw std::invalid_argument("SpscLiveHandoffQueue capacity must be at least 2");302        return capacity;303    }304    static std::uint64_t next_id(std::uint64_t& id) noexcept {305        if (id == std::numeric_limits<std::uint64_t>::max()) std::terminate();306        return ++id;307    }308    template<bool TC> Ticket<TC> make_ticket(Loan loan, std::uint64_t id) const noexcept {309        Ticket<TC> ticket;310        ticket.owner_ = this;311        ticket.id_ = id;312        ticket.loan_ = loan;313        return ticket;314    }315    template<bool TC> bool matches(const Ticket<TC>& input, const Ticket<TC>& active) const noexcept {316        return active && input.owner_ == this && input.id_ == active.id_;317    }318    // Data members: preserve cache-line separation.319    TPool& pool_;320    SpscQueue<Entry> entries_;321    const Callbacks callbacks_;322    const bool enable_stats_;323    NotifyCallback notify_ = nullptr;324    void* notify_context_ = nullptr;325    std::atomic<bool> closed_{false};326    alignas(64) ProducerLoan producer_;327    std::uint64_t producer_id_ = 0;328    alignas(64) Claim claim_;329    std::uint64_t claim_id_ = 0;330    alignas(64) Loan current_;331    std::uint64_t popped_id_ = 0, finished_id_ = 0;332    alignas(64) std::atomic<std::uint64_t> acquired_{0}, queue_full_{0}, pool_limit_{0}, resource_error_{0};333};334335// Convenience resource-type specialization: owns a default WarmPool internally.336// Passing WarmPool<TResource> still uses the primary template and supports an externally shared pool.337template<PoolResource TResource>338class SpscLiveHandoffQueue<TResource>339    : private WarmPool<TResource>, public SpscLiveHandoffQueue<WarmPool<TResource>> {340    using OwnedPool = WarmPool<TResource>;341    using QueueBase = SpscLiveHandoffQueue<OwnedPool>;342public:343    using Callbacks = typename QueueBase::Callbacks;344345    SpscLiveHandoffQueue(std::size_t min_count, std::size_t max_count, std::size_t capacity,346                        Callbacks callbacks = {}, bool enable_stats = false)347        : OwnedPool(min_count, max_count),348          QueueBase(static_cast<OwnedPool&>(*this), capacity, callbacks, enable_stats) {}349350    [[nodiscard]] bool all_released() const noexcept { return OwnedPool::all_released(); }351};352353} // namespace TwoBufQueue