lineariser.hpp

 
#pragma once
 
#include <cstddef>
#include <memory>
#include <utility>
 
#include <rapidsmpf/streaming/core/actor.hpp>
#include <rapidsmpf/streaming/core/channel.hpp>
#include <rapidsmpf/streaming/core/context.hpp>
#include <rapidsmpf/streaming/core/coro_utils.hpp>
#include <rapidsmpf/streaming/core/queue.hpp>
 
namespace rapidsmpf::streaming {
 
class Lineariser {
  public:
    Lineariser(
        std::shared_ptr<Context> ctx,
        std::shared_ptr<Channel> ch_out,
        std::size_t num_producers,
        std::size_t buffer_size = 1
    )
        : ctx_{std::move(ctx)}, ch_out_{std::move(ch_out)} {
        queues_.reserve(num_producers);
        for (std::size_t i = 0; i < num_producers; i++) {
            queues_.push_back(ctx->create_bounded_queue(buffer_size));
        }
    }
 
    std::vector<std::shared_ptr<BoundedQueue>>& get_queues() {
        return queues_;
    }
 
    Actor drain() {
        ShutdownAtExit c{ch_out_};
        co_await ctx_->executor()->schedule();
        while (!queues_.empty()) {
            for (auto& q : queues_) {
                auto [receipt, msg] = co_await q->receive();
                if (msg.empty()) {
                    q = nullptr;
                    continue;
                }
                if (!co_await ch_out_->send(std::move(msg))) {
                    // Output channel is shut down, tell the producers to shutdown.
                    break;
                }
                co_await receipt;
            }
            std::erase(queues_, nullptr);
        }
        // We either exited the loop because all the queues are gone, or the output
        // channel is shutdown (in which case we want no more inputs), so either way, just
        // shut down the remaining queues.
        std::vector<coro::task<void>> tasks;
        tasks.reserve(1 + queues_.size());
        for (auto& q : queues_) {
            tasks.push_back(q->shutdown());
        }
        tasks.push_back(ch_out_->drain(ctx_->executor()));
        coro_results(co_await coro::when_all(std::move(tasks)));
    }
 
    coro::task<void> shutdown() {
        std::vector<coro::task<void>> tasks;
        tasks.reserve(1 + queues_.size());
        for (auto& q : queues_) {
            tasks.push_back(q->shutdown());
        }
        tasks.push_back(ch_out_->shutdown());
        coro_results(co_await coro::when_all(std::move(tasks)));
    }
 
  private:
    std::shared_ptr<Context> ctx_;
    std::vector<std::shared_ptr<BoundedQueue>> queues_;
    std::shared_ptr<Channel> ch_out_;  
};
 
}  // namespace rapidsmpf::streaming