LNXRNT/Sources/connection_pool.cpp

#include "connection_pool.h"
#include "socket_optimization.h"
#include <kinc/log.h>

#ifdef _WIN32
#include <ws2tcpip.h>
#pragma comment(lib, "ws2_32.lib")
#else
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <netinet/tcp.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <sys/socket.h>
#define INVALID_SOCKET -1
#define SOCKET_ERROR -1
#endif

// global
ConnectionPool g_connection_pool;

ConnectionPool::~ConnectionPool() {
    std::lock_guard<std::mutex> lock(pool_mutex_);
    
    for (auto& host_pair : pool_) {
        while (!host_pair.second.empty()) {
            auto conn = host_pair.second.front();
            host_pair.second.pop();
#ifdef _WIN32
            closesocket(conn.socket);
#else
            close(conn.socket);
#endif
        }
    }
}

socket_t ConnectionPool::acquire_connection(const std::string& host, int port, bool is_ssl,
                                            std::atomic<bool>* cancel) {
    std::string key = make_key(host, port, is_ssl);
    
    {
        std::lock_guard<std::mutex> lock(pool_mutex_);
        auto it = pool_.find(key);
        if (it != pool_.end() && !it->second.empty()) {
            auto conn = it->second.front();
            it->second.pop();
            
            if (is_socket_alive(conn.socket)) {
                kinc_log(KINC_LOG_LEVEL_INFO, "Reusing pooled connection for %s:%d", host.c_str(), port);
                return conn.socket;
            } else {
#ifdef _WIN32
                closesocket(conn.socket);
#else
                close(conn.socket);
#endif
            }
        }
    }
    
    socket_t sock = ::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sock == INVALID_SOCKET) {
        kinc_log(KINC_LOG_LEVEL_ERROR, "Failed to create socket for %s:%d", host.c_str(), port);
        return INVALID_SOCKET;
    }
    
    SocketOptimization::optimizeHttp(static_cast<int>(sock));
    
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);
    
    struct hostent* host_entry = gethostbyname(host.c_str());
    if (!host_entry) {
        kinc_log(KINC_LOG_LEVEL_ERROR, "Failed to resolve hostname %s", host.c_str());
#ifdef _WIN32
        closesocket(sock);
#else
        close(sock);
#endif
        return INVALID_SOCKET;
    }
    
    memcpy(&server_addr.sin_addr, host_entry->h_addr, host_entry->h_length);
    
#ifdef _WIN32
    u_long nb_mode = 1;
    ioctlsocket(sock, FIONBIO, &nb_mode);
#else
    int old_flags = fcntl(sock, F_GETFL, 0);
    fcntl(sock, F_SETFL, old_flags | O_NONBLOCK);
#endif

    bool connected = (::connect(sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) == 0);
    if (!connected) {
#ifdef _WIN32
        int connect_err = WSAGetLastError();
        bool in_progress = (connect_err == WSAEWOULDBLOCK || connect_err == WSAEINPROGRESS);
#else
        bool in_progress = (errno == EINPROGRESS);
#endif
        auto deadline = std::chrono::steady_clock::now() +
                        std::chrono::milliseconds(15000);
        while (in_progress) {
            if (cancel != nullptr && cancel->load()) {
                in_progress = false;
                break;
            }
            if (std::chrono::steady_clock::now() >= deadline) {
                in_progress = false;
                break;
            }
            fd_set wfds, efds;
            FD_ZERO(&wfds);
            FD_ZERO(&efds);
            FD_SET(sock, &wfds);
            FD_SET(sock, &efds);
            struct timeval tv;
            tv.tv_sec = 0;
            tv.tv_usec = 50000;
#ifdef _WIN32
            int sel = ::select(0, nullptr, &wfds, &efds, &tv);
#else
            int sel = ::select(static_cast<int>(sock) + 1, nullptr, &wfds, &efds, &tv);
#endif
            if (sel > 0) {
                int so_err = 0;
                socklen_t so_len = sizeof(so_err);
                getsockopt(sock, SOL_SOCKET, SO_ERROR, (char*)&so_err, &so_len);
                connected = (so_err == 0);
                break;
            }
            if (sel < 0) break;
        }
    }

#ifdef _WIN32
    nb_mode = 0;
    ioctlsocket(sock, FIONBIO, &nb_mode);
#else
    fcntl(sock, F_SETFL, old_flags);
#endif

    if (!connected) {
        kinc_log(KINC_LOG_LEVEL_ERROR, "Failed to connect to %s:%d", host.c_str(), port);
#ifdef _WIN32
        closesocket(sock);
#else
        close(sock);
#endif
        return INVALID_SOCKET;
    }
    
    kinc_log(KINC_LOG_LEVEL_INFO, "Created new connection for %s:%d", host.c_str(), port);
    return sock;
}

void ConnectionPool::release_connection(const std::string& host, int port, socket_t socket, bool is_ssl) {
    if (socket == INVALID_SOCKET) return;
    
    std::string key = make_key(host, port, is_ssl);
    
    std::lock_guard<std::mutex> lock(pool_mutex_);
    auto& queue = pool_[key];
    
    if (queue.size() >= MAX_CONNECTIONS_PER_HOST) {
#ifdef _WIN32
        closesocket(socket);
#else
        close(socket);
#endif
        return;
    }
    
    PooledConnection conn;
    conn.socket = socket;
    conn.last_used = std::chrono::steady_clock::now();
    conn.is_ssl = is_ssl;
    
    queue.push(conn);
    kinc_log(KINC_LOG_LEVEL_INFO, "Pooled connection for %s:%d", host.c_str(), port);
}

void ConnectionPool::cleanup_expired() {
    std::lock_guard<std::mutex> lock(pool_mutex_);
    auto now = std::chrono::steady_clock::now();
    
    for (auto& host_pair : pool_) {
        auto& queue = host_pair.second;
        std::queue<PooledConnection> new_queue;
        
        while (!queue.empty()) {
            auto conn = queue.front();
            queue.pop();
            
            auto age = std::chrono::duration_cast<std::chrono::seconds>(now - conn.last_used);
            if (age.count() < CONNECTION_TIMEOUT_SECONDS && is_socket_alive(conn.socket)) {
                new_queue.push(conn);
            } else {
#ifdef _WIN32
                closesocket(conn.socket);
#else
                close(conn.socket);
#endif
            }
        }
        
        queue = std::move(new_queue);
    }
}

size_t ConnectionPool::get_pool_size() const {
    std::lock_guard<std::mutex> lock(pool_mutex_);
    size_t total = 0;
    for (const auto& host_pair : pool_) {
        total += host_pair.second.size();
    }
    return total;
}

std::string ConnectionPool::make_key(const std::string& host, int port, bool is_ssl) const {
    return host + ":" + std::to_string(port) + (is_ssl ? ":ssl" : "");
}

bool ConnectionPool::is_socket_alive(socket_t socket) const {
    fd_set rfds;
    FD_ZERO(&rfds);
    FD_SET(socket, &rfds);
    struct timeval tv;
    tv.tv_sec = 0;
    tv.tv_usec = 0;
#ifdef _WIN32
    int sel = ::select(0, &rfds, nullptr, nullptr, &tv);
#else
    int sel = ::select(static_cast<int>(socket) + 1, &rfds, nullptr, nullptr, &tv);
#endif
    if (sel < 0) return false;
    if (sel == 0) return true; 
    char test_byte;
    int result = ::recv(socket, &test_byte, 1, MSG_PEEK);
    if (result == 0) return false;

#ifdef _WIN32
    if (result == SOCKET_ERROR) {
        int error = WSAGetLastError();
        return (error == WSAEWOULDBLOCK);
    }
#else
    if (result < 0) {
        return (errno == EWOULDBLOCK || errno == EAGAIN);
    }
#endif

    return false;
}