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