libdpf/include/dpf/net/connect.hpp

/// @file dpf/net/connect.hpp
/// @brief Deadline-bounded connect, accept, and handshake I/O for setup.
/// @details Setup is blocking by design (it happens before protocol traffic),
///          but every step has a wall-clock bound. `connect_until` retries
///          refused connections until the deadline, so processes may start in
///          any order.
#ifndef LIBDPF_INCLUDE_DPF_NET_CONNECT_HPP__
#define LIBDPF_INCLUDE_DPF_NET_CONNECT_HPP__

#include <cerrno>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <stdexcept>
#include <string>
#include <system_error>
#include <thread>

#include <fcntl.h>
#include <poll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include "dpf/net/asio_ns.hpp"

#include "dpf/log.hpp"

namespace dpf
{
namespace net
{

using setup_clock = std::chrono::steady_clock;

namespace detail
{

inline int remaining_ms(setup_clock::time_point deadline)
{
    const auto left = std::chrono::duration_cast<std::chrono::milliseconds>(
        deadline - setup_clock::now());
    if (left.count() <= 0)
        return 0;
    if (left.count() > 0x7fffffff)
        return 0x7fffffff;
    return static_cast<int>(left.count());
}

/// @brief Wait until `fd` is ready for `events` or throw at `deadline`.
inline void wait_fd(int fd, short events, setup_clock::time_point deadline,
    const std::string & what)
{
    for (;;)
    {
        const int ms = remaining_ms(deadline);
        if (ms == 0)
            throw std::system_error(std::make_error_code(std::errc::timed_out),
                what + ": timed out");
        pollfd pfd{};
        pfd.fd = fd;
        pfd.events = events;
        const int rc = ::poll(&pfd, 1, ms);
        if (rc > 0)
        {
            if ((pfd.revents & (POLLERR | POLLNVAL)) != 0
                && (pfd.revents & events) == 0)
                throw std::system_error(
                    std::make_error_code(std::errc::connection_reset), what);
            return;
        }
        if (rc < 0 && errno != EINTR)
            throw std::system_error(errno, std::generic_category(), what);
    }
}

inline void send_all_until(int fd, const void * data, std::size_t n,
    setup_clock::time_point deadline, const std::string & what)
{
    const auto * p = static_cast<const std::uint8_t *>(data);
    std::size_t done = 0;
    while (done < n)
    {
        const ssize_t r = ::send(fd, p + done, n - done,
            MSG_DONTWAIT | MSG_NOSIGNAL);
        if (r > 0)
        {
            done += static_cast<std::size_t>(r);
            continue;
        }
        if (r < 0 && (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR))
        {
            wait_fd(fd, POLLOUT, deadline, what);
            continue;
        }
        throw std::system_error(errno, std::generic_category(), what);
    }
}

inline void recv_all_until(int fd, void * data, std::size_t n,
    setup_clock::time_point deadline, const std::string & what)
{
    auto * p = static_cast<std::uint8_t *>(data);
    std::size_t done = 0;
    while (done < n)
    {
        const ssize_t r = ::recv(fd, p + done, n - done, MSG_DONTWAIT);
        if (r > 0)
        {
            done += static_cast<std::size_t>(r);
            continue;
        }
        if (r == 0)
            throw std::system_error(std::make_error_code(
                std::errc::connection_reset), what + ": peer closed");
        if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
        {
            wait_fd(fd, POLLIN, deadline, what);
            continue;
        }
        throw std::system_error(errno, std::generic_category(), what);
    }
}

inline void put_u32(std::uint8_t * dst, std::uint32_t v) noexcept
{
    dst[0] = static_cast<std::uint8_t>(v & 0xffu);
    dst[1] = static_cast<std::uint8_t>((v >> 8) & 0xffu);
    dst[2] = static_cast<std::uint8_t>((v >> 16) & 0xffu);
    dst[3] = static_cast<std::uint8_t>((v >> 24) & 0xffu);
}

inline std::uint32_t get_u32(const std::uint8_t * src) noexcept
{
    return static_cast<std::uint32_t>(src[0]) | (static_cast<std::uint32_t>(src[1]) << 8)
        | (static_cast<std::uint32_t>(src[2]) << 16)
        | (static_cast<std::uint32_t>(src[3]) << 24);
}

}  // namespace detail

/// @brief Connect `sock` to `host:port`, retrying refusals until `budget`.
inline void connect_until(asio::ip::tcp::socket & sock, const std::string & host,
    unsigned short port, std::chrono::milliseconds budget)
{
    const auto started = setup_clock::now();
    const auto deadline = started + budget;
    const std::string what = "connect " + host + ":" + std::to_string(port);
    asio::ip::tcp::resolver res(sock.get_executor());
    std::error_code rec;
    auto eps = res.resolve(host, std::to_string(port), rec);
    if (rec)
        throw std::system_error(rec, what + ": resolve");
    std::error_code last = std::make_error_code(std::errc::timed_out);
    std::size_t attempts = 0;
    for (;;)
    {
        for (const auto & entry : eps)
        {
            std::error_code ec;
            if (sock.is_open())
                sock.close(ec);
            sock.open(entry.endpoint().protocol(), ec);
            if (ec)
            {
                last = ec;
                continue;
            }
            const int fd = sock.native_handle();
            const int fl = ::fcntl(fd, F_GETFL, 0);
            (void)::fcntl(fd, F_SETFL, fl | O_NONBLOCK);
            ++attempts;
            int rc = ::connect(fd, entry.endpoint().data(),
                static_cast<socklen_t>(entry.endpoint().size()));
            int err = rc == 0 ? 0 : errno;
            if (rc != 0 && err == EINPROGRESS)
            {
                pollfd pfd{};
                pfd.fd = fd;
                pfd.events = POLLOUT;
                const int prc = ::poll(&pfd, 1, detail::remaining_ms(deadline));
                if (prc <= 0)
                    err = ETIMEDOUT;
                else
                {
                    socklen_t len = sizeof(err);
                    err = 0;
                    ::getsockopt(fd, SOL_SOCKET, SO_ERROR, &err, &len);
                }
            }
            if (err == 0)
            {
                (void)::fcntl(fd, F_SETFL, fl);
                DPF_LOG(debug, "connect").kv("target", host + ":" + std::to_string(port))
                    .kv("resolved", entry.endpoint().address().to_string())
                    .kv("attempts", attempts)
                    .kv("elapsed_ms", std::chrono::duration<double, std::milli>(
                        setup_clock::now() - started).count());
                return;
            }
            last = std::error_code(err, std::generic_category());
            sock.close(ec);
        }
        if (setup_clock::now() >= deadline)
        {
            DPF_LOG(error, "connect.failed").kv("target", host + ":" + std::to_string(port))
                .kv("attempts", attempts).kv("budget_ms", budget.count())
                .kv("last_error", last.message());
            throw std::system_error(last, what + " failed after "
                + std::to_string(budget.count()) + " ms");
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(20));
    }
}

/// @brief Accept one connection on `acc` within `budget`.
inline void accept_until(asio::ip::tcp::acceptor & acc,
    asio::ip::tcp::socket & sock, std::chrono::milliseconds budget)
{
    const auto deadline = setup_clock::now() + budget;
    const std::string what = "accept on port "
        + std::to_string(acc.local_endpoint().port());
    detail::wait_fd(acc.native_handle(), POLLIN, deadline, what);
    acc.accept(sock);
}

/// @brief Open, bind, and listen on `port` (0 = ephemeral), reusing the address.
inline void open_listener(asio::ip::tcp::acceptor & acc, unsigned short port)
{
    const asio::ip::tcp::endpoint ep(asio::ip::tcp::v4(), port);
    acc.open(ep.protocol());
    acc.set_option(asio::socket_base::reuse_address(true));
    acc.bind(ep);
    acc.listen();
}

/// @brief Exchange one `u32` each way within `budget` (setup handshake).
inline std::uint32_t exchange_u32(int fd, std::uint32_t mine,
    std::chrono::milliseconds budget, const std::string & what)
{
    const auto deadline = setup_clock::now() + budget;
    std::uint8_t out[4];
    detail::put_u32(out, mine);
    detail::send_all_until(fd, out, 4, deadline, what);
    std::uint8_t in[4];
    detail::recv_all_until(fd, in, 4, deadline, what);
    return detail::get_u32(in);
}

/// @brief Send then receive a fixed-size record within `budget`.
inline void exchange_record(int fd, const void * mine, void * theirs,
    std::size_t n, std::chrono::milliseconds budget, const std::string & what)
{
    const auto deadline = setup_clock::now() + budget;
    detail::send_all_until(fd, mine, n, deadline, what);
    detail::recv_all_until(fd, theirs, n, deadline, what);
}

}  // namespace net
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_NET_CONNECT_HPP__