queue.hpp

 
#pragma once
 
#include <cstddef>
#include <memory>
#include <optional>
#include <utility>
 
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/core/coro_utils.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
#include <rapidsmpf/streaming/core/node.hpp>
 
#include <coro/queue.hpp>
#include <coro/sync_wait.hpp>
 
namespace rapidsmpf::streaming {
 
class BoundedQueue {
    friend Context;
 
  private:
    class Shutdown {
      public:
        explicit Shutdown(BoundedQueue& q) : q_{q} {}
 
        ~Shutdown() noexcept {
            coro::sync_wait(q_.shutdown());
        }
 
      private:
        BoundedQueue& q_;
    };
 
    class Ticket {
      public:
        Ticket& operator=(Ticket const&) = delete;
        Ticket(Ticket const&) = delete;
        Ticket& operator=(Ticket&&) = default;
        Ticket(Ticket&&) = default;
        ~Ticket() = default;
 
        Ticket(coro::queue<Message>* q, Semaphore* semaphore)
            : q_{q}, semaphore_{semaphore} {}
 
        [[nodiscard]] coro::task<bool> send(Message msg) {
            RAPIDSMPF_EXPECTS(
                q_ && semaphore_, "Ticket has already been used", std::logic_error
            );
            auto result = co_await q_->push(std::move(msg));
            q_ = nullptr;
            semaphore_ = nullptr;
            co_return result == coro::queue_produce_result::produced;
        }
 
      private:
        coro::queue<Message>* q_;
        Semaphore* semaphore_;
    };
 
    BoundedQueue(std::size_t buffer_size)
        : semaphore_{static_cast<std::ptrdiff_t>(buffer_size)} {};
 
  public:
    [[nodiscard]] coro::task<std::optional<Ticket>> acquire() {
        auto result = co_await semaphore_.acquire();
        if (result == coro::semaphore_acquire_result::shutdown) {
            co_return std::nullopt;
        }
        co_return std::make_optional<Ticket>(&q_, &semaphore_);
    }
 
    [[nodiscard]] coro::task<std::pair<coro::task<void>, Message>> receive() {
        auto msg = co_await q_.pop();
        if (msg.has_value()) {
            co_return {semaphore_.release(), std::move(*msg)};
        } else {
            co_return {[]() -> coro::task<void> { co_return; }(), Message{}};
        }
    }
 
    [[nodiscard]] coro::task<void> drain(std::unique_ptr<coro::thread_pool>& executor) {
        co_await q_.shutdown_drain(executor);
        co_await semaphore_.shutdown();
    }
 
    [[nodiscard]] coro::task<void> shutdown() {
        coro_results(co_await coro::when_all(q_.shutdown(), semaphore_.shutdown()));
    }
 
    [[nodiscard]] BoundedQueue::Shutdown raii_shutdown() noexcept {
        return BoundedQueue::Shutdown{*this};
    }
 
  private:
    coro::queue<Message> q_{};  
    Semaphore semaphore_;  
};
 
}  // namespace rapidsmpf::streaming