Topic 05: Preemptive Multitasking, Context Switching & Custom Blocking Mutex

Focus: Process Control Block (PCB), 8 KiB isolated KernelStack frames, SysV64 naked assembly context switching, round-robin scheduler queues, voluntary preemption, and custom descheduling sync::Mutex<T> with TaskId wait-lists.

1. High-Level Summary Table

2. Spoken Interview Answer (Word-for-Word Script)

"Preemptive multitasking is the beating heart of an operating system, allowing multiple threads to share a single CPU core without cooperating voluntarily.

In RustOS, each task is represented by a Process Control Block (PCB) consisting of a unique TaskId, an execution state (Ready, Running, Blocked, or Terminated), a saved register struct, and an isolated 8 KiB private kernel stack. When a task is spawned, I manually synthesize a fake x86 hardware interrupt stack frame at the top of its stack: setting the Stack Segment to 0, RSP to the stack top, RFLAGS to 0x202 (interrupts enabled), CS to the kernel code selector, RIP to the task entrypoint, and pushing 15 zeros representing all general-purpose registers.

Preemption is driven by the 8253 timer interrupt on IRQ 0 (Vector 32). In interrupts.rs, I wrote a naked assembly trampoline using #[unsafe(naked)]. When the timer ticks, the CPU hardware pushes the interrupted task's flags and instruction pointer. Our naked assembly then pushes all 15 general-purpose registers (rax through r15), moves the current stack pointer rsp into rdi as the first argument, and calls our Rust scheduler function timer_schedule.

The scheduler pushes the interrupted task into the back of ready_queue, pops the next ready task, and returns the new task's stack pointer in rax. The assembly routine simply moves rax into rsp—switching stacks in a single instruction—pops the 15 general-purpose registers of the new task, sends an End of Interrupt signal (0x20) to the PIC, and issues iretq. The CPU restores flags and jumps into the new task's code as if nothing happened.

Furthermore, on top of this scheduler, I built a custom blocking Mutex. Unlike a spinlock that spins in a tight loop and starves other threads on a single core, our mutex checks the lock with compare_exchange. If contended, it captures the current TaskId, enqueues it into an internal wait_list, marks the task as Blocked, and triggers software interrupt 32 to yield immediately. When the lock guard is dropped, it pops the next waiting task and moves it back to the ready queue."

3. Deep-Dive Mechanics & Architecture Walkthrough

3.1 Process Control Block & Synthetic Stack Frame

In src/process.rs, ProcessControlBlock::new constructs the initial stack:

// src/process.rs
let stack = KernelStack::new(8192); // 8 KB private stack
let mut stack_ptr = stack.end.as_mut_ptr::();

let cs = CS::get_reg().0 as u64;

// 1. Hardware Interrupt Stack Frame (pushed by CPU on real interrupt)
stack_ptr = stack_ptr.offset(-1); stack_ptr.write(0);                    // SS
stack_ptr = stack_ptr.offset(-1); stack_ptr.write(stack.end.as_u64());   // RSP
stack_ptr = stack_ptr.offset(-1); stack_ptr.write(0x202);                // RFLAGS (IF=1)
stack_ptr = stack_ptr.offset(-1); stack_ptr.write(cs);                   // CS
stack_ptr = stack_ptr.offset(-1); stack_ptr.write(entry_point as u64);   // RIP

// 2. 15 General Purpose Registers (pushed by naked ISR)
for _ in 0..15 {
    stack_ptr = stack_ptr.offset(-1);
    stack_ptr.write(0); // r15 down to rax
}
registers.rsp = stack_ptr as u64;

By mimicking the exact memory layout of an interrupted thread, the context restoration code requires zero special-casing for brand-new threads.

3.2 The Naked Assembly Context Switcher

The timer interrupt handler in src/interrupts.rs is declared naked to prevent the Rust compiler from generating function preludes that corrupt stack offsets:

// src/interrupts.rs
#[unsafe(naked)]
extern "x86-interrupt"
fn timer_interrupt_handler(_stack_frame: InterruptStackFrame) {
    naked_asm!(
        // 1. Push all general-purpose registers
        "push rax", "push rbx", "push rcx", "push rdx",
        "push rsi", "push rdi", "push rbp",
        "push r8",  "push r9",  "push r10", "push r11",
        "push r12", "push r13", "push r14", "push r15",

        // 2. Call Rust scheduler: passes current RSP in RDI (SysV64 ABI)
        "mov rdi, rsp",
        "call timer_schedule",

        // 3. Switch stack: RAX contains next task's RSP
        "mov rsp, rax",

        // 4. Pop next task's general-purpose registers
        "pop r15", "pop r14", "pop r13", "pop r12",
        "pop r11", "pop r10", "pop r9",  "pop r8",
        "pop rbp", "pop rdi", "pop rsi", "pop rdx",
        "pop rcx", "pop rbx", "pop rax",

        // 5. Send End of Interrupt (EOI) to PIC
        "mov al, 0x20",
        "out 0x20, al",

        // 6. Return from interrupt and resume task
        "iretq",
    );
}

3.3 The Round-Robin Scheduler (timer_schedule)

In src/process.rs, timer_schedule(current_rsp) manages lifecycle queues:

// src/process.rs
#[no_mangle]
pub unsafe extern "sysv64" fn timer_schedule(current_rsp: u64) -> u64 {
    let mut scheduler = SCHEDULER.lock();

    // 1. Save current task RSP and place into appropriate queue
    if let Some(mut current) = scheduler.current_task.take() {
        current.registers.rsp = current_rsp;
        match current.state {
            TaskState::Running | TaskState::Ready => {
                current.state = TaskState::Ready;
                scheduler.ready_queue.push_back(current);
            }
            TaskState::Blocked => scheduler.blocked_queue.push_back(current),
            TaskState::Terminated => scheduler.terminated_queue.push_back(current),
        }
    }

    // 2. Select next ready task
    scheduler.schedule();

    // 3. Return stack pointer of next task
    if let Some(ref next) = scheduler.current_task {
        next.registers.rsp
    } else {
        current_rsp
    }
}

3.4 Custom Blocking Mutex (Sleep-Lock)

Unlike spinlocks, src/sync.rs implements a scheduler-aware sleep-lock:

// src/sync.rs
pub struct Mutex {
    locked: AtomicBool,
    wait_list: spin::Mutex>,
    data: UnsafeCell,
}

impl Mutex {
    pub fn lock(&self) -> MutexGuard<'_, T> {
        loop {
            if self.locked.compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed).is_ok() {
                break;
            }
            // Lock contended: enqueue current task and sleep
            let current_id = {
                let sched = SCHEDULER.lock();
                sched.current_task().unwrap().id
            };
            self.wait_list.lock().push_back(current_id);
            sleep(); // Sets state to Blocked and triggers software INT 32
        }
        MutexGuard { lock: self }
    }
}

impl<'a, T> Drop for MutexGuard<'a, T> {
    fn drop(&mut self) {
        self.lock.locked.store(false, Ordering::Release);
        if let Some(task_id) = self.lock.wait_list.lock().pop_front() {
            wake(task_id); // Moves task from blocked_queue to ready_queue
        }
    }
}

4. Real Edge Cases, Pitfalls & Debugging

4.1 The 16-Byte Stack Alignment Bug

The SysV64 ABI mandates that the stack pointer (RSP) must be 16-byte aligned before executing a call instruction. If RSP is not 16-byte aligned, executing SSE/AVX vector instructions (like floating point or string copy optimizations emitted by LLVM) triggers a General Protection Fault (#GP).
Fix: In src/process.rs (L78), KernelStack::new aligns end downwards to the nearest 16-byte boundary: (end_raw.as_u64() / 16) * 16. Counting the 5 interrupt frame words and 15 register words ($20 \times 8 = 160$ bytes, an exact multiple of 16) ensures that when iretq executes, RSP remains perfectly 16-byte aligned.

5. Concrete Verification & Test Traces

When RustOS boots in QEMU, multiple concurrent tasks run simultaneously alongside the interactive shell:

Welcome to RustOS, from The Rusty Crew!!
Multitasking Scheduler starting...
Exiting task starting, about to exit...
Counter task incremented data to: 1
Clock task added 10, data now: 11
Counter task incremented data to: 12
/ >> 

This trace proves:

1. exiting_task called exit() and was safely moved to terminated_queue without corrupting other tasks.
2. counter_task and clock_task concurrently acquired SHARED_DATA via sync::Mutex, incremented values, and yielded cleanly.
3. The interactive shell task runs concurrently without blocking background scheduler ticks.

6. Interview Questions & Tactical Answers

Q1: "Why must the timer interrupt handler be declared with #[unsafe(naked)]?"

Tactical Answer: "Normal Rust and C functions include compiler-generated prologues and epilogues that create stack frames (e.g. push rbp; mov rbp, rsp; sub rsp, N). In a context switch, we must have byte-for-byte control over the exact stack layout so that RSP points precisely to the saved register structure. A naked function prevents LLVM from generating any setup code, ensuring our manual push and pop instructions mirror the register layout exactly."

Q2: "What is the difference between cooperative multitasking and preemptive multitasking?"

Tactical Answer: "In cooperative multitasking, a task runs until it explicitly relinquishes the CPU by calling a yield function. If a task enters an infinite loop, the entire OS hangs permanently. In preemptive multitasking, a periodic hardware timer interrupt involuntarily interrupts the running task, saves its state, and gives another task a turn. Preemption guarantees bounded execution time and system responsiveness regardless of task behavior."