proxy_v2_threadpool.cpp

// proxy_v2_threadpool.cpp
// Usage: ./proxy_v2 <listen_port> <backend_host> <backend_port> [num_threads]

#include <arpa/inet.h>
#include <cstring>
#include <iostream>
#include <netdb.h>
#include <netinet/in.h>
#include <string>
#include <sys/socket.h>
#include <unistd.h>

#include "thread_pool.h"

namespace {

constexpr int kListenBacklog = 128;
constexpr size_t kBufferSize = 4096;
constexpr size_t kDefaultNumThreads = 4;

int connect_to_backend(const std::string &host, int port) {
  addrinfo hints{};
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_STREAM;

  addrinfo *result = nullptr;
  std::string port_str = std::to_string(port);
  int rc = getaddrinfo(host.c_str(), port_str.c_str(), &hints, &result);
  if (rc != 0) {
    std::cerr << "getaddrinfo failed for " << host << ":" << port << " - "
              << gai_strerror(rc) << "\n";
    return -1;
  }

  int fd = -1;
  for (addrinfo *p = result; p != nullptr; p = p->ai_next) {
    fd = socket(p->ai_family, p->ai_socktype, p->ai_protocol);
    if (fd < 0)
      continue;
    if (connect(fd, p->ai_addr, p->ai_addrlen) == 0) {
      break;
    }
    close(fd);
    fd = -1;
  }
  freeaddrinfo(result);
  return fd;
}

ssize_t forward_once(int src_fd, int dst_fd) {
  char buffer[kBufferSize];
  ssize_t n_read = recv(src_fd, buffer, sizeof(buffer), 0);
  if (n_read <= 0) {
    return n_read;
  }

  ssize_t total_sent = 0;
  while (total_sent < n_read) {
    ssize_t n_sent = send(dst_fd, buffer + total_sent, n_read - total_sent, 0);
    if (n_sent < 0) {
      perror("send");
      return -1;
    }
    total_sent += n_sent;
  }
  return total_sent;
}

void handle_client(int client_fd, const std::string &backend_host,
                   int backend_port) {
  int backend_fd = connect_to_backend(backend_host, backend_port);
  if (backend_fd < 0) {
    std::cerr << "[proxy] failed to connect to backend " << backend_host << ":"
              << backend_port << "\n";
    close(client_fd);
    return;
  }

  ssize_t forwarded = forward_once(client_fd, backend_fd);
  if (forwarded <= 0) {
    close(client_fd);
    close(backend_fd);
    return;
  }

  forward_once(backend_fd, client_fd);

  close(backend_fd);
  close(client_fd);
}

} // namespace

int main(int argc, char *argv[]) {
  if (argc < 4) {
    std::cerr << "Usage: " << argv[0]
              << " <listen_port> <backend_host> <backend_port> [num_threads]\n";
    return 1;
  }

  int listen_port = std::stoi(argv[1]);
  std::string backend_host = argv[2];
  int backend_port = std::stoi(argv[3]);
  size_t num_threads = (argc >= 5) ? static_cast<size_t>(std::stoi(argv[4]))
                                   : kDefaultNumThreads;

  int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
  if (listen_fd < 0) {
    perror("socket");
    return 1;
  }

  int opt = 1;
  if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
    perror("setsockopt");
  }

  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_addr.s_addr = INADDR_ANY;
  addr.sin_port = htons(listen_port);

  if (bind(listen_fd, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) < 0) {
    perror("bind");
    close(listen_fd);
    return 1;
  }

  if (listen(listen_fd, kListenBacklog) < 0) {
    perror("listen");
    close(listen_fd);
    return 1;
  }

  std::cout << "[proxy v2 - thread pool, " << num_threads
            << " threads] listening on port " << listen_port
            << ", forwarding to " << backend_host << ":" << backend_port
            << std::endl;

  // Note: Pool currently runs until process termination. Graceful SIGINT
  // shutdown pending.
  ThreadPool pool(num_threads);

  while (true) {
    sockaddr_in client_addr{};
    socklen_t client_len = sizeof(client_addr);
    int client_fd = accept(
        listen_fd, reinterpret_cast<sockaddr *>(&client_addr), &client_len);
    if (client_fd < 0) {
      perror("accept");
      continue;
    }

    char client_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, sizeof(client_ip));
    std::cout << "[proxy] accepted connection from " << client_ip << ":"
              << ntohs(client_addr.sin_port) << std::endl;

    // Dispatch connection to worker pool and immediately resume accepting
    pool.submit([client_fd, &backend_host, backend_port] {
      handle_client(client_fd, backend_host, backend_port);
    });
  }

  close(listen_fd);
  return 0;
}