Processes & Context Switching: Teaching the Kernel to Multitask

In the last post, we gave the kernel the ability to create things at runtime. Every future data structure in this kernel will be born out of kmalloc, and now we’re about to call it on something important: a process control block.

Right now, our kernel runs as a single-threaded process. kernel_main runs, calls timer_init, sets up the heap, enables interrupts, and then parks in a wfi loop waiting for the timer to fire. Right now, there is only one stack, one program counter and one task. That’s not an operating system; it’s a very fancy main().

After this post, your kernel will have two tasks running concurrently, interleaving their output on the UART using a round-robin scheduler. The single-threaded boot sequence ends here.

What We’re Building Today

The end state of this post is a kernel that creates two tasks, hands control to the first one, and lets the scheduler bounce between them. Each task prints a counter to the UART and voluntarily yields the CPU. The timer interrupt drives preemptive switching on top of that.

We need to create three new source files for this to work:

• include/kernel/scheduler.h: defines pcb_t (the process control block), cpu_context_t (the saved register state), and the scheduler’s public API.
• arch/arm64/context_switch.S: 20 lines of assembly that save and restore the CPU register state when switching tasks.
• kernel/scheduler.c: the round-robin scheduler: task_create, scheduler_yield, and scheduler_tick for timer-driven preemption.

Here’s where the two new task stacks live in physical memory after task_create runs twice:

The kernel stack (the one boot.S set up) is abandoned after scheduler_start(). From that point on, every stack frame lives in one of the per-task stacks that kmalloc provided. The kernel heap we created in the previous post is now doing real work.

The Process Control Block

Every task needs somewhere to store its “I’ll be right back” note. This note contains the register values, stack pointer, and metadata that let the scheduler resume it exactly where it left off. That structure is the PCB, or process control block. We can define a structure for it like this:

/* include/kernel/scheduler.h */

typedef enum {
    TASK_READY,     /* in the runqueue, waiting for CPU time  */
    TASK_RUNNING,   /* currently executing on the CPU         */
    TASK_BLOCKED,   /* waiting for an event (future work)     */
    TASK_DEAD,      /* finished; slot can be reclaimed        */
} task_state_t;

typedef struct {
    uint64_t x19, x20, x21, x22, x23, x24;
    uint64_t x25, x26, x27, x28;
    uint64_t fp;    /* x29 — frame pointer */
    uint64_t lr;    /* x30 — link register / return address   */
    uint64_t sp;    /* stack pointer                          */
} cpu_context_t;

typedef struct pcb {
    uint32_t        pid;           /* unique task id              */
    task_state_t    state;         /* current scheduling state    */
    cpu_context_t   ctx;           /* saved register state        */
    uint8_t        *stack_base;    /* bottom of this task's stack */
    size_t          stack_size;    /* stack size in bytes         */
    void          (*entry)(void);  /* entry function (for debug)  */
} pcb_t;

This struct is allocated with kmalloc. We should have one call to kmalloc(sizeof(pcb_t)) per task. The fields are straightforward except for cpu_context_t, which I will explain in the next section.

The CPU Context: Why Only Callee-Saved?

The most common point of confusion when building a first-time scheduler is deciding which registers to save during a context switch. While a CPU might have dozens of registers, a context switch in a typical OS occurs within a function call (e.g., yield() -> schedule() -> context_switch()). Because we are jumping between tasks from within a C function, we can lean on the Application Binary Interface (ABI) calling conventions to do the heavy lifting.

The ARM AAPCS64 convention splits registers into two categories. Understanding this split explains why our cpu_context_t is so small.

• Caller-saved (x0–x18): These are “volatile” registers. The compiler assumes any function call might overwrite them. Therefore, if the calling task had important data in x0, it already saved that data to its own stack before calling scheduler_yield(). By the time our context switch code runs, these registers are effectively “dead” or already protected.
• Callee-saved (x19–x28, FP, LR): These are “static” registers. The ABI guarantees that the values returned by a function will be the same after it returns as they were before it was called. Since our context switch is the function call, we must manually save these to the task’s context buffer so we can restore them when the task resumes.
• Stack Pointer (SP): While not technically a “saved” register in the ABI sense, we must save the SP to know where the task’s unique stack begins when we switch back to it.

By leveraging the ABI, we only need to track 13 registers, rather than all 31 general-purpose registers or the massive NEON/floating-point state. We are only saving the values that a C function expects to survive across a call boundary.

The field order in your C struct must perfectly match the assembly offsets in context_switch.S. If they drift apart, you will restore the wrong values to the wrong registers, leading to immediate crashes.

Memory Layout (104 Bytes Total):

By saving only what is strictly necessary, we keep the context_switch function fast and each task’s memory footprint small.

The Context Switch Assembly

The entire switch mechanism is implemented in 20 lines of assembly code. Half of it saves the old task’s registers. The other half loads the new task’s registers. The ret at the end is what makes it feel like magic. You can step through each instruction in this context switch example:

context_switch:

  stp  x19, x20, [x0, #0]

  stp  x21, x22, [x0, #16]

  stp  x23, x24, [x0, #32]

  stp  x25, x26, [x0, #48]

  stp  x27, x28, [x0, #64]

  stp  x29, x30, [x0, #80]

  mov  x2, sp

  str  x2,       [x0, #96]

  ldp  x19, x20, [x1, #0]

  ldp  x21, x22, [x1, #16]

  ldp  x23, x24, [x1, #32]

  ldp  x25, x26, [x1, #48]

  ldp  x27, x28, [x1, #64]

  ldp  x29, x30, [x1, #80]

  ldr  x2,       [x1, #96]

  mov  sp, x2

  ret

The function prototype is:

void context_switch(cpu_context_t *old, cpu_context_t *new_ctx);

The key thing to hold in your head: context_switch is a function call from the old task’s perspective. When it eventually “returns” to the old task (after a future switch back), it behaves like any other function return. The compiler’s caller-saved registers are already handled by the compiler, and the callee-saved ones are restored by context_switch. From the old task’s perspective, scheduler_yield() simply took a long time to return.

The Initial Context Trick

task_create sets up a new PCB and allocates a stack from the heap, but there’s a subtlety: the first time context_switch restores a brand-new task, no C code has run in that task yet. We need ret to jump to the task’s entry function, not to some random saved return address. The trick is in task_create’s initialisation of ctx.lr:

pcb_t *task_create(void (*entry)(void)) {
    pcb_t *pcb   = kmalloc(sizeof(pcb_t));
    uint8_t *stack = kmalloc(TASK_STACK_SIZE);

    /* Zero everything first — callee-saved regs default to 0 */
    zero_memory(pcb, sizeof(pcb_t));

    /* Stack grows downward. Initial sp = top of the allocated block.
     * Mask the low 4 bits: AArch64 requires 16-byte sp alignment. */
    uint64_t stack_top = (uint64_t)(stack + TASK_STACK_SIZE) & ~(uint64_t)0xF;

    pcb->pid        = task_count;
    pcb->state      = TASK_READY;
    pcb->entry      = entry;
    pcb->stack_base = stack;
    pcb->stack_size = TASK_STACK_SIZE;

    /* The critical setup: lr = entry function address.
     * When context_switch first restores this task, it loads lr = entry,
     * then executes ret, which jumps to entry(). The task starts running. */
    pcb->ctx.sp = stack_top;
    pcb->ctx.lr = (uint64_t)entry;

    tasks[task_count++] = pcb;
    return pcb;
}

The two magic lines are pcb->ctx.sp = stack_top and pcb->ctx.lr = (uint64_t)entry. It’s just a stack pointer pointing to fresh memory and a return address pointing to the task’s entry function. When context_switch runs ret, the new task starts executing from entry() with a clean stack. Very elegant if you ask me.

The Round-Robin Scheduler

The scheduler maintains a fixed array of PCB pointers and a current_task cursor. On each call to schedule_next(), the cursor advances by one, wrapping around, and context_switch swaps in the next task.

/* kernel/scheduler.c */

static pcb_t *tasks[MAX_TASKS];
static int    task_count   = 0;
static int    current_task = 0;

static void schedule_next(void) {
    if (task_count <= 1) return;

    int old_idx = current_task;
    int new_idx = (old_idx + 1) % task_count;

    /* Skip DEAD tasks */
    int scanned = 0;
    while (tasks[new_idx]->state == TASK_DEAD) {
        new_idx = (new_idx + 1) % task_count;
        if (++scanned >= task_count) return;
    }

    if (new_idx == old_idx) return; /* only one runnable task */

    pcb_t *old = tasks[old_idx];
    pcb_t *new = tasks[new_idx];

    old->state   = TASK_READY;
    new->state   = TASK_RUNNING;
    current_task = new_idx;

    context_switch(&old->ctx, &new->ctx);
}

void scheduler_yield(void) { schedule_next(); }
void scheduler_tick(void)  { schedule_next(); }

scheduler_yield() and scheduler_tick() both call schedule_next(). The difference is who calls them: scheduler_yield() is called voluntarily by a task (cooperative); scheduler_tick() is called by the timer IRQ handler (preemptive). The scheduling logic is identical either way.

The Register State Before and After

Here’s what the CPU register context looks like at the moment context_switch is called when task A yields for the first time. These are the values written into task_a->ctx:

Compare that to task B’s initial cpu_context_t, set up by task_create before task B has ever run:

Starting the First Task

There’s one bootstrapping problem: to call context_switch(old, new), we need a valid “old” PCB to save into. But before the scheduler starts, there is no current task. We’re still running in kernel_main. The solution is a stripped-down version of context_switch called start_first_task. It skips the save half entirely, it’s just the restore and ret:

/* arch/arm64/context_switch.S */

start_first_task:
    ldp  x19, x20, [x0, #0]
    ldp  x21, x22, [x0, #16]
    ldp  x23, x24, [x0, #32]
    ldp  x25, x26, [x0, #48]
    ldp  x27, x28, [x0, #64]
    ldp  x29, x30, [x0, #80]
    ldr  x2,       [x0, #96]
    mov  sp, x2
    ret

scheduler_start() calls this once and never returns. kernel_main’s stack is abandoned nd from this point onward, every execution context lives on a task-owned stack.

/* kernel/scheduler.c */

void scheduler_start(void) {
    pcb_t *first = tasks[0];
    first->state   = TASK_RUNNING;
    current_task   = 0;
    start_first_task(&first->ctx);
    /* unreachable */
}

Wiring It Into kernel_main

The updated kernel/main.c for this post is straightforward. After everything is initialised, we create two tasks and hand control to the scheduler:

/* kernel/main.c */

#include "drivers/uart.h"
#include "kernel/mmu.h"
#include "kernel/exceptions.h"
#include "kernel/gic.h"
#include "kernel/timer.h"
#include "kernel/heap.h"
#include "kernel/scheduler.h"

static void task_a(void) {
    uint32_t count = 0;
    while (1) {
        kprint("[task A] count=");
        kprint_uint(count++);
        kprint("\n");
        scheduler_yield();
    }
}

static void task_b(void) {
    uint32_t count = 0;
    while (1) {
        kprint("[task B] count=");
        kprint_uint(count++);
        kprint("\n");
        scheduler_yield();
    }
}

void kernel_main(void) {
    mmu_init();
    uart_init();
    kprint("PurgatoryOS booting...\n");

    exceptions_init();
    kprint("Exceptions: ready\n");

    gic_init();
    kprint("GIC: ready\n");

    timer_init(500);
    kprint("Timer: armed\n");

    heap_init();
    kprint("Heap: ready\n");

    scheduler_init();

    task_create(task_a);
    task_create(task_b);

    asm volatile("msr daifclr, #2");
    kprint("Interrupts enabled.\n");

    kprint("Starting scheduler...\n");
    scheduler_start();

    /* never reached */
    while (1) { asm volatile("wfi"); }
}

We also need to hook scheduler_tick() into the timer interrupt. In kernel/timer.c, add a call to scheduler_tick() at the end of timer_tick():

/* kernel/timer.c — update timer_tick() */

#include "kernel/scheduler.h"

void timer_tick(void) {
    tick_count++;

    /* ... re-arm the timer ... */
    asm volatile("msr cntp_tval_el0, %0" :: "r"(reload_ticks));

    /* Preemptive context switch on every timer tick */
    scheduler_tick();
}

Now build and run:

PurgatoryOS booting...
Exceptions: ready
GIC: ready
Timer: armed
[heap] init: free-list allocator
Heap: ready
[sched] init: round-robin scheduler
[sched] task created: pid=0
[sched] task created: pid=1
Interrupts enabled.
Starting scheduler...
[sched] starting first task (pid=0)
[task A] count=0
[task B] count=0
[task A] count=1
[task B] count=1
[task A] count=2
[task B] count=2
[task A] count=3
[task B] count=3
tick #1
[task A] count=4
[task B] count=4
[task A] count=5
[task B] count=5
tick #2
[task A] count=6
[task B] count=6

Task A and task B interleave on every yield. The timer tick appears every 500 ms and triggers a preemptive switch when it fires. The kernel is finally multitasking: both cooperatively (via scheduler_yield()) and preemptively (via the timer interrupt).

A Note on the Timer + Preemption

You might be wondering how the timer-triggered preemption actually works, since timer_tick() calls scheduler_tick(), which calls a function that just switches the stack pointer, context_switch(). Let’s trace through it. When the timer fires while task A is in its loop:

1. vectors.S saves all general-purpose registers (x0–x30) onto task A’s current stack via SAVE_REGS. The stack pointer is now task_a_sp - N bytes (N depends on the number of registers saved).
2. irq_handler runs, dispatches to timer_tick(), which calls scheduler_tick().
3. scheduler_tick() calls context_switch(&taskA->ctx, &taskB->ctx).
4. context_switch saves task A’s callee-saved registers including the current sp (which points to the IRQ frame on task A’s stack) into taskA->ctx.
5. context_switch restores task B’s registers and switches the stack pointer to task B’s stack. ret jumps to task B’s saved lr.
6. Task B runs — either from its entry function (first run) or from inside its own schedule_next() call (subsequent runs).
7. Task B eventually yields. context_switch restores task A’s registers, setting sp back to the value that included the IRQ frame.
8. We return through scheduler_tick() → timer_tick() → irq_handler() → vectors.S RESTORE_REGS, which pops the saved registers from task A’s stack. eret resumes task A exactly where the timer interrupted it.

The IRQ frame on task A’s stack is left intact throughout the preemptive switch. When we return to task A via context_switch, we re-enter the IRQ handler chain and unwind it cleanly via RESTORE_REGS and eret. The saved sp in taskA->ctx includes the stack space consumed by the IRQ frame. Doing all these steps is what makes it work.

One caveat: a real kernel disables interrupts (msr daifset, #2) before entering context_switch and re-enables them after. We’re not doing that here, which means a second timer interrupt could fire mid-switch. For our demo tasks, the 500 ms interval is long enough that this is extremely unlikely. For a production kernel, the fix is a single instruction pair around the switch.

What Broke (And Why)

The bug I spent the longest time on wasn’t in context_switch.S. It was in task_create. My first version of the initial stack pointer calculation looked like this:

/* BROKEN: off-by-one in stack setup */
uint64_t *stack_top = (uint64_t *)(stack + TASK_STACK_SIZE);
pcb->ctx.sp = (uint64_t)stack_top;

On AArch64, the stack pointer must be 16-byte aligned at all times. kmalloc returns 8-byte-aligned memory. TASK_STACK_SIZE is $8 * 1024 = 8192$, a multiple of 8 but also a multiple of 16, so stack + TASK_STACK_SIZE is 16-byte aligned if stack itself is 16-byte aligned. Our allocator aligns to 8 bytes, not 16. So the stack might be 8-byte-aligned but not 16-byte-aligned.

The symptom was subtle: task B ran its first iteration perfectly, but the second iteration triggered a Data Abort exception with EC = 0x25 (PC alignment fault or SP alignment fault). The exception handler printed ESR_EL1 = 0x96000086, which decodes to EC 0x26 (Data Abort from the current EL). FAR_EL1 pointed somewhere in task B’s stack. After a long time, I finally looked at the SP_EL1 register in GDB:

(gdb) info register sp
sp  0x4002e7f8  0x4002e7f8

0x4002e7f8 the low nibble is 8. On AArch64, the SP alignment check fires (SCTLR_EL1.SA bit) whenever an ldp or stp instruction uses the SP as a base and the SP is not 16-byte aligned. A misaligned SP is a guaranteed crash on any stp xN, xM, [sp, #-16]! in the function prologue.

The fix was to mask the low bits unconditionally:

uint64_t stack_top = ((uint64_t)(stack + TASK_STACK_SIZE)) & ~(uint64_t)0xF;
pcb->ctx.sp = stack_top;

The & ~0xF rounds the top of the stack down to the nearest 16-byte boundary. If stack + TASK_STACK_SIZE was already aligned, this is a no-op. If it were misaligned by 8 bytes, this wastes 8 bytes but prevents the alignment fault. On AArch64, the SP must be 16-byte aligned before any function call. This is a hard architectural requirement.

QEMU is inconsistent about enforcing this. Some kernel configurations enforce it from the start; others let you get away with it for a while. The alignment check is controlled by SCTLR_EL1.SA (bit 3). We set SCTLR_EL1 during mmu_init(); check that bit if you’re seeing alignment faults you didn’t expect.

What’s Next

We have two tasks that run and yield, and the timer preempts them. But they can’t ask the kernel for anything. They can’t request memory, open a “file,” or exit cleanly. They are completely isolated from the rest of the kernel with no communication channel.

In the next post, we will make the change so they are no longer isolated. The svc instruction that drops from EL0 (user space) to EL1 (kernel space). Right now, our tasks run in EL1 alongside the kernel, which is fine for a demo but not for a real OS. The next post also introduces the privilege-level boundary and a small syscall table so that tasks can ask the kernel to do things on their behalf.

The two building blocks we have now, a working heap (kmalloc) and a working scheduler (scheduler_yield), are what every other OS primitive will depend on from here. New filesystem nodes will be kmalloc’d. New processes will be born through task_create. The kernel is no longer a single-threaded boot sequence. It is, unambiguously, an operating system.

Sources

• Context Switch — Wikipedia — clear overview of what a context switch is, including the distinction between hardware and software context and the role of the PCB. A good first-principles read.
• Linux Context Switching Internals: Part 1 — codingconfessions.com — deep-dive into how the Linux kernel handles process state and memory during a context switch. More complex than our implementation, but the mental model transfers directly.
• SO2 Lecture 03: Processes — linux-kernel-labs.github.io — the Systems Operating lecture covering PCB structures, task_struct, and the Linux scheduler’s relationship to context_switch. Good academic companion to this post.
• CPU Scheduler Implementation Hints — kernel.org — the Linux kernel’s own documentation on what an architecture port must provide for scheduler support. Our context_switch.S and cpu_context_t satisfy the minimum contract described here.
• AAPCS64: ABI, Calling Conventions & Machine Registers — medium.com/@tunacici7 — thorough walkthrough of the AArch64 Procedure Call Standard. The register classification (caller-saved vs callee-saved) is central to understanding why cpu_context_t only saves x19–x28.
• ARM64 Registers and Basic Instructions — ohyaan.github.io — concise reference for AArch64 registers, including the calling convention roles of each register. Useful if the AAPCS64 spec feels heavy.
• AArch64 Procedure Call Standard — ARM Software ABI (aapcs64.rst) — the authoritative specification. Section 6.1 covers register roles and preservation requirements. This is the ground truth for why we save exactly the registers we save.
• Lab 4: Preemptive Multitasking — CS-3210 Spring 2020, GaTech — the closest public teaching resource to what we built in this post: PCB, context switch, and a round-robin scheduler for a bare-metal kernel. Worth reading to compare approaches.
• Rust-OS Kernel: Task Scheduler — nfil.dev — a Rust no_std implementation of a task scheduler with context switching. Good preview of what Post 11 (Bringing in Rust) will look like when we port the scheduler.
• Context Switch Overheads for Linux on ARM Platforms — ACM DL — empirical measurement of context switch overhead on ARM. Establishes that saving/restoring registers is the dominant cost, which is why our minimal 13-register save matters even in a teaching kernel.
• AArch64 MMU Programming — lowenware.com — useful reference for the memory system context in which our scheduler runs, particularly the ASID (Address Space ID) concept that allows the TLB to tag entries per-process, avoiding full flushes on context switch.