libdpf/party/spawn.hpp

/// @file party/spawn.hpp
/// @brief Spawn p0/p1/p2 for one flow (used by gtest and party_bench).
#ifndef LIBDPF_PARTY_SPAWN_HPP__
#define LIBDPF_PARTY_SPAWN_HPP__

#include <chrono>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <stdexcept>
#include <string>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>

#ifndef LIBDPF_PARTY_BIN_DIR
#define LIBDPF_PARTY_BIN_DIR "."
#endif

namespace dpf
{
namespace party
{

inline std::string party_bin(const char * name)
{
    const char * env = std::getenv("LIBDPF_PARTY_BIN_DIR");
    std::string dir = env ? env : LIBDPF_PARTY_BIN_DIR;
    return dir + "/" + name;
}

struct spawn_result
{
    int rc[3] = {0, 0, 0};
    double wall_ms = 0;
    /// @brief Per-role metric lines (from --metrics), ordered p0/p1/p2.
    std::vector<std::string> metrics_lines;
};

struct spawn_opts
{
    std::uint64_t repeat = 1;
    std::uint64_t warmup = 0;
    bool metrics = false;
    /// @brief Pass `case_name` as `--tag` so one mesh runs every match.
    bool by_tag = false;
};

inline spawn_result spawn_trio_flow(const std::string & case_name,
    const spawn_opts & opts = {})
{
    namespace fs = std::filesystem;
    auto dir = fs::temp_directory_path()
        / ("libdpf_party_" + std::to_string(::getpid()) + "_" + case_name);
    fs::create_directories(dir);

    std::string rep = std::to_string(opts.repeat);
    std::string warm = std::to_string(opts.warmup);

    const auto t0 = std::chrono::steady_clock::now();
    pid_t kids[3];
    const char * bins[3] = {"p0", "p1", "p2"};
    for (int i = 0; i < 3; ++i)
    {
        kids[i] = fork();
        if (kids[i] < 0)
        {
            fs::remove_all(dir);
            throw std::runtime_error("fork failed");
        }
        if (kids[i] == 0)
        {
            auto bin = party_bin(bins[i]);
            const char * sel = opts.by_tag ? "--tag" : "--case";
            if (opts.metrics)
            {
                execl(bin.c_str(), bin.c_str(),
                    "--dir", dir.string().c_str(),
                    sel, case_name.c_str(),
                    "--repeat", rep.c_str(),
                    "--warmup", warm.c_str(),
                    "--metrics",
                    static_cast<char *>(nullptr));
            }
            else
            {
                execl(bin.c_str(), bin.c_str(),
                    "--dir", dir.string().c_str(),
                    sel, case_name.c_str(),
                    "--repeat", rep.c_str(),
                    "--warmup", warm.c_str(),
                    static_cast<char *>(nullptr));
            }
            _exit(126);
        }
    }

    spawn_result out;
    for (int i = 0; i < 3; ++i)
    {
        int status = 0;
        if (waitpid(kids[i], &status, 0) < 0)
            out.rc[i] = 125;
        else
            out.rc[i] = WIFEXITED(status) ? WEXITSTATUS(status) : 124;
    }
    const auto t1 = std::chrono::steady_clock::now();
    out.wall_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();

    if (opts.metrics)
    {
        const char * roles[3] = {"p0", "p1", "p2"};
        for (const char * role : roles)
        {
            auto path = dir / (std::string("metrics.") + role);
            std::ifstream in(path);
            std::string line;
            while (std::getline(in, line))
            {
                if (!line.empty())
                    out.metrics_lines.push_back(line);
            }
        }
    }

    fs::remove_all(dir);
    return out;
}

/// @brief Spawn only p0 and p1 (no dealer). For IKNP / two-party flows.
inline spawn_result spawn_pair_flow(const std::string & case_name,
    const spawn_opts & opts = {})
{
    namespace fs = std::filesystem;
    auto dir = fs::temp_directory_path()
        / ("libdpf_pair_" + std::to_string(::getpid()) + "_" + case_name);
    fs::create_directories(dir);

    std::string rep = std::to_string(opts.repeat);
    std::string warm = std::to_string(opts.warmup);

    const auto t0 = std::chrono::steady_clock::now();
    pid_t kids[2];
    const char * bins[2] = {"p0", "p1"};
    for (int i = 0; i < 2; ++i)
    {
        kids[i] = fork();
        if (kids[i] < 0)
        {
            fs::remove_all(dir);
            throw std::runtime_error("fork failed");
        }
        if (kids[i] == 0)
        {
            auto bin = party_bin(bins[i]);
            const char * sel = opts.by_tag ? "--tag" : "--case";
            if (opts.metrics)
            {
                execl(bin.c_str(), bin.c_str(),
                    "--dir", dir.string().c_str(),
                    sel, case_name.c_str(),
                    "--repeat", rep.c_str(),
                    "--warmup", warm.c_str(),
                    "--metrics",
                    "--pair",
                    static_cast<char *>(nullptr));
            }
            else
            {
                execl(bin.c_str(), bin.c_str(),
                    "--dir", dir.string().c_str(),
                    sel, case_name.c_str(),
                    "--repeat", rep.c_str(),
                    "--warmup", warm.c_str(),
                    "--pair",
                    static_cast<char *>(nullptr));
            }
            _exit(126);
        }
    }

    spawn_result out;
    for (int i = 0; i < 2; ++i)
    {
        int status = 0;
        if (waitpid(kids[i], &status, 0) < 0)
            out.rc[i] = 125;
        else
            out.rc[i] = WIFEXITED(status) ? WEXITSTATUS(status) : 124;
    }
    out.rc[2] = 0;
    const auto t1 = std::chrono::steady_clock::now();
    out.wall_ms = std::chrono::duration<double, std::milli>(t1 - t0).count();

    if (opts.metrics)
    {
        const char * roles[2] = {"p0", "p1"};
        for (const char * role : roles)
        {
            auto path = dir / (std::string("metrics.") + role);
            std::ifstream in(path);
            std::string line;
            while (std::getline(in, line))
            {
                if (!line.empty())
                    out.metrics_lines.push_back(line);
            }
        }
    }

    fs::remove_all(dir);
    return out;
}

}  // namespace party
}  // namespace dpf

#endif  // LIBDPF_PARTY_SPAWN_HPP__