main.cpp

// Author: Suchismita Bose (psychomita)
// Year: 2025
// License: MIT
// This file is part of the HTTP Load Balancer project.
// You may use, modify, and distribute this file under the terms of the MIT License.

#include "connection_handler.h"
#include <iostream>
#include <thread>
#include <chrono>
#include <csignal>
#include <atomic>
#include <memory>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <cstring>
#include <sstream>
#include <iomanip>
#include <fcntl.h>     
#include <sys/resource.h> 
#include <unistd.h>    


class HTTPLoadBalancer {
private:
    ConnectionHandler connection_handler_;
    int listen_port_;
    int listen_socket_;
    std::atomic<bool> running_;
    std::thread health_thread_;
    std::thread stats_thread_;
    
public:
    HTTPLoadBalancer(int port, int health_check_interval = 15, bool enable_logging = true)
        : connection_handler_(health_check_interval, enable_logging), 
          listen_port_(port), listen_socket_(-1), running_(false) {
    }
    
    ~HTTPLoadBalancer() {
        stop();
    }
    
    bool add_backend(const std::string& ip, int port, int weight = 1, int max_connections = 1000) {
        try {
            connection_handler_.add_backend_server(ip, port, weight, max_connections);
            return true;
        } catch (const std::exception& e) {
            std::cerr << "[ERROR] Failed to add backend " << ip << ":" << port << " - " << e.what() << std::endl;
            return false;
        }
    }
    
    bool start() {
        if (connection_handler_.get_backend_count() == 0) {
            std::cerr << "[ERROR] No backend servers configured" << std::endl;
            return false;
        }
        
        // Create socket
        listen_socket_ = socket(AF_INET, SOCK_STREAM, 0);
        if (listen_socket_ < 0) {
            perror("[ERROR] Failed to create socket");
            return false;
        }
        
        // Enable socket reuse and set socket options
        int opt = 1;
        if (setsockopt(listen_socket_, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) {
            perror("[ERROR] setsockopt SO_REUSEADDR failed");
            close(listen_socket_);
            return false;
        }
        
        // Set socket to non-blocking for better performance
        int flags = fcntl(listen_socket_, F_GETFL, 0);
        if (flags >= 0) {
            fcntl(listen_socket_, F_SETFL, flags | O_NONBLOCK);
        }
        
        // Configure server address
        struct sockaddr_in server_addr;
        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(listen_port_);
        
        // Bind socket
        if (bind(listen_socket_, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
            perror("[ERROR] Bind failed");
            close(listen_socket_);
            return false;
        }
        
        // Listen for connections
        if (listen(listen_socket_, SOMAXCONN) < 0) {
            perror("[ERROR] Listen failed");
            close(listen_socket_);
            return false;
        }
        
        running_ = true;
        
        print_startup_banner();
        
        return true;
    }
    
    void run() {
        if (!running_) {
            std::cerr << "[ERROR] Load balancer not started" << std::endl;
            return;
        }
        
        std::cout << "[INFO] Load balancer accepting connections..." << std::endl;
        
        while (running_) {
            struct sockaddr_in client_addr;
            socklen_t client_len = sizeof(client_addr);
            
            // Accept client connection
            int client_socket = accept(listen_socket_, (struct sockaddr*)&client_addr, &client_len);
            
            if (client_socket < 0) {
                if (running_ && errno != EAGAIN && errno != EWOULDBLOCK) {
                    perror("[ERROR] Accept failed");
                }
                // Small delay to prevent busy waiting
                std::this_thread::sleep_for(std::chrono::milliseconds(1));
                continue;
            }
            
            // Get client IP for logging
            char client_ip[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &client_addr.sin_addr, client_ip, INET_ADDRSTRLEN);
            
            // Handle client in a separate thread
            std::thread client_thread(&ConnectionHandler::handle_client_connection, 
                                    &connection_handler_, client_socket, std::string(client_ip));
            client_thread.detach();
        }
        
        std::cout << "[INFO] Stopped accepting new connections" << std::endl;
    }
    
    void stop() {
        if (!running_) {
            return;
        }
        
        std::cout << "[INFO] Initiating graceful shutdown..." << std::endl;
        running_ = false;
        
        // Close listening socket
        if (listen_socket_ >= 0) {
            close(listen_socket_);
            listen_socket_ = -1;
        }
        
        // Wait for monitoring threads to finish
        if (health_thread_.joinable()) {
            health_thread_.join();
        }
        if (stats_thread_.joinable()) {
            stats_thread_.join();
        }
        
        // Give active connections time to finish
        std::this_thread::sleep_for(std::chrono::seconds(2));
        
        std::cout << "[INFO] Load balancer stopped gracefully" << std::endl;
    }
    
    void start_monitoring_threads() {
        if (!running_) {
            std::cerr << "[ERROR] Cannot start monitoring threads - load balancer not running" << std::endl;
            return;
        }
        
        // Health check thread
        health_thread_ = std::thread([this]() {
            std::cout << "[INFO] Health monitoring thread started" << std::endl;
            while (running_) {
                std::this_thread::sleep_for(std::chrono::seconds(15));
                if (running_) {
                    connection_handler_.perform_health_checks();
                    
                    // Check if we have any healthy servers
                    if (!connection_handler_.has_healthy_servers()) {
                        std::cerr << "[CRITICAL] No healthy backend servers available!" << std::endl;
                    }
                }
            }
            std::cout << "[INFO] Health monitoring thread stopped" << std::endl;
        });
        
        // Statistics thread
        stats_thread_ = std::thread([this]() {
            std::cout << "[INFO] Statistics monitoring thread started" << std::endl;
            while (running_) {
                std::this_thread::sleep_for(std::chrono::seconds(10));
                if (running_) {
                    connection_handler_.print_statistics();
                }
            }
            std::cout << "[INFO] Statistics monitoring thread stopped" << std::endl;
        });
    }
    
    bool is_running() const {
        return running_;
    }
    
    const ConnectionHandler& get_connection_handler() const {
        return connection_handler_;
    }
    
    int get_listen_port() const {
        return listen_port_;
    }
    
private:
    void print_startup_banner() {
        std::cout << "\n" << std::string(60, '=') << std::endl;
        std::cout << "HTTP LOAD BALANCER - LEAST CONNECTIONS ALGORITHM" << std::endl;
        std::cout << std::string(60, '=') << std::endl;
        std::cout << "Status: STARTED" << std::endl;
        std::cout << "Listen Port: " << listen_port_ << std::endl;
        std::cout << "Backend Servers: " << connection_handler_.get_backend_count() << std::endl;
        std::cout << "Algorithm: Weighted Least Connections" << std::endl;
        std::cout << "Health Checks: Enabled (15s interval)" << std::endl;
        std::cout << "Statistics: Enabled (10s interval)" << std::endl;
        
        auto now = std::chrono::system_clock::now();
        auto time_t = std::chrono::system_clock::to_time_t(now);
        std::cout << "Started: " << std::put_time(std::localtime(&time_t), "%Y-%m-%d %H:%M:%S") << std::endl;
        std::cout << std::string(60, '=') << std::endl;
        std::cout << "Press Ctrl+C to stop gracefully" << std::endl;
        std::cout << std::string(60, '=') << "\n" << std::endl;
    }
};

// Global variable for signal handling
std::unique_ptr<HTTPLoadBalancer> g_load_balancer;

// Signal handler for graceful shutdown
void signal_handler(int signal) {
    std::cout << "\n[SIGNAL] Received signal " << signal;
    
    switch (signal) {
        case SIGINT:
            std::cout << " (SIGINT - Ctrl+C)";
            break;
        case SIGTERM:
            std::cout << " (SIGTERM - Termination request)";
            break;
        case SIGHUP:
            std::cout << " (SIGHUP - Hangup)";
            break;
        default:
            std::cout << " (Unknown signal)";
            break;
    }
    
    std::cout << ". Initiating graceful shutdown..." << std::endl;
    
    if (g_load_balancer) {
        g_load_balancer->stop();
    }
    
    std::cout << "[SIGNAL] Shutdown complete. Exiting." << std::endl;
    exit(0);
}

void print_usage(const char* program_name) {
    std::cout << std::string(70, '=') << std::endl;
    std::cout << "HTTP LOAD BALANCER WITH LEAST CONNECTIONS ALGORITHM" << std::endl;
    std::cout << std::string(70, '=') << std::endl;
    std::cout << std::endl;
    std::cout << "USAGE:" << std::endl;
    std::cout << "  " << program_name << " <listen_port> <backend1> [backend2] [backend3] ..." << std::endl;
    std::cout << std::endl;
    std::cout << "PARAMETERS:" << std::endl;
    std::cout << "  listen_port    Port to listen for incoming connections (1024-65535)" << std::endl;
    std::cout << "  backend        Backend server specification" << std::endl;
    std::cout << std::endl;
    std::cout << "BACKEND FORMAT:" << std::endl;
    std::cout << "  IP:PORT              Basic backend (weight = 1, max_conn = 1000)" << std::endl;
    std::cout << "  IP:PORT:WEIGHT       Weighted backend (max_conn = 1000)" << std::endl;
    std::cout << "  IP:PORT:WEIGHT:MAX   Full specification" << std::endl;
    std::cout << std::endl;
    std::cout << "EXAMPLES:" << std::endl;
    std::cout << "  # Basic load balancing (equal weights)" << std::endl;
    std::cout << "  " << program_name << " 8080 127.0.0.1:9001 127.0.0.1:9002 127.0.0.1:9003" << std::endl;
    std::cout << std::endl;
    std::cout << "  # Weighted load balancing (server 1 gets 2x traffic)" << std::endl;
    std::cout << "  " << program_name << " 8080 127.0.0.1:9001:2 127.0.0.1:9002:1 127.0.0.1:9003:1" << std::endl;
    std::cout << std::endl;
    std::cout << "  # Mixed local and remote backends" << std::endl;
    std::cout << "  " << program_name << " 8080 127.0.0.1:9001:2 192.168.1.100:8080:1 10.0.0.5:3000:3" << std::endl;
    std::cout << std::endl;
    std::cout << "  # Full specification with connection limits" << std::endl;
    std::cout << "  " << program_name << " 8080 127.0.0.1:9001:2:500 127.0.0.1:9002:1:200" << std::endl;
    std::cout << std::endl;
    std::cout << "FEATURES:" << std::endl;
    std::cout << "  ✓ Least connections load balancing algorithm" << std::endl;
    std::cout << "  ✓ Weighted load balancing support" << std::endl;
    std::cout << "  ✓ Automatic health monitoring and failover" << std::endl;
    std::cout << "  ✓ Real-time statistics and monitoring" << std::endl;
    std::cout << "  ✓ High-performance multi-threaded architecture" << std::endl;
    std::cout << "  ✓ Graceful shutdown (Ctrl+C)" << std::endl;
    std::cout << "  ✓ Configurable connection limits per backend" << std::endl;
    std::cout << std::endl;
    std::cout << "TESTING:" << std::endl;
    std::cout << "  # Start test backend servers" << std::endl;
    std::cout << "  python3 backend_server.py 9001 &" << std::endl;
    std::cout << "  python3 backend_server.py 9002 &" << std::endl;
    std::cout << "  python3 backend_server.py 9003 &" << std::endl;
    std::cout << std::endl;
    std::cout << "  # Test the load balancer" << std::endl;
    std::cout << "  curl http://localhost:8080" << std::endl;
    std::cout << "  ./test_script.sh" << std::endl;
    std::cout << std::endl;
    std::cout << std::string(70, '=') << std::endl;
}

bool parse_backend_string(const std::string& backend_str, std::string& ip, int& port, int& weight, int& max_connections) {
    std::istringstream iss(backend_str);
    std::string token;
    std::vector<std::string> parts;
    
    // Split by colon
    while (std::getline(iss, token, ':')) {
        parts.push_back(token);
    }
    
    if (parts.size() < 2 || parts.size() > 4) {
        return false;
    }
    
    try {
        // Parse IP
        ip = parts[0];
        if (ip.empty()) {
            return false;
        }
        
        // Parse port
        port = std::stoi(parts[1]);
        if (port <= 0 || port > 65535) {
            return false;
        }
        
        // Parse weight (optional, default = 1)
        if (parts.size() >= 3) {
            weight = std::stoi(parts[2]);
            if (weight <= 0) {
                return false;
            }
        } else {
            weight = 1;
        }
        
        // Parse max_connections (optional, default = 1000)
        if (parts.size() >= 4) {
            max_connections = std::stoi(parts[3]);
            if (max_connections <= 0) {
                return false;
            }
        } else {
            max_connections = 1000;
        }
        
        return true;
        
    } catch (const std::exception& e) {
        return false;
    }
}

void print_configuration_summary(const HTTPLoadBalancer& lb) {
    std::cout << "\n" << std::string(50, '-') << std::endl;
    std::cout << "CONFIGURATION SUMMARY" << std::endl;
    std::cout << std::string(50, '-') << std::endl;
    std::cout << "Listen Port: " << lb.get_listen_port() << std::endl;
    std::cout << "Backend Count: " << lb.get_connection_handler().get_backend_count() << std::endl;
    std::cout << "Load Balancing: Weighted Least Connections" << std::endl;
    std::cout << "Health Checks: Every 15 seconds" << std::endl;
    std::cout << "Statistics: Every 10 seconds" << std::endl;
    std::cout << std::string(50, '-') << std::endl;
}

void validate_environment() {
    // Check if we can create sockets
    int test_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (test_socket < 0) {
        std::cerr << "[ERROR] Cannot create sockets. Check system permissions." << std::endl;
        exit(1);
    }
    close(test_socket);
    
    // Check file descriptor limits
    struct rlimit rl;
    if (getrlimit(RLIMIT_NOFILE, &rl) == 0) {
        if (rl.rlim_cur < 1024) {
            std::cout << "[WARNING] Low file descriptor limit (" << rl.rlim_cur 
                      << "). Consider increasing with 'ulimit -n 4096'" << std::endl;
        }
    }
}

int main(int argc, char* argv[]) {
    // Print header
    std::cout << "HTTP Load Balancer v1.0 - Starting..." << std::endl;
    
    // Validate command line arguments
    if (argc < 3) {
        print_usage(argv[0]);
        return 1;
    }
    
    // Validate environment
    validate_environment();
    
    // Setup signal handlers for graceful shutdown
    signal(SIGINT, signal_handler);   // Ctrl+C
    signal(SIGTERM, signal_handler);  // Termination request
    signal(SIGHUP, signal_handler);   // Hangup
    
    // Ignore SIGPIPE (broken pipe) to handle client disconnections gracefully
    signal(SIGPIPE, SIG_IGN);
    
    // Parse listen port
    int listen_port = 0;
    try {
        listen_port = std::stoi(argv[1]);
        if (listen_port <= 1023 || listen_port > 65535) {
            std::cerr << "[ERROR] Invalid listen port: " << listen_port 
                      << " (must be between 1024-65535)" << std::endl;
            return 1;
        }
    } catch (const std::exception& e) {
        std::cerr << "[ERROR] Invalid port number: " << argv[1] << std::endl;
        return 1;
    }
    
    // Create load balancer instance
    g_load_balancer = std::make_unique<HTTPLoadBalancer>(listen_port);
    
    // Parse and add backend servers
    int backends_added = 0;
    std::cout << "\n[CONFIG] Parsing backend servers..." << std::endl;
    
    for (int i = 2; i < argc; i++) {
        std::string backend_str = argv[i];
        std::string ip;
        int port, weight, max_connections;
        
        if (!parse_backend_string(backend_str, ip, port, weight, max_connections)) {
            std::cerr << "[ERROR] Invalid backend format: " << backend_str 
                     << " (expected IP:PORT[:WEIGHT[:MAX_CONN]])" << std::endl;
            continue;
        }
        
        if (g_load_balancer->add_backend(ip, port, weight, max_connections)) {
            backends_added++;
            std::cout << "[CONFIG] Added backend #" << backends_added << ": " 
                      << ip << ":" << port 
                      << " (weight=" << weight << ", max_conn=" << max_connections << ")" << std::endl;
        } else {
            std::cerr << "[ERROR] Failed to add backend: " << backend_str << std::endl;
        }
    }
    
    if (backends_added == 0) {
        std::cerr << "[ERROR] No valid backend servers configured" << std::endl;
        return 1;
    }
    
    // Print configuration summary
    print_configuration_summary(*g_load_balancer);
    
    // Start the load balancer
    std::cout << "\n[INFO] Starting load balancer..." << std::endl;
    if (!g_load_balancer->start()) {
        std::cerr << "[ERROR] Failed to start load balancer" << std::endl;
        return 1;
    }
    
    // Start monitoring threads
    std::cout << "[INFO] Starting monitoring threads..." << std::endl;
    g_load_balancer->start_monitoring_threads();
    
    // Initial health check
    std::cout << "[INFO] Performing initial health checks..." << std::endl;
    g_load_balancer->get_connection_handler().perform_health_checks();
    
    // Check if we have any healthy servers
    if (!g_load_balancer->get_connection_handler().has_healthy_servers()) {
        std::cerr << "[WARNING] No healthy backend servers found during initial check!" << std::endl;
        std::cerr << "[WARNING] Make sure your backend servers are running and accessible." << std::endl;
    }
    
    // Run the main server loop
    try {
        g_load_balancer->run();
    } catch (const std::exception& e) {
        std::cerr << "[ERROR] Exception in main loop: " << e.what() << std::endl;
        return 1;
    }
    
    std::cout << "[INFO] Load balancer exited normally" << std::endl;
    return 0;
}