WorkerPool.hpp

/**
 * @file WorkerPool.hpp
 * @brief Process pool. Manages pre-warmed Python worker processes via a
 *        CoW-aware template process. Workers are forked from a pre-loaded
 *        template (Copy-on-Write), reducing per-worker cold start from
 *        ~800ms (fresh import) to ~20ms (fork overhead only).
 *        Includes a scavenger thread for scale-down: workers idle longer
 *        than idle_timeout_s are killed with SIGKILL.
 *
 * CSCI 599: Network Systems for Cloud Computing
 * University of Southern California
 */

#ifndef __WORKER_POOL_HPP__
#define __WORKER_POOL_HPP__

#include <iostream>
#include <string>
#include <queue>
#include <unordered_map>
#include <vector>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <thread>
#include <mutex>
#include <atomic>
#include <chrono>
#include <cerrno>
#include <cstring>
#include "log.hpp"

namespace fengcheng {

// Path of the Unix socket the template process listens on for fork requests.
static constexpr const char* TEMPLATE_CTRL_SOCK = "/tmp/faas_template_ctrl.sock";

class WorkerPool {
    using Clock     = std::chrono::steady_clock;
    using TimePoint = Clock::time_point;

    struct IdleEntry {
        int       id;
        TimePoint became_idle_at;
    };

private:
    std::string                       _sock_path_prefix;
    std::unordered_map<int, pid_t>    _workers;        // id -> pid (alive: warm + busy)
    std::queue<IdleEntry>             _idle_queue;      // FIFO with timestamps
    std::mutex                        _mtx;

    int                               _next_id;         // monotonic, never reused
    int                               _alive_count;     // alive workers (warm + busy)
    int                               _idle_timeout_s;  // 0 = never kill

    // CoW template process (skipped in --no-cow ablation mode)
    bool                              _use_cow_template;
    pid_t                             _template_pid;

    std::thread                       _scavenger;
    std::atomic<bool>                 _stop_scavenger;

public:
    WorkerPool(const std::string& prefix = "/tmp/faas_worker_",
               int idle_timeout_s = 0,
               bool use_cow_template = true)
        : _sock_path_prefix(prefix),
          _next_id(0), _alive_count(0),
          _idle_timeout_s(idle_timeout_s),
          _use_cow_template(use_cow_template),
          _template_pid(-1),
          _stop_scavenger(false) {

        if (_use_cow_template) {
            _start_template_process();
        } else {
            // --no-cow ablation: every worker pays the full ~900 ms cost of
            // booting Python and re-importing Pillow on each spin-up.
            // Used to isolate CoW's contribution from the predictor's.
            logMessage(NORMAL, "[WorkerPool] CoW template DISABLED (--no-cow). "
                               "Workers will exec fresh Python on each spawn (~900ms each).");
        }

        if (_idle_timeout_s > 0) {
            _scavenger = std::thread(&WorkerPool::_scavenge, this);
            logMessage(NORMAL, "[WorkerPool] Scale-down enabled: idle_timeout=%ds", _idle_timeout_s);
        }
    }

    ~WorkerPool() {
        _stop_scavenger.store(true);
        if (_scavenger.joinable()) _scavenger.join();

        // Worker reaping differs by mode:
        //   CoW:    workers are children of the template process, which has
        //           SIGCHLD=SIG_IGN, so it auto-reaps zombies. We send SIGKILL
        //           but never waitpid in this loop.
        //   no-CoW: workers are direct children of the C++ gateway. We must
        //           waitpid (non-blocking) to prevent zombies.
        {
            std::lock_guard<std::mutex> lock(_mtx);
            for (auto& [id, pid] : _workers) {
                kill(pid, SIGKILL);
                if (!_use_cow_template) {
                    waitpid(pid, nullptr, WNOHANG);
                }
                unlink(GetSockPath(id).c_str());
            }
        }

        // Kill the template process (CoW mode only — direct child of C++).
        if (_use_cow_template && _template_pid > 0) {
            kill(_template_pid, SIGKILL);
            waitpid(_template_pid, nullptr, 0);
            unlink(TEMPLATE_CTRL_SOCK);
            logMessage(NORMAL, "[WorkerPool] Template process (pid=%d) terminated.", _template_pid);
        }
    }

    // Two-phase fork: collect IDs under lock, request workers from template outside lock.
    // Safety cap: healthy runs peak at ~23 workers (Spike=30 RPS × predicted_lambda_bonus
    // 1.5 × T=0.5s + margin = 23). A cap of 64 gives ~2.8x headroom for transient
    // bursts but blocks fork bombs when T gets polluted (observed once: T drifted to
    // 91s -> Target=443 -> OOM -> hard shutdown). The cap is a safety rail, not a
    // functional limit; if we ever see an ERROR log for it, investigate T pollution.
    static constexpr int kMaxWorkers = 64;

    void ScaleTo(int target) {
        if (target > kMaxWorkers) {
            logMessage(ERROR, "[WorkerPool] Target=%d exceeds safety cap %d — clamping. "
                              "Check for T/lambda pollution in Predictor logs.",
                       target, kMaxWorkers);
            target = kMaxWorkers;
        }
        std::vector<std::pair<int, std::string>> to_fork;
        {
            std::lock_guard<std::mutex> lock(_mtx);
            if (target > _alive_count) {
                int need = target - _alive_count;
                logMessage(NORMAL, "[WorkerPool] Scale up: +%d workers (alive %d -> %d)",
                           need, _alive_count, target);
                for (int i = 0; i < need; ++i) {
                    ++_next_id;
                    ++_alive_count;
                    to_fork.push_back({ _next_id, GetSockPath(_next_id) });
                }
            }
        }

        // Warm-up timer differs by mode:
        //   CoW:    ~100 ms (fork from template + socket bind)
        //   no-CoW: ~900 ms (fork + execlp + Python boot + Pillow import)
        // The B4 ablation hinges on this difference: without CoW the worker
        // is not ready in time for the spike, even when the predictor fires
        // correctly during the ramp.
        const useconds_t warmup_us = _use_cow_template ? 100000 : 900000;

        for (auto& [id, sock_path] : to_fork) {
            pid_t pid;
            if (_use_cow_template) {
                // Ask the template process to fork via CoW — worker inherits
                // pre-loaded Pillow pages.
                pid = _request_worker_from_template(sock_path);
            } else {
                // --no-cow: fork+execlp a fresh Python worker.
                pid = _spawn_worker_naive(sock_path);
            }

            if (pid <= 0) {
                logMessage(ERROR, "[WorkerPool] Worker %d spawn failed.", id);
                std::lock_guard<std::mutex> lock(_mtx);
                --_alive_count;
                continue;
            }

            {
                std::lock_guard<std::mutex> lock(_mtx);
                _workers[id] = pid;
            }

            const char* tag = _use_cow_template ? "CoW fork" : "naive exec";
            std::thread([this, id, warmup_us, tag]() {
                usleep(warmup_us);
                std::lock_guard<std::mutex> lock(_mtx);
                if (_workers.count(id)) {
                    _idle_queue.push({ id, Clock::now() });
                    logMessage(NORMAL, "[WorkerPool] Worker %d ready (%s). Idle=%zu Alive=%d",
                               id, tag, _idle_queue.size(), _alive_count);
                }
            }).detach();
        }
    }

    // Returns a warm worker id (removes from idle queue), or -1 if empty.
    // Caller MUST call ReturnWorkerId() when done, or the slot leaks.
    int GetIdleWorkerId() {
        std::lock_guard<std::mutex> lock(_mtx);
        if (_idle_queue.empty()) return -1;
        int id = _idle_queue.front().id;
        _idle_queue.pop();
        return id;
    }

    // Return a worker to the idle pool after finishing a request.
    void ReturnWorkerId(int id) {
        std::lock_guard<std::mutex> lock(_mtx);
        if (_workers.count(id)) {
            _idle_queue.push({ id, Clock::now() });
        }
    }

    std::string GetSockPath(int id) const {
        return _sock_path_prefix + std::to_string(id) + ".sock";
    }

    int IdleCount() {
        std::lock_guard<std::mutex> lock(_mtx);
        return (int)_idle_queue.size();
    }

    int AliveCount() {
        std::lock_guard<std::mutex> lock(_mtx);
        return _alive_count;
    }

private:
    // ---------------------------------------------------------------------------
    // Template process management
    // ---------------------------------------------------------------------------

    // Fork and exec worker_template.py. Block until the template signals READY
    // (meaning Pillow is loaded and the control socket is up).
    void _start_template_process() {
        int pipefd[2];
        if (pipe(pipefd) < 0) {
            logMessage(FATAL, "[WorkerPool] pipe() failed: %s", strerror(errno));
            return;
        }

        pid_t pid = fork();
        if (pid == 0) {
            // ---- CHILD: become the template process ----
            close(pipefd[0]);  // child does not read from the pipe

            // Close all fds except stderr (2) and the ready pipe write end.
            // This prevents the template from inheriting the epoll fd, listen socket, etc.
            for (int fd = 3; fd < 1024; fd++) {
                if (fd != pipefd[1]) close(fd);
            }

            char fd_str[16];
            snprintf(fd_str, sizeof(fd_str), "%d", pipefd[1]);
            execlp("python3", "python3", "worker_template.py", fd_str, nullptr);
            // exec failed
            _exit(1);
        }

        // ---- PARENT: wait for READY signal ----
        close(pipefd[1]);  // parent does not write to the pipe

        if (pid < 0) {
            logMessage(FATAL, "[WorkerPool] fork() for template process failed: %s", strerror(errno));
            close(pipefd[0]);
            return;
        }

        logMessage(NORMAL, "[WorkerPool] Template process forked (pid=%d). Waiting for READY...", pid);

        // Blocking read — template writes "READY\n" after Pillow is loaded (~800ms).
        char buf[32] = {};
        ssize_t n = read(pipefd[0], buf, sizeof(buf) - 1);
        close(pipefd[0]);

        if (n <= 0 || strstr(buf, "READY") == nullptr) {
            logMessage(FATAL, "[WorkerPool] Template process did not signal READY (got: '%s').", buf);
            kill(pid, SIGKILL);
            waitpid(pid, nullptr, 0);
            return;
        }

        _template_pid = pid;
        logMessage(NORMAL, "[WorkerPool] Template process READY (pid=%d). "
                           "Workers will now fork via CoW (~100ms warm-up vs ~900ms before).", pid);
    }

    // Connect to the template's control socket, send sock_path, receive worker PID.
    // Returns the new worker's PID, or -1 on error.
    pid_t _request_worker_from_template(const std::string& sock_path) {
        int ctrl_fd = socket(AF_UNIX, SOCK_STREAM, 0);
        if (ctrl_fd < 0) {
            logMessage(ERROR, "[WorkerPool] socket() for template ctrl failed: %s", strerror(errno));
            return -1;
        }

        struct sockaddr_un addr = {};
        addr.sun_family = AF_UNIX;
        strncpy(addr.sun_path, TEMPLATE_CTRL_SOCK, sizeof(addr.sun_path) - 1);

        if (connect(ctrl_fd, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
            logMessage(ERROR, "[WorkerPool] connect() to template ctrl failed: %s", strerror(errno));
            close(ctrl_fd);
            return -1;
        }

        // Send: "<sock_path>\n"
        std::string msg = sock_path + "\n";
        if (send(ctrl_fd, msg.c_str(), msg.size(), MSG_NOSIGNAL) < 0) {
            logMessage(ERROR, "[WorkerPool] send() to template ctrl failed: %s", strerror(errno));
            close(ctrl_fd);
            return -1;
        }

        // Receive: "<pid>\n"
        char buf[32] = {};
        ssize_t n = recv(ctrl_fd, buf, sizeof(buf) - 1, 0);
        close(ctrl_fd);

        if (n <= 0) {
            logMessage(ERROR, "[WorkerPool] recv() from template ctrl failed: %s", strerror(errno));
            return -1;
        }

        pid_t worker_pid = (pid_t)atoi(buf);
        if (worker_pid <= 0) {
            logMessage(ERROR, "[WorkerPool] Template returned invalid PID: '%s'", buf);
            return -1;
        }

        logMessage(NORMAL, "[WorkerPool] Template forked worker pid=%d -> %s", worker_pid, sock_path.c_str());
        return worker_pid;
    }

    // ---------------------------------------------------------------------------
    // --no-cow ablation path: fork+execlp a fresh Python worker.
    // Each spawn pays the full Python boot + Pillow import cost (~800ms),
    // matching pre-CoW behavior. Used for the B4 ablation experiment to
    // isolate CoW's contribution from the predictor's.
    // ---------------------------------------------------------------------------
    pid_t _spawn_worker_naive(const std::string& sock_path) {
        pid_t pid = fork();
        if (pid == 0) {
            // ---- CHILD: become a fresh Python worker ----
            // Close all fds except stdio so we don't inherit epoll/listen sockets.
            for (int fd = 3; fd < 1024; fd++) close(fd);
            execlp("python3", "python3", "worker.py", sock_path.c_str(), nullptr);
            // exec failed
            _exit(1);
        }
        if (pid < 0) {
            logMessage(ERROR, "[WorkerPool/naive] fork() failed: %s", strerror(errno));
            return -1;
        }
        logMessage(NORMAL, "[WorkerPool/naive] Spawned worker pid=%d -> %s", pid, sock_path.c_str());
        return pid;
    }

    // ---------------------------------------------------------------------------
    // Scavenger thread: kill workers idle longer than _idle_timeout_s.
    //
    // Reaping path differs by mode:
    //   CoW:    worker is a child of the template (SIGCHLD=SIG_IGN there),
    //           so SIGKILL alone reaps it.
    //   no-CoW: worker is a direct child of the C++ gateway. We must call
    //           waitpid(WNOHANG) here to prevent zombies.
    // ---------------------------------------------------------------------------
    void _scavenge() {
        while (!_stop_scavenger.load()) {
            std::this_thread::sleep_for(std::chrono::seconds(1));

            auto now = Clock::now();
            std::lock_guard<std::mutex> lock(_mtx);

            while (!_idle_queue.empty()) {
                auto& front = _idle_queue.front();
                double idle_s = std::chrono::duration<double>(now - front.became_idle_at).count();
                if (idle_s < _idle_timeout_s) break;

                int id = front.id;
                _idle_queue.pop();

                auto it = _workers.find(id);
                if (it != _workers.end()) {
                    kill(it->second, SIGKILL);
                    if (!_use_cow_template) {
                        // no-CoW: reap immediately to prevent zombies.
                        waitpid(it->second, nullptr, WNOHANG);
                    }
                    _workers.erase(it);
                }
                unlink(GetSockPath(id).c_str());
                --_alive_count;

                logMessage(NORMAL,
                    "[WorkerPool] Worker %d expired (idle %.1fs > %ds), killed. Alive=%d",
                    id, idle_s, _idle_timeout_s, _alive_count);
            }
        }
    }
};

} // namespace fengcheng

#endif // __WORKER_POOL_HPP__