channel.hpp

 
#pragma once
 
#include <cstddef>
#include <limits>
#include <memory>
#include <stdexcept>
#include <utility>
 
#include <coro/coro.hpp>
#include <coro/queue.hpp>
#include <coro/semaphore.hpp>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/streaming/core/actor.hpp>
#include <rapidsmpf/streaming/core/coro_executor.hpp>
#include <rapidsmpf/streaming/core/message.hpp>
#include <rapidsmpf/streaming/core/spillable_messages.hpp>
 
namespace rapidsmpf::streaming {
 
class Context;
 
using Semaphore = coro::semaphore<std::numeric_limits<std::ptrdiff_t>::max()>;
 
class Channel {
    friend Context;
 
  public:
    [[nodiscard]] coro::task<bool> send(Message msg);
 
    [[nodiscard]] coro::task<Message> receive();
 
    [[nodiscard]] coro::task<bool> send_metadata(Message msg);
 
    [[nodiscard]] coro::task<Message> receive_metadata();
 
    [[nodiscard]] Actor drain_metadata(std::shared_ptr<CoroThreadPoolExecutor> executor);
 
    [[nodiscard]] Actor drain(std::shared_ptr<CoroThreadPoolExecutor> executor);
 
    [[nodiscard]] Actor shutdown();
 
    [[nodiscard]] Actor shutdown_metadata();
 
    [[nodiscard]] bool empty() const noexcept;
 
    [[nodiscard]] bool is_shutdown() const noexcept;
 
  private:
    Channel(std::shared_ptr<SpillableMessages> spillable_messages)
        : sm_{std::move(spillable_messages)} {}
 
    coro::ring_buffer<SpillableMessages::MessageId, 1> rb_;
    std::shared_ptr<SpillableMessages> sm_;
    coro::queue<Message> metadata_;
};
 
class ThrottlingAdaptor {
  private:
    class Ticket {
      public:
        Ticket& operator=(Ticket const&) = delete;
        Ticket(Ticket const&) = delete;
        Ticket& operator=(Ticket&&) = default;
        Ticket(Ticket&&) = default;
        ~Ticket() = default;
 
        [[nodiscard]] coro::task<std::pair<bool, coro::task<void>>> send(Message msg) {
            RAPIDSMPF_EXPECTS(ch_, "Ticket has already been used", std::logic_error);
            auto sent = co_await ch_->send(std::move(msg));
            ch_ = nullptr;
            if (sent) {
                co_return {sent, semaphore_->release()};
            } else {
                // If the channel is closed we want to wake any waiters so shutdown the
                // semaphore.
                co_await semaphore_->shutdown();
                co_return {sent, []() -> coro::task<void> { co_return; }()};
            }
        }
 
        Ticket(Channel* channel, Semaphore* semaphore)
            : ch_{channel}, semaphore_{semaphore} {};
 
      private:
        Channel* ch_;
        Semaphore* semaphore_;
    };
 
  public:
    explicit ThrottlingAdaptor(
        std::shared_ptr<Channel> channel, std::ptrdiff_t max_tickets
    )
        : ch_{std::move(channel)}, semaphore_(max_tickets) {
        RAPIDSMPF_EXPECTS(
            max_tickets > 0, "ThrottlingAdaptor must have at least one ticket"
        );
    }
 
    [[nodiscard]] coro::task<Ticket> acquire() {
        auto result = co_await semaphore_.acquire();
        RAPIDSMPF_EXPECTS(
            result == coro::semaphore_acquire_result::acquired,
            "Semaphore was shutdown",
            std::runtime_error
        );
        co_return {ch_.get(), &semaphore_};
    }
 
  private:
    std::shared_ptr<Channel> ch_;
    Semaphore semaphore_;
};
 
class ShutdownAtExit {
  public:
    explicit ShutdownAtExit(std::vector<std::shared_ptr<Channel>> channels)
        : channels_{std::move(channels)} {
        for (auto& ch : channels_) {
            RAPIDSMPF_EXPECTS(ch, "channel cannot be null", std::invalid_argument);
        }
    }
 
    template <class... T>
    explicit ShutdownAtExit(T&&... channels)
        requires(std::convertible_to<T, std::shared_ptr<Channel>> && ...)
        : ShutdownAtExit(
              std::vector<std::shared_ptr<Channel>>{std::forward<T>(channels)...}
          ) {}
 
    // Non-copyable, non-movable.
    ShutdownAtExit(ShutdownAtExit const&) = delete;
    ShutdownAtExit& operator=(ShutdownAtExit const&) = delete;
    ShutdownAtExit(ShutdownAtExit&&) = delete;
    ShutdownAtExit& operator=(ShutdownAtExit&&) = delete;
 
    ~ShutdownAtExit() noexcept {
        for (auto& ch : channels_) {
            coro::sync_wait(ch->shutdown());
        }
    }
 
  private:
    std::vector<std::shared_ptr<Channel>> channels_;
};
 
}  // namespace rapidsmpf::streaming