weaseldb/src/server.hpp

#pragma once

#include <atomic>
#include <span>
#include <thread>
#include <vector>

#include "config.hpp"
#include "connection_handler.hpp"
#include "connection_registry.hpp"
#include "reference.hpp"

/**
 * High-performance multi-threaded server for handling network connections.
 *
 * The Server class encapsulates all networking logic including:
 * - Socket management and configuration
 * - Multi-threaded epoll-based I/O multiplexing
 * - Connection lifecycle management
 * - Graceful shutdown handling
 *
 * The server uses a configurable thread pool architecture:
 * - Accept threads: Handle incoming connections with load balancing
 * - Network threads: Process I/O events for established connections
 *
 * All protocol-specific logic is delegated to the provided ConnectionHandler,
 * maintaining clean separation between networking and application logic.
 *
 * IMPORTANT: Server uses a factory pattern and MUST be created via
 * Server::create(). This ensures:
 * - Proper Ref<Server> semantics for reference counting
 * - Safe WeakRef<Server> references from Connection objects
 * - Prevention of accidental stack allocation that would break safety
 * guarantees
 */
struct Server {
  /**
   * Factory method to create a Server instance.
   *
   * This is the only way to create a Server - ensures proper Ref<Server>
   * semantics and prevents accidental stack allocation that would break
   * WeakRef<Server> safety.
   *
   * @param config Server configuration (threads, ports, limits, etc.)
   * @param handler Protocol handler for processing connection data
   * @param listen_fds Vector of file descriptors to accept connections on.
   *                   Server takes ownership and will close them on
   * destruction. Server will set these to non-blocking mode for safe epoll
   * usage. Empty vector means no listening sockets.
   * @return Ref to the newly created Server
   */
  static Ref<Server> create(const weaseldb::Config &config,
                            ConnectionHandler &handler,
                            const std::vector<int> &listen_fds);

  /**
   * Destructor ensures proper cleanup of all resources.
   */
  ~Server();

  /**
   * Start the server and begin accepting connections.
   *
   * This method:
   * - Creates and configures the listen socket
   * - Starts all worker threads
   * - Blocks until shutdown() is called or an error occurs
   *
   * Aborts the process on socket creation or configuration errors
   */
  void run();

  /**
   * Initiate graceful server shutdown.
   *
   * This method is async-signal-safe and can be called from signal handlers.
   * It signals all threads to stop processing and begin cleanup.
   *
   * The run() method will return after all threads have completed shutdown.
   */
  void shutdown();

  /**
   * Creates a local connection using socketpair() for testing or local IPC.
   *
   * Creates a socketpair, registers one end as a Connection in the server,
   * and returns the other end to the caller for communication.
   *
   * The caller takes ownership of the returned file descriptor and must close
   * it.
   *
   * @return File descriptor for the client end of the socketpair, or -1 on
   * error
   */
  int create_local_connection();

private:
  friend struct Connection;
  /**
   * Private constructor - use create() factory method instead.
   *
   * @param config Server configuration (threads, ports, limits, etc.)
   * @param handler Protocol handler for processing connection data. Must
   * outlive the server.
   * @param listen_fds Vector of file descriptors to accept connections on.
   * Server takes ownership and will close them on
   * destruction. Server will set these to non-blocking mode for safe epoll
   * usage.
   */
  explicit Server(const weaseldb::Config &config, ConnectionHandler &handler,
                  const std::vector<int> &listen_fds);
  template <typename T, typename... Args>
  friend Ref<T> make_ref(Args &&...args);

  WeakRef<Server> self_;

  weaseldb::Config config_;
  ConnectionHandler &handler_;

  // Connection registry
  ConnectionRegistry connection_registry_;

  // Connection management
  std::atomic<int64_t> connection_id_{0};
  std::atomic<int> active_connections_{0};

  // Round-robin counter for connection distribution
  std::atomic<size_t> connection_distribution_counter_{0};

  // Shutdown coordination
  int shutdown_pipe_[2] = {-1, -1};

  // Multiple epoll file descriptors to reduce contention
  std::vector<int> epoll_fds_;
  std::vector<int>
      listen_fds_; // FDs to accept connections on (Server owns these)

  // Private helper methods
  void setup_shutdown_pipe();
  void setup_signal_handling();
  void create_epoll_instances();
  void start_io_threads(std::vector<std::thread> &threads);

  // Helper to get epoll fd for a thread using round-robin
  int get_epoll_for_thread(int thread_id) const;

  // Helper for processing connection I/O
  void process_connection_reads(Ref<Connection> &conn, int events);
  void process_connection_writes(Ref<Connection> &conn, int events);

  void close_connection(Ref<Connection> &conn);

  // Helper for processing a batch of connections with their events
  void process_connection_batch(std::span<Ref<Connection>> batch,
                                std::span<const int> events);

  // Make non-copyable and non-movable
  Server(const Server &) = delete;
  Server &operator=(const Server &) = delete;
  Server(Server &&) = delete;
  Server &operator=(Server &&) = delete;
};