socket_backend.hpp

 
#pragma once
 
#include <array>
#include <memory>
#include <mutex>
#include <string>
#include <string_view>
#include <thread>
#include <vector>
 
#include <rapidsmpf/bootstrap/backend.hpp>
#include <rapidsmpf/bootstrap/bootstrap.hpp>
 
// NOTE: Do not use RAPIDSMPF_EXPECTS or RAPIDSMPF_FAIL in this file.
// Using these macros introduces a CUDA dependency via rapidsmpf/error.hpp.
// Prefer throwing standard exceptions instead.
 
namespace rapidsmpf::bootstrap::detail {
 
class SocketServer {
  public:
    explicit SocketServer(int nranks);
 
    ~SocketServer();
 
    [[nodiscard]] std::string const& address() const noexcept {
        return address_;
    }
 
    [[nodiscard]] std::string const& token() const noexcept {
        return token_;
    }
 
    SocketServer(SocketServer const&) = delete;
    SocketServer& operator=(SocketServer const&) = delete;
    SocketServer(SocketServer&&) = delete;
    SocketServer& operator=(SocketServer&&) = delete;
 
  private:
    struct State;
    std::shared_ptr<State> state_;
    int listen_fd_{-1};
    std::array<int, 2> wakeup_pipe_{-1, -1};
    std::string address_;
    std::string token_;
    std::thread accept_thread_;
    std::mutex handler_mutex_;
    std::vector<std::thread> handler_threads_;
 
    void accept_loop();
    void handle_connection(int client_fd);
};
 
class SocketBackend : public Backend {
  public:
    explicit SocketBackend(Context ctx);
 
    ~SocketBackend() override;
 
    void put(std::string const& key, std::string_view value) override;
 
    std::string get(std::string const& key, Duration timeout) override;
 
    void barrier() override;
 
    void sync() override;
 
  private:
    Context ctx_;
    int fd_{-1};
 
    void send_line(std::string const& line);
    std::string recv_line();
    void send_bytes(void const* data, std::size_t n);
    void recv_bytes(void* data, std::size_t n);
};
 
}  // namespace rapidsmpf::bootstrap::detail