tcp_server.hpp

/**
 * @file tcp_server.hpp
 * @brief Epoll-based reactor. Accepts connections and dispatches HTTP requests
 *        to web_server via callback. All sockets are non-blocking in ET mode.
 *
 * CSCI 599: Network Systems for Cloud Computing
 * University of Southern California
 */

#ifndef __TCP_SERVER_HPP__
#define __TCP_SERVER_HPP__

#include "log.hpp"
#include "poll.hpp"
#include "sock.hpp"
#include <assert.h>
#include <cerrno>
#include <functional>
#include <iostream>
#include <string>
#include <unordered_map>
#include <atomic>
#include <mutex>

namespace fengcheng {

class tcp_server;
class connection;
using func_t = std::function<void(connection*)>;
using business_func_t = std::function<void(connection*, std::string& request)>;

class connection {
public:
    uint64_t __conn_id;

    connection(int sock = -1) : __sock(sock), __tsvr(nullptr) {
        static std::atomic<uint64_t> id_gen(0);
        __conn_id = ++id_gen;
    }
    ~connection() { }
    void set_callback(func_t recv_cb, func_t send_cb, func_t except_cb) {
        __recv_callback = recv_cb;
        __send_callback = send_cb;
        __except_callback = except_cb;
    }

public:
    int __sock;
    func_t __recv_callback;
    func_t __send_callback;
    func_t __except_callback;
    std::string __in_buffer;
    std::string __out_buffer;
    tcp_server* __tsvr;
};

class tcp_server {
public:
    const static int gport = 8080;
    const static int gnum = 1024; // [FIX] was 128, too small for 60+ concurrent conns

private:
    int __listen_sock;
    int __port;
    __epoll __poll;
    std::unordered_map<int, connection*> __connection_map;
    std::mutex __map_mtx;
    struct epoll_event* __revs;
    int __revs_num;
    business_func_t __callback_func;

public:
    tcp_server(int port = gport) : __port(port), __poll(-1), __revs_num(gnum) { // [FIX] timeout -1, was 0 (busy loop)
        __listen_sock = Sock::Socket();
        Sock::Bind(__listen_sock, __port);
        Sock::Listen(__listen_sock);
        __poll.create_poll();
        __add_connection(__listen_sock, std::bind(&tcp_server::__accepter, this, std::placeholders::_1), nullptr, nullptr);
        __revs = new struct epoll_event[__revs_num];
    }
    ~tcp_server() {
        if (__listen_sock >= 0) close(__listen_sock);
        if (__revs) delete[] __revs;
    }

    void dispather(business_func_t cb) {
        __callback_func = cb;
        while (true) loop_once();
    }

    void loop_once() {
        int n = __poll.wait_poll(__revs, __revs_num);
        for (int i = 0; i < n; i++) {
            int sock = __revs[i].data.fd;
            uint32_t revents = __revs[i].events;
            if (revents & EPOLLERR) revents |= (EPOLLIN | EPOLLOUT);
            if (revents & EPOLLHUP) revents |= (EPOLLIN | EPOLLOUT);

            if (revents & EPOLLIN) {
                connection* conn = nullptr;
                {
                    std::lock_guard<std::mutex> lock(__map_mtx);
                    auto it = __connection_map.find(sock);
                    if (it != __connection_map.end()) conn = it->second;
                }
                if (conn && conn->__recv_callback) conn->__recv_callback(conn);
            }
            if (revents & EPOLLOUT) {
                connection* conn = nullptr;
                {
                    std::lock_guard<std::mutex> lock(__map_mtx);
                    auto it = __connection_map.find(sock);
                    if (it != __connection_map.end()) conn = it->second;
                }
                if (conn && conn->__send_callback) conn->__send_callback(conn);
            }
        }
    }

    void __add_connection(int sock, func_t recv_cb, func_t send_cb, func_t except_cb) {
        Sock::SetNonBlock(sock);
        connection* conn = new connection(sock);
        conn->set_callback(recv_cb, send_cb, except_cb);
        conn->__tsvr = this;
        __poll.add_sock_to_poll(sock, EPOLLIN | EPOLLET);
        {
            std::lock_guard<std::mutex> lock(__map_mtx);
            __connection_map.insert({ sock, conn });
        }
    }

    // [NEW] Transfer ownership of a client fd to a dispatch thread.
    // Removes the fd from epoll and connection_map, deletes the connection
    // object, but does NOT close the fd. The caller is now responsible for close().
    void detach_connection(int sock) {
        connection* conn = nullptr;
        {
            std::lock_guard<std::mutex> lock(__map_mtx);
            auto it = __connection_map.find(sock);
            if (it == __connection_map.end()) return;
            conn = it->second;
            __connection_map.erase(it);
        }
        __poll.delete_from_epoll(sock);
        delete conn; // delete connection object only, NOT the underlying fd
    }

private:
    void __accepter(connection* conn) {
        while (true) {
            std::string client_ip;
            uint16_t client_port;
            int accept_errno = 0;
            int sock = Sock::Accept(conn->__sock, &client_ip, &client_port, &accept_errno);
            if (sock < 0) {
                if (accept_errno == EAGAIN || accept_errno == EWOULDBLOCK) break;
                else if (accept_errno == EINTR) continue;
                else break;
            }
            __add_connection(sock,
                std::bind(&tcp_server::__recver, this, std::placeholders::_1),
                std::bind(&tcp_server::__sender, this, std::placeholders::_1),
                std::bind(&tcp_server::__excepter, this, std::placeholders::_1));
        }
    }

    void __recver(connection* conn) {
        bool is_read_err = false;
        bool is_closed = false;

        while (true) {
            char buffer[4096];
            ssize_t n = recv(conn->__sock, buffer, sizeof(buffer) - 1, 0);
            if (n < 0) {
                if (errno == EAGAIN || errno == EWOULDBLOCK) break;
                else if (errno == EINTR) continue;
                else { is_read_err = true; break; }
            } else if (n == 0) {
                is_closed = true;
                break;
            } else {
                buffer[n] = 0;
                conn->__in_buffer += buffer;
            }
        }

        if (is_read_err || is_closed) {
            conn->__except_callback(conn);
            return;
        }
        __callback_func(conn, conn->__in_buffer);
    }

    void __sender(connection* conn) {
        while (true) {
            ssize_t n = send(conn->__sock, conn->__out_buffer.c_str(), conn->__out_buffer.size(), 0);
            if (n > 0) {
                conn->__out_buffer.erase(0, n);
                if (conn->__out_buffer.empty()) break;
            } else {
                if (errno == EAGAIN || errno == EWOULDBLOCK) break;
                else if (errno == EINTR) continue;
                else { conn->__except_callback(conn); return; }
            }
        }
        if (conn->__out_buffer.empty()) enable_read_write(conn, true, false);
        else enable_read_write(conn, true, true);
    }

public:
    void __excepter(connection* conn) {
        {
            std::lock_guard<std::mutex> lock(__map_mtx);
            if (__connection_map.find(conn->__sock) == __connection_map.end()) return;
            __connection_map.erase(conn->__sock);
        }
        __poll.delete_from_epoll(conn->__sock);
        close(conn->__sock);
        delete conn;
    }

    void enable_read_write(connection* conn, bool readable, bool writable) {
        uint32_t events = (readable ? EPOLLIN : 0) | (writable ? EPOLLOUT : 0);
        __poll.control_poll(conn->__sock, events);
    }

    bool is_valid_connection(int sock, uint64_t conn_id) {
        std::lock_guard<std::mutex> lock(__map_mtx);
        auto it = __connection_map.find(sock);
        if (it == __connection_map.end()) return false;
        return it->second->__conn_id == conn_id;
    }
};

} // namespace fengcheng
#endif