libdpf/include/dpf/app_runtime.hpp

/// @file dpf/app_runtime.hpp
/// @brief File prep and two-party synchronous mux helpers.
/// @details `stream_runtime` configures only this synchronous TCP mux path.
///          For any other transport, party count, or a dealer, use
///          `app::run_config` with `app::run_parties` (`party_runner.hpp`).
///          The synchronous mux speaks the same frames as the async mux.
#ifndef LIBDPF_INCLUDE_DPF_APP_RUNTIME_HPP__
#define LIBDPF_INCLUDE_DPF_APP_RUNTIME_HPP__

#include <atomic>
#include <cstddef>
#include <cstdint>
#include <map>
#include <mutex>
#include <stdexcept>
#include <string>
#include <thread>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/compose.hpp"
#include "dpf/net/stream_array.hpp"
#include "dpf/net/sync_stream_array.hpp"
#include "dpf/party_run.hpp"
#include "dpf/prep_source.hpp"
#include "dpf/protocol_roles.hpp"

namespace dpf
{
namespace app
{

/// @brief File prep basename and synchronous mux settings (two parties).
struct stream_runtime
{
    std::string prep_basename;
    std::string mux_host = "127.0.0.1";
    unsigned short mux_port = 0;
    std::size_t mux_nstreams = 1;
};

/// @brief Read one party's prep blob from `basename-p0` / `basename-p1`.
HEDLEY_WARN_UNUSED_RESULT
inline std::vector<std::uint8_t> open_file_prep(const std::string & basename,
    unsigned party)
{
    return prep::read_file(basename + "-p" + std::to_string(party));
}

/// @brief Same as `open_file_prep`, parsed as a prep cursor.
HEDLEY_WARN_UNUSED_RESULT
inline prep::cursor open_file_prep_cursor(const std::string & basename,
    unsigned party)
{
    return prep::cursor(open_file_prep(basename, party));
}

/// @brief Run `fn(party, mux_stream_array &)` over one TCP connection with mux.
template <typename Fn>
void with_tcp_mux_peer(unsigned party, std::string host,
    std::atomic<unsigned short> & port, std::size_t nstreams, Fn && fn)
{
    dpf::run::tcp_pair_mux(party, std::move(host), port, nstreams,
        std::forward<Fn>(fn));
}

/// @brief Convenience: mux settings from `stream_runtime`.
template <typename Fn>
void with_tcp_mux_peer(unsigned party, const stream_runtime & rt,
    std::size_t nstreams, Fn && fn)
{
    std::atomic<unsigned short> port{rt.mux_port};
    with_tcp_mux_peer(party, rt.mux_host, port, nstreams, std::forward<Fn>(fn));
}

/// @brief One party: TCP mux + `drive_plan_on_streams`.
inline void drive_plan_mux(const protocol::plan & plan,
    net::stream_array & mux, std::vector<std::vector<std::uint8_t>> & values,
    const std::map<std::uint32_t, protocol::kernel_fn> & kernels, std::size_t party,
    std::size_t lanes = 1, const protocol::drive_options & opt = {})
{
    protocol::drive_plan_on_streams(plan, mux, values, kernels, party, lanes, opt);
}

/// @brief Two localhost threads on one mux TCP link, both plans driven on streams.
inline void drive_plan_mux_both(const protocol::plan & p0,
    const protocol::plan & p1, std::vector<std::vector<std::uint8_t>> & v0,
    std::vector<std::vector<std::uint8_t>> & v1,
    const std::map<std::uint32_t, protocol::kernel_fn> & kernels = {},
    std::size_t lanes = 1, const protocol::drive_options & opt = {},
    const stream_runtime & rt = {})
{
    const auto slots0 = p0.slot_bytes_all();
    const auto slots1 = p1.slot_bytes_all();
    if (slots0 != slots1)
        throw std::invalid_argument("drive_plan_mux_both: slot shapes differ");
    const std::size_t nstreams =
        slots0.empty() ? 1 : protocol::lanes_for_plan(slots0.size(), opt);
    std::atomic<unsigned short> port{rt.mux_port};
    std::exception_ptr err;
    std::mutex err_mu;
    auto note = [&](std::exception_ptr e) {
        std::lock_guard<std::mutex> lock(err_mu);
        if (!err)
            err = std::move(e);
    };
    std::thread t0([&] {
        try
        {
            with_tcp_mux_peer(0, rt.mux_host, port, nstreams,
                [&](unsigned, net::mux_stream_array & mux) {
                    drive_plan_mux(p0, mux, v0, kernels, 0, lanes, opt);
                });
        }
        catch (...)
        {
            note(std::current_exception());
        }
    });
    std::thread t1([&] {
        try
        {
            with_tcp_mux_peer(1, rt.mux_host, port, nstreams,
                [&](unsigned, net::mux_stream_array & mux) {
                    drive_plan_mux(p1, mux, v1, kernels, 1, lanes, opt);
                });
        }
        catch (...)
        {
            note(std::current_exception());
        }
    });
    t0.join();
    t1.join();
    if (err)
        std::rethrow_exception(err);
}

}  // namespace app
}  // namespace dpf

#endif