CPP_Core/Core/socket/TCP_Server.cpp

#include "TCP_Server.h"

#include <array>
#include <string_view>

namespace Psc::asio_socket {

TCP_Server::~TCP_Server() { close(); }

void TCP_Server::tick() {
  if (state == Working) {
    start_accept();
    flush_clients();
  }

  io_context.restart();
  io_context.poll();

  if (state == Working) {
    flush_clients();
    cleanup_closed_clients();
  }
}

TCP_Server &TCP_Server::set_tcp_no_delay(bool value) {
  no_delay = value;
  return *this;
}

TCP_Server &TCP_Server::set_connect_system_buffer_size(size_t size) {
  connect_system_buffer_size = size;
  return *this;
}

TCP_Server &TCP_Server::set_connect_user_buffer_size(size_t size) {
  connect_user_buffer_size = size;
  return *this;
}

TCP_Server &TCP_Server::set_recv_system_buffer_size(size_t size) {
  recv_system_buffer_size = size;
  return *this;
}

TCP_Server &TCP_Server::create() {
  acceptor = std::make_unique<asio::ip::tcp::acceptor>(io_context);
  socket_fd = socket_id(acceptor.get());
  state = Not_bind;
  return *this;
}

bool TCP_Server::listen(std::string_view ip, std::uint32_t port) {
  return listen(Sockaddr_In(ip, port));
}

bool TCP_Server::listen(const Sockaddr_In &addr) {
  if (!acceptor)
    create();
  if (addr.ip.empty() || addr.port == 0) {
    state = Not_Set_Field;
    return false;
  }

  asio::error_code ec;
  const auto address_value = asio::ip::make_address(addr.ip, ec);
  if (ec) {
    state = Not_bind;
    return false;
  }

  asio::ip::tcp::endpoint endpoint(address_value,
                                   static_cast<unsigned short>(addr.port));
  acceptor->open(endpoint.protocol(), ec);
  if (ec)
    return false;
  acceptor->set_option(asio::socket_base::reuse_address(true), ec);
  acceptor->bind(endpoint, ec);
  if (ec) {
    state = Not_bind;
    return false;
  }
  acceptor->listen(asio::socket_base::max_listen_connections, ec);
  if (ec) {
    state = Not_Listen;
    return false;
  }
  address = addr;
  state = Working;
  start_accept();
  return true;
}

void TCP_Server::close() {
  accept_pending = false;
  if (pending_accept_socket) {
    asio::error_code ec;
    pending_accept_socket->cancel(ec);
    pending_accept_socket->close(ec);
  }
  pending_accept_socket.reset();

  for (auto &[_, conn] : tcp_clients) {
    close_client(conn);
  }
  tcp_clients.clear();
  accepted_clients.clear();

  if (acceptor) {
    asio::error_code ec;
    acceptor->cancel(ec);
    acceptor->close(ec);
  }
  socket_fd = static_cast<Socket_FD>(-1);
  state = Not_Created;
}

std::shared_ptr<TCP_Connect> TCP_Server::accept() const {
  while (!accepted_clients.empty()) {
    auto conn = accepted_clients.front().lock();
    accepted_clients.erase(accepted_clients.begin());
    if (conn && !conn->closing)
      return conn;
  }
  return nullptr;
}

void TCP_Server::start_accept() {
  if (accept_pending || !acceptor || state != Working || !acceptor->is_open())
    return;

  pending_accept_socket = std::make_shared<asio::ip::tcp::socket>(io_context);
  accept_pending = true;
  acceptor->async_accept(
      *pending_accept_socket, [this](const asio::error_code &ec) {
        accept_pending = false;
        if (state != Working)
          return;
        if (ec == asio::error::operation_aborted)
          return;
        if (!ec && pending_accept_socket) {
          asio::error_code option_ec;
          pending_accept_socket->set_option(asio::ip::tcp::no_delay(no_delay),
                                            option_ec);
          pending_accept_socket->set_option(
              asio::socket_base::send_buffer_size(
                  static_cast<int>(connect_system_buffer_size)),
              option_ec);
          pending_accept_socket->set_option(
              asio::socket_base::receive_buffer_size(
                  static_cast<int>(recv_system_buffer_size)),
              option_ec);

          auto conn = std::make_shared<TCP_Connect>();
          conn->socket = pending_accept_socket;
          conn->send_buffer.init(connect_user_buffer_size);
          conn->recv_buffer.init(connect_user_buffer_size);
          conn->info.fd = socket_id(conn->socket.get());
          conn->info.sockaddr =
              endpoint_to_sockaddr(conn->socket->remote_endpoint(option_ec));
          tcp_clients[conn->info.fd] = conn;
          accepted_clients.push_back(conn);
          start_read(conn);
          start_write(conn);
        }
        pending_accept_socket.reset();
        start_accept();
      });
}

void TCP_Server::start_read(const std::shared_ptr<TCP_Connect> &conn) {
  if (!conn || conn->closing || conn->read_pending || !conn->socket ||
      !conn->socket->is_open())
    return;

  conn->read_pending = true;
  conn->socket->async_read_some(
      asio::buffer(conn->recv_storage),
      [this, conn](const asio::error_code &ec, std::size_t n) {
        conn->read_pending = false;
        if (conn->closing)
          return;
        if (ec == asio::error::operation_aborted)
          return;
        if (ec == asio::error::eof || ec == asio::error::connection_reset ||
            ec) {
          close_client(conn);
          return;
        }

        conn->recv_buffer.write_best_effort(conn->recv_storage.data(), n);
        start_read(conn);
      });
}

void TCP_Server::start_write(const std::shared_ptr<TCP_Connect> &conn) {
  if (!conn || conn->closing || conn->write_pending || !conn->socket ||
      !conn->socket->is_open())
    return;

  conn->send_storage.resize(16 * 1024);
  auto size = conn->send_buffer.peek_best_effort(conn->send_storage.data(),
                                                 conn->send_storage.size());
  if (size == 0)
    return;
  conn->send_storage.resize(size);

  conn->write_pending = true;
  conn->socket->async_write_some(
      asio::buffer(conn->send_storage.data(), conn->send_storage.size()),
      [this, conn](const asio::error_code &ec, std::size_t sent) {
        conn->write_pending = false;
        if (conn->closing)
          return;
        if (ec == asio::error::operation_aborted)
          return;
        if (ec) {
          close_client(conn);
          return;
        }

        conn->send_buffer.skip(sent);
        conn->push_speed.update(static_cast<unsigned int>(sent));
        conn->send_num.update(static_cast<double>(sent));
        if (sent != 0)
          start_write(conn);
      });
}

void TCP_Server::close_client(const std::shared_ptr<TCP_Connect> &conn) {
  if (!conn || conn->closing)
    return;
  conn->closing = true;
  if (conn->socket) {
    asio::error_code ec;
    conn->socket->cancel(ec);
    conn->socket->shutdown(asio::ip::tcp::socket::shutdown_both, ec);
    conn->socket->close(ec);
  }
}

void TCP_Server::cleanup_closed_clients() {
  for (auto it = tcp_clients.begin(); it != tcp_clients.end();) {
    if (!it->second || it->second->closing || !it->second->socket ||
        !it->second->socket->is_open()) {
      it = tcp_clients.erase(it);
    } else {
      ++it;
    }
  }

  for (auto it = accepted_clients.begin(); it != accepted_clients.end();) {
    auto conn = it->lock();
    if (!conn || conn->closing) {
      it = accepted_clients.erase(it);
    } else {
      ++it;
    }
  }
}

void TCP_Server::flush_clients() {
  start_accept();
  for (auto &[_, conn] : tcp_clients) {
    start_read(conn);
    start_write(conn);
  }
}

void TCP_Server::write_to_all_clients(std::string_view data) {
  if (data.empty())
    return;
  for (auto &[_, conn] : tcp_clients) {
    if (!conn || conn->closing)
      continue;
    auto written =
        conn->send_buffer.write_best_effort(data.data(), data.size());
    if (written < data.size())
      conn->lose_speed.update(static_cast<unsigned int>(data.size() - written));
    start_write(conn);
  }
}

std::vector<TCP_Server::Read_Info> TCP_Server::read_from_all_clients() {
  std::vector<Read_Info> ret;
  std::array<char, 16 * 1024> buffer{};

  for (auto &[_, conn] : tcp_clients) {
    if (!conn || conn->closing)
      continue;

    std::string data;
    for (;;) {
      auto n = conn->recv_buffer.read_best_effort(buffer.data(), buffer.size());
      if (n == 0)
        break;
      data.append(buffer.data(), n);
    }
    if (!data.empty())
      ret.push_back({conn->info, std::move(data)});
  }

  cleanup_closed_clients();
  return ret;
}

std::vector<Socket_FD> TCP_Server::client_fds() {
  std::vector<Socket_FD> ret;
  for (const auto &[fd, conn] : tcp_clients) {
    if (conn && !conn->closing)
      ret.push_back(fd);
  }
  return ret;
}

std::vector<std::shared_ptr<TCP_Connect>> TCP_Server::get_all_clients() {
  std::vector<std::shared_ptr<TCP_Connect>> ret;
  for (const auto &[_, conn] : tcp_clients) {
    if (conn && !conn->closing)
      ret.push_back(conn);
  }
  return ret;
}

std::string TCP_Server::to_string() {
  return "TCP_Server:[" + (address ? address->to_string() : std::string{}) +
         "]";
}

} // namespace Psc::asio_socket