server_concurrent.cpp

#include <iostream>
#include <vector>
#include <cstring>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <sys/epoll.h>
#include <thread>
#include <mutex>
#include <queue>
#include <condition_variable>
#include <functional>
#include "hash_map.hpp"

const int MAX_EVENTS = 64;
const int PORT = 8080;
const int BUFFER_SIZE = 1024;

// 1. Sharded Mutex to minimize contention
const int NUM_SHARDS = 16;
std::mutex shard_mutexes[NUM_SHARDS];

int get_shard(const std::string& key) {
    return (int)(std::hash<std::string>{}(key) % NUM_SHARDS);
}

// 2. Thread Pool for processing requests
class ThreadPool {
    std::vector<std::thread> workers;
    std::queue<std::function<void()>> tasks;
    std::mutex queue_mutex;
    std::condition_variable cv;
    bool stop = false;

public:
    ThreadPool(size_t threads) {
        for(size_t i = 0; i < threads; ++i)
            workers.emplace_back([this] {
                while(true) {
                    std::function<void()> task;
                    {
                        std::unique_lock<std::mutex> lock(this->queue_mutex);
                        this->cv.wait(lock, [this]{ return this->stop || !this->tasks.empty(); });
                        if(this->stop && this->tasks.empty()) return;
                        task = std::move(this->tasks.front());
                        this->tasks.pop();
                    }
                    task();
                }
            });
    }

    void enqueue(std::function<void()> task) {
        {
            std::unique_lock<std::mutex> lock(queue_mutex);
            tasks.emplace(std::move(task));
        }
        cv.notify_one();
    }

    ~ThreadPool() {
        { std::unique_lock<std::mutex> lock(queue_mutex); stop = true; }
        cv.notify_all();
        for(std::thread &worker : workers) worker.join();
    }
};

// Global Store
LowLatencyHashMap global_kv_store(100000);

// Logic to process a request (this runs in a worker thread)
void process_request(int fd, const std::string& command, const std::string& key, const std::string& value) {
    if (key.empty()) {
        std::string response = "ERROR: Missing key\n";
        write(fd, response.c_str(), response.size());
        return;
    }

    int shard = get_shard(key);
    
    // 3. Include server logs for hash functions
    std::cout << "[Worker] Command: " << command << " | Key: " << key 
              << " | Hash Shard: " << shard << std::endl;

    std::lock_guard<std::mutex> lock(shard_mutexes[shard]);
    
    std::string response;
    
    // 1 & 2. Handle SET, GET, and DEL methods
    if (command == "SET") {
        global_kv_store.insert(key, value);
        response = "OK\n";
    } else if (command == "GET") {
        std::string out_val;
        if (global_kv_store.get(key, out_val)) {
            response = out_val + "\n";
        } else {
            response = "(nil)\n"; // Standard format for not found
        }
    } else if (command == "DEL") {
        if (global_kv_store.remove(key)) {
            response = "1\n"; // 1 item deleted
        } else {
            response = "0\n"; // 0 items deleted (not found)
        }
    } else {
        response = "ERROR: Unknown command. Use SET, GET, or DEL.\n";
    }

    write(fd, response.c_str(), response.size());
}

// Take current socket settings
// Add NONBLOCK flag
// Save settings
bool set_non_blocking(int fd) {
    int flags = fcntl(fd, F_GETFL, 0);
    if (flags == -1) return false;
    return fcntl(fd, F_SETFL, flags | O_NONBLOCK) == 0;
}

int main() {
    std::cout << "Starting Concurrent Server with Thread Pool..." << std::endl;
    ThreadPool pool(4); // 4 worker threads

    // 1. Create listening socket
        // socket(int domain, int type, int protocol);
        // AF_INET = IPv4, SOCK_STREAM = TCP, 0 = default protocol no.
        // this will return 0, 1, 2, 3 or -1, significance:
            // 0 -> stdin
            // 1 -> stdout
            // 2 -> stderr
            // 3 -> socket
            // -1 -> socket creation fails
    int listen_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (listen_fd == -1) {
        std::perror("Socket creation failed");
        return 1;
    }

    // Allow immediate reuse of local address/port
    // you tell the kernel, "Allow me to reuse this port immediately.""
    int opt = 1; // opt = 1 means allow, opt = 0 means dont allow
    setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    // 2. Bind to port
    sockaddr_in server_addr{}; // A structure used to store an IPv4 address and port.
    server_addr.sin_family = AF_INET; // Address Family = IPv4
    server_addr.sin_addr.s_addr = INADDR_ANY; // Specifying IP Address, INADDR_ANY = 0.0.0.0 i.e. accept connections on all network interfaces
        // WiFi      → 192.168.1.10
        // Ethernet  → 10.0.0.5
        // Loopback  → 127.0.0.1
        // INADDR_ANY will listen for all at once

    server_addr.sin_port = htons(PORT); // htons means host to network short, converts 8080 into proper network byte order.

    // now, server_addr looks like: 
    //     Family = AF_INET
    //     IP     = 0.0.0.0
    //     Port   = 8080

    // socket() -> creates socket object, bind() -> assigns address to socket
    if (bind(listen_fd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
        std::perror("Bind failed");
        return 1;
    }

    // 3. Listen for incoming connections
        // SOMAXCONN -> Controls the connection queue size.
        // explanation: client1, client2, client3 sends connection request.
        //              Before server accepts, it temporarily stores these requests in queue.
        //              So, SOMAXCONN means -> use the maximum queue size allowed by the OS, instead of hard coded value
        //              Return value is 0 for success, -1 for failure
    if (listen(listen_fd, SOMAXCONN) < 0) {
        std::perror("Listen failed");
        return 1;
    }

    // Need for non blocking socket:
    //      blocking socket waits forever unless there are no clients.
    //      for non-blocking -> if no clients then return immediately
    if (!set_non_blocking(listen_fd)) {
        std::perror("Failed to set non-blocking on listen socket");
        return 1;
    }

    std::cout << "Server starting on port " << PORT << "...\n";

    
    // epoll
    // ↓
    // Tell me which sockets are ready
    // ↓
    // Process only those sockets
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        std::perror("epoll_create1 failed");
        return 1;
    }

    // Add listening socket to epoll instance
    epoll_event ev{};
    ev.events = EPOLLIN; // Watch for read events (new connections)
    ev.data.fd = listen_fd;
    // Monitor this socket and notify me when something interesting happens.
    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, listen_fd, &ev) == -1) {
        std::perror("epoll_ctl failed for listen_fd");
        return 1;
    }
    std::vector<epoll_event> events(MAX_EVENTS);

    // 5. Core event loop
    while(true) {
        // Block until events occur (timeout = -1 means infinite wait)
        // epoll_wait() puts your program to sleep until something happens.
        // 4th parameter = -1, which means wait forever
        int num_events = epoll_wait(epoll_fd, events.data(), MAX_EVENTS, -1); // Put my server to sleep until at least one event occurs.
        if (num_events == -1) {
            std::perror("epoll_wait failed");
            break;
        }

        for(int i=0; i<num_events; ++i) {
            int current_fd = events[i].data.fd;

            // Scenario A: New Client connection on listening socket
                // listen_fd = Reception Desk

                // Client arrives
                // ↓
                // accept()
                // ↓
                // client_fd = Private room for that client
            if (current_fd == listen_fd) {
                sockaddr_in client_addr{};
                socklen_t client_len = sizeof(client_addr);

                // creates a new socket dedicated to that client.
                int client_fd = accept(listen_fd, (struct sockaddr*)&client_addr, &client_len);
                
                if (client_fd == -1) {
                    if (errno != EAGAIN && errno != EWOULDBLOCK) {
                        std::perror("Accept error");
                    }
                    continue;
                }

                set_non_blocking(client_fd);

                // Start monitoring this client socket too.
                // Watch this client socket for incoming data (EPOLLIN)
                // Using Edge-Triggered (EPOLLET) mode for optimal performance (Notifies ONLY once "data has arrived")
                ev.events = EPOLLIN | EPOLLET;
                ev.data.fd = client_fd;
                if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, client_fd, &ev) == -1) {
                    std::perror("epoll_ctl ADD client failed");
                    close(client_fd);
                } else {
                    std::cout << "Connected new client (FD: " << client_fd << ")\n";
                }
            }

            // Scenario B: Existing client sent data
            else if (events[i].events & EPOLLIN) {
                char buffer[BUFFER_SIZE];
                bool close_connection = false;

                // Since we use Edge-Triggered mode, we must loop and read ALL data
                while (true) {
                    std::memset(buffer, 0, sizeof(buffer));
                    ssize_t bytes_read = read(current_fd, buffer, sizeof(buffer) - 1);

                    if (bytes_read > 0) {
                        std::string payload(buffer, bytes_read);
                        
                        // Strip trailing newlines
                        while (!payload.empty() && (payload.back() == '\n' || payload.back() == '\r')) {
                            payload.pop_back();
                        }

                        // Enqueue task to ThreadPool instead of echoing directly
                        pool.enqueue([current_fd, payload] {
                            std::string command, key, value;
                            size_t first_space = payload.find(' ');
                            
                            // Parse COMMAND KEY VALUE format
                            if (first_space != std::string::npos) {
                                command = payload.substr(0, first_space);
                                size_t second_space = payload.find(' ', first_space + 1);
                                
                                if (second_space != std::string::npos) {
                                    // Found a second space: extract key and value
                                    key = payload.substr(first_space + 1, second_space - first_space - 1);
                                    value = payload.substr(second_space + 1);
                                } else {
                                    // Only one space: e.g., "GET AAPL" or "DEL AAPL"
                                    key = payload.substr(first_space + 1);
                                }
                            } else {
                                command = payload; // Just a single word, like "GET" (which will trigger the missing key error)
                            }
                            
                            process_request(current_fd, command, key, value);
                        });
                    } else if (bytes_read == 0) {
                        // Client closed connection cleanly
                        close_connection = true;
                        break;
                    } else {
                        // If errno is EAGAIN, it means we read all available data for now
                        if (errno == EAGAIN || errno == EWOULDBLOCK) {
                            break;
                        }
                        std::perror("Read error");
                        close_connection = true;
                        break;
                    }
                }

                if (close_connection) {
                    std::cout << "Client disconnected (FD: " << current_fd << ")\n";
                    epoll_ctl(epoll_fd, EPOLL_CTL_DEL, current_fd, nullptr);
                    close(current_fd);
                }
            }
        }
    }

    close(listen_fd);
    close(epoll_fd);
    return 0;
}