MontaukOS/kernel/src/Sched/Scheduler.hpp

/*
    * Scheduler.hpp
    * Preemptive process scheduler with SMP support
    * Copyright (c) 2025-2026 Daniel Hammer
*/

#pragma once
#include <cstdint>
#include <Api/Syscall.hpp>

namespace Sched {

    // Process and thread slots share the same table: the main thread of a
    // process owns the per-process resources (PML4, IPC handles, args/cwd/user,
    // redirected I/O), while sibling threads occupy additional slots whose
    // `primarySlot` points back to the main thread. The scheduler treats
    // every slot uniformly -- the primary/sibling distinction matters only
    // for resource ownership on exit.
    static constexpr int MaxProcesses = 256;
    static constexpr uint64_t StackPages = 4;  // 16 KiB kernel stack per thread
    static constexpr uint64_t StackSize = StackPages * 0x1000;
    static constexpr uint64_t UserStackPages = 8;  // 32 KiB user stack (main thread)
    static constexpr uint64_t UserStackSize = UserStackPages * 0x1000;
    static constexpr uint64_t UserStackTop = 0x7FFFFFF000ULL;  // Main-thread user stack top VA
    static constexpr uint64_t UserHeapBase = 0x40000000ULL;    // User heap start VA
    static constexpr uint32_t UserReadDirSlots = 64;           // rotating scratch pages for SYS_READDIR
    static constexpr uint64_t UserReadDirBase =
        UserHeapBase - (uint64_t)UserReadDirSlots * 0x1000ULL;
    static constexpr uint64_t ExitStubAddr = 0x3FF000ULL;     // User-space exit stub page
    static constexpr uint64_t TimeSliceMs = 10; // 10 ms time slice

    enum class ProcessState {
        Free,
        Ready,
        Running,
        Blocked,       // Waiting (e.g. waitpid) -- not schedulable
        Terminated
    };

    struct Process {
        int pid;
        ProcessState state;
        int waitingForPid;     // PID this process is blocked on (-1 if none)
        uint64_t sleepUntilTick; // Tick deadline for sleep/object wait timeout (0 = none)
        void* waitingOnObject; // IPC/scheduler object this process is blocked on (nullptr if none)
        char name[64];
        uint64_t savedRsp;
        uint64_t stackBase;       // Bottom of allocated kernel stack (lowest address)
        uint64_t entryPoint;
        uint64_t sliceRemaining;  // Milliseconds left in current time slice
        uint64_t cpuTimeMs;       // Accumulated scheduler runtime in milliseconds
        uint64_t pml4Phys;        // Physical address of per-process PML4
        uint64_t kernelStackTop;  // Top of kernel stack (for TSS RSP0 / SYSCALL)
        uint64_t userStackTop;    // User-space stack top
        uint64_t heapNext;        // Simple bump allocator for user heap
        uint32_t readdirCursor;   // Next SYS_READDIR scratch slot
        char args[256];           // Command-line arguments (set by parent via Spawn)
        char user[32];            // Owner user name (inherited from parent on spawn)
        char cwd[256];            // Absolute current working directory

        int runningOnCpu;         // CPU index running this process (-1 if not running)
        bool killPending = false; // Set by Sys_Kill when target is running on another CPU
        bool reapReady = false;   // Set once teardown is complete and BSP may free slot resources
        bool startPending = false; // Spawned but not yet made runnable

        // Threading. Every slot is a schedulable thread. The "main thread" of a
        // process has primarySlot == own slot and owns per-process resources
        // (pml4Phys, IPC handles, libs, heap, redirected I/O, cwd/user/args).
        // Sibling threads share those resources via primarySlot.
        int primarySlot = -1;       // Index into processTable; -1 if Free
        int joinerSlot = -1;        // Slot blocked in JoinThread on this slot
        int exitCode = 0;           // Returned to joiner / accumulated for the process
        bool joinable = true;       // Sibling thread keeps Terminated state for the joiner
        uint64_t threadArg = 0;     // Argument passed into entryPoint on first dispatch

        // I/O redirection for GUI terminal
        bool redirected = false;
        int parentPid = -1;
        uint8_t* outBuf = nullptr;   // 4KB ring: child writes (print/putchar), parent reads
        volatile uint32_t outHead = 0;
        volatile uint32_t outTail = 0;
        uint8_t* inBuf = nullptr;    // 4KB ring: parent writes, child reads (getchar)
        volatile uint32_t inHead = 0;
        volatile uint32_t inTail = 0;
        montauk::abi::KeyEvent keyBuf[64]; // parent injects, child reads (getkey/iskeyavailable)
        volatile uint32_t keyHead = 0;
        volatile uint32_t keyTail = 0;
        static constexpr uint32_t IoBufSize = 4096;

        // GUI terminal dimensions (set by desktop, read by SYS_TERMSIZE)
        int termCols = 0;
        int termRows = 0;

        // IPC-backed redirected terminal channels
        int ioOutHandle = -1;
        int ioInHandle = -1;
        int ioKeyHandle = -1;
        int ioWaitsetHandle = -1;

        // FPU/SSE state (FXSAVE format, must be 16-byte aligned)
        uint8_t fpuState[512] __attribute__((aligned(16)));
    };

    void Initialize();
    int Spawn(const char* vfsPath, const char* args = nullptr, bool startReady = true);
    int StartProcess(int pid);
    void Schedule();

    // True when there is runnable work somewhere in the process table.
    bool HasReadyProcesses();

    // Called from the APIC timer handler with the elapsed time for that CPU's
    // tick interval. The BSP runs at 1 ms; APs may use a coarser interval.
    void Tick(uint32_t elapsedMs = 1);

    // Get the PID of the currently running process (-1 if idle).
    // For sibling threads this returns the primary slot's PID -- i.e. the
    // process-level identifier, which is what userspace sees.
    int GetCurrentPid();

    // Get the per-thread id of the currently running thread (-1 if idle).
    // For the main thread this equals GetCurrentPid().
    int GetCurrentTid();

    // Get a pointer to the currently running process (primary slot, nullptr if idle).
    // Always returns the slot that owns per-process state -- never a sibling thread.
    Process* GetCurrentProcessPtr();

    // Get a pointer to the currently running thread's slot (may be a sibling).
    Process* GetCurrentThreadPtr();

    // Called by terminated processes to mark themselves done.
    // Tears down the entire process (kills all sibling threads, frees address space).
    void ExitProcess();

    // ====================================================================
    // Threading
    // ====================================================================

    // Spawn a sibling thread inside the current process.
    // `entry` is a user-mode RIP, `arg` is delivered to entry in %rdi, and
    // `userStackTop` points to the top-of-stack the caller (libc) allocated
    // out of the user heap. Returns the new TID, or -1 on failure.
    int SpawnThread(uint64_t entry, uint64_t arg, uint64_t userStackTop);

    // Terminate the currently executing thread. If this is the main thread,
    // the entire process exits (equivalent to ExitProcess).
    [[noreturn]] void ExitCurrentThread(int exitCode);

    // Block the current thread until `tid` (a sibling of the same process)
    // terminates. On success, writes the exit code via outExitCode (may be
    // null) and frees the joined thread's slot. Returns 0 on success, -1 if
    // tid is invalid or not joinable by the caller.
    int JoinThread(int tid, int* outExitCode);

    // Check if a process is still alive (Ready, Running, or Blocked)
    bool IsAlive(int pid);

    // Block the current process until the given PID exits.
    void BlockOnPid(int pid);

    // Block the current process for the given number of milliseconds.
    void BlockForSleep(uint64_t ms);

    // Block the current process until the given object is signaled or timed out.
    // timeoutMs == 0 means wait indefinitely.
    void BlockOnObject(void* object, uint64_t timeoutMs = 0);

    // Atomically check a condition under the scheduler lock and block only if
    // it still holds. This prevents lost wakeups for edge-triggered waiters.
    bool BlockOnObjectIf(void* object, uint64_t timeoutMs,
                         bool (*shouldBlock)(void*), void* context);

    // BSP-only scheduler housekeeping: wake expired sleepers and reclaim
    // terminated process resources.
    void RunBspMaintenance();

    // Return the earliest blocked sleep/object timeout deadline in ticks,
    // or 0 when no timed waits are pending.
    uint64_t GetNextDeadlineTick();

    // Wake any processes blocked on the given object.
    void WakeObjectWaiters(void* object);

    // Kill a process by PID. If the process is running on another CPU,
    // sets a kill-pending flag checked on the next timer tick.
    // Returns 0 on success, -1 on failure.
    int KillProcess(int pid);

    // Find a process by PID (returns nullptr if not found or not alive)
    Process* GetProcessByPid(int pid);

    // Spawn a crashpad process with the given crash file path as argument.
    // Called by the exception handler when a process faults.
    int SpawnCrashPad();

    // Get a pointer to slot i in the process table (for enumeration)
    Process* GetProcessSlot(int slot);

    // Per-process allocated page count (tracked by Heap syscalls, separate from Process struct)
    inline uint64_t g_allocatedPages[MaxProcesses] = {};

}