libdpf/include/dpf/compose_async.hpp

/// @file dpf/compose_async.hpp
/// @brief Drive compose `plan`s on `async_stream_array`.
/// @details The same `drive_via_schedule` path used on sync `stream_array`
///          runs on `async_round_sink`, so peer reads are event-driven.
///          `drive_options::n_lanes` sizes the stream pool and
///          `drive_options::framing` picks the wire layout.
#ifndef LIBDPF_INCLUDE_DPF_COMPOSE_ASYNC_HPP__
#define LIBDPF_INCLUDE_DPF_COMPOSE_ASYNC_HPP__

#include <cstddef>
#include <cstdint>
#include <exception>
#include <map>
#include <mutex>
#include <thread>
#include <utility>
#include <vector>

#include "dpf/net/asio_ns.hpp"

#include "dpf/compose.hpp"
#include "dpf/net/async_round_sink.hpp"
#include "dpf/net/async_stream_array.hpp"
#include "dpf/net/round_lane.hpp"

namespace dpf
{
namespace protocol
{

/// @brief Sink options that follow a drive's framing and drain budget.
inline net::sink_options sink_options_for(const drive_options & opt)
{
    net::sink_options so;
    so.framing = opt.framing;
    if (opt.wait_timeout.count() != 0)
        so.drain_timeout = opt.wait_timeout;
    return so;
}

/// @brief Drive a peer-only plan on an `async_stream_array` (overlapped I/O).
/// @param instances  RoundSink batch width (protocol instances per round).
inline void drive_plan_on_async_streams(const plan & p,
    net::async_stream_array & streams,
    std::vector<std::vector<std::uint8_t>> & values,
    const std::map<std::uint32_t, kernel_fn> & kernels, std::size_t party = 0,
    std::size_t instances = 1, const drive_options & opt = {},
    net::sink_options sopt = {})
{
    auto slots = p.slot_bytes_all();
    if (!slots.empty() && streams.size() == 0)
        throw std::invalid_argument(
            "drive_plan_on_async_streams: empty stream array");
    sopt.framing = opt.framing;
    net::async_round_sink sink(streams, std::move(slots), instances,
        std::move(sopt));
    drive_options local = opt;
    if (local.workers != nullptr && local.pump == nullptr)
        local.pump = &streams.context();
    drive_via_schedule(p, sink, values, kernels, party, local);
}

/// @brief Two parties, two threads, split-io async memory pair; drive both plans.
/// @details A party that fails closes its end, so the other fails fast.
inline void drive_both_on_async_streams(const plan & p0, const plan & p1,
    std::vector<std::vector<std::uint8_t>> & v0,
    std::vector<std::vector<std::uint8_t>> & v1,
    const std::map<std::uint32_t, kernel_fn> & kernels = {},
    std::size_t instances = 1, const drive_options & opt = {})
{
    const auto slots = p0.slot_bytes_all();
    if (slots != p1.slot_bytes_all())
        throw std::invalid_argument(
            "drive_both_on_async_streams: party slot shapes differ");
    const std::size_t nstreams =
        slots.empty() ? 1 : lanes_for_plan(slots.size(), opt);
    asio::io_context io0;
    asio::io_context io1;
    auto peer = net::make_async_dual_memory_stream_pair(io0, io1, nstreams);
    auto work0 = asio::make_work_guard(io0);
    auto work1 = asio::make_work_guard(io1);
    std::mutex err_mu;
    std::exception_ptr err;
    auto note = [&](std::exception_ptr e) {
        std::lock_guard<std::mutex> lock(err_mu);
        if (!err)
            err = std::move(e);
    };
    std::thread t0([&] {
        try
        {
            drive_plan_on_async_streams(p0, peer.first, v0, kernels, 0,
                instances, opt);
        }
        catch (...)
        {
            note(std::current_exception());
            peer.first.close();
        }
        work0.reset();
    });
    std::thread t1([&] {
        try
        {
            drive_plan_on_async_streams(p1, peer.second, v1, kernels, 1,
                instances, opt);
        }
        catch (...)
        {
            note(std::current_exception());
            peer.second.close();
        }
        work1.reset();
    });
    t0.join();
    t1.join();
    if (err)
        std::rethrow_exception(err);
}

/// @brief Same plan on both parties (symmetric compose graphs).
inline void drive_both_on_async_streams(const plan & p,
    std::vector<std::vector<std::uint8_t>> & v0,
    std::vector<std::vector<std::uint8_t>> & v1,
    const std::map<std::uint32_t, kernel_fn> & kernels = {},
    std::size_t instances = 1, const drive_options & opt = {})
{
    drive_both_on_async_streams(p, p, v0, v1, kernels, instances, opt);
}

}  // namespace protocol
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_COMPOSE_ASYNC_HPP__