From 23285ec1b8d00dd8753eefe193d0547659cc3c4d Mon Sep 17 00:00:00 2001
From: FlareCoding <albert.slepak@gmail.com>
Date: Sun, 8 Sep 2024 13:32:11 -0400
Subject: [PATCH] added full proper smp support and integrated it into the
 mandatory kernel initialization sequence

---
 efi/OVMF_VARS.fd                              | Bin 540672 -> 540672 bytes
 kernel/src/arch/x86/ap_startup.cpp            | 188 ++++++++++++++++++
 kernel/src/arch/x86/ap_startup.h              |  28 +++
 kernel/src/arch/x86/asm/ap_startup.s          |   2 +-
 kernel/src/entry/kernel_entry.cpp             |  12 +-
 .../tests/kernel_entry_ap_startup_test.cpp    | 147 ++------------
 kernel/src/memory/kheap.cpp                   |  21 ++
 kernel/src/sched/sched.cpp                    |  51 ++++-
 kernel/src/sched/sched.h                      |   8 +-
 9 files changed, 301 insertions(+), 156 deletions(-)
 create mode 100644 kernel/src/arch/x86/ap_startup.cpp
 create mode 100644 kernel/src/arch/x86/ap_startup.h

diff --git a/efi/OVMF_VARS.fd b/efi/OVMF_VARS.fd
index b482cfec3e0cb3a40408857f1e12e64e0fcabc02..c9f629067b58f3695aed345dac0dd3d22cbeeba2 100644
GIT binary patch
delta 130
zcmZo@P;6*W-0(%4xlWW}^LkTPX2$x>w))#S85<x}BZO*#P|Xmk1wyq#sJ6|2Yk(Ts
zfxJ?-X06_KtzJeTX4<aR%Y3bAv%`cPoZAFg1lYD85M}wm*`~nKroh^!z}BX~zD$8b
F0RYrgER_HN

delta 135
zcmZo@P;6*W-0(%4SxuB-^LkTPW=8eRw))#S88sl3CWO+0P}&em2SVvWC_M<J52Q*t
zC-16pVl-&}Thsoxh7pLFw*ReR_U_y!z#_o5Jwc3R0q5j~wtzMTmNo^}HU+je1@>hM
G90~xQlq=By

diff --git a/kernel/src/arch/x86/ap_startup.cpp b/kernel/src/arch/x86/ap_startup.cpp
new file mode 100644
index 0000000..fc147ec
--- /dev/null
+++ b/kernel/src/arch/x86/ap_startup.cpp
@@ -0,0 +1,188 @@
+#include "ap_startup.h"
+#include <acpi/acpi_controller.h>
+#include <paging/page_frame_allocator.h>
+#include <paging/page.h>
+#include <time/ktime.h>
+#include <kelevate/kelevate.h>
+#include <gdt/gdt.h>
+#include <interrupts/idt.h>
+#include <sched/sched.h>
+
+//
+// ------------------------------------ IMPORTANT -------------------------------------
+// Below addresses are also hardcoded and must match those in arch/x86/asm/ap_startup.s
+//
+
+#define AP_STARTUP_ASM_ADDRESS              (void*)0x8000
+#define AP_STARTUP_C_ENTRY_ADDRESS          (void*)0x9000
+#define AP_STARTUP_PAGE_TABLE_PTR_ADDRESS   (void*)0x15000
+
+#define AP_STARTUP_AP_RUNNING_COUNT_PTR     (volatile uint8_t*)0x11000
+#define AP_STARTUP_BSP_ID_PTR               (volatile uint8_t*)0x11008
+#define AP_STARTUP_BSP_SPINLOCK_PTR         (volatile uint8_t*)0x11010
+
+//
+// Each core will have 512 bytes of stack space during bootup
+// until a new 2k stack gets allocated in C entry code. 
+//
+#define AP_STARTUP_STACK_POOL_BASE          (void*)0x18000
+#define AP_STARTUP_STACK_POOL_TOP           (void*)0x70000
+#define AP_STARTUP_STACK_POOL_PAGE_COUNT    8
+#define AP_STARTUP_STACK_POOL_STACK_SIZE    512
+
+EXTERN_C __PRIVILEGED_CODE void __ap_startup_asm();
+EXTERN_C __PRIVILEGED_CODE void apStartupEntryC(int apicid);
+void apStartupEntryLowered();
+
+void _acquireApStartupSpinlockFlag() {
+    *AP_STARTUP_BSP_SPINLOCK_PTR = 0;
+}
+
+void _releaseApStartupSpinlockFlag() {
+    *AP_STARTUP_BSP_SPINLOCK_PTR = 1;
+}
+
+/**
+ * @brief Prepares memory mappings for Application Processor (AP) startup.
+ *
+ * This function locks the necessary physical pages and sets up the memory mappings 
+ * required for the AP startup process. It copies the AP startup assembly code 
+ * to a 16-bit real mode address and maps the C entry point to a lower physical 
+ * address. These mappings ensure that the AP can transition between real mode 
+ * and protected mode during the boot sequence.
+ * 
+ * Specifically, it:
+ * - Locks physical pages for AP startup code, stack, and runtime data.
+ * - Copies the startup assembly code to the real mode address.
+ * - Maps the C entry point and sets up synchronization pointers for tracking AP states.
+ *
+ * This setup allows AP cores to boot properly and transition to the system's SMP 
+ * environment.
+ */
+void _prepareApStartupMemoryMappings() {
+    auto& globalPageFrameAllocator = paging::getGlobalPageFrameAllocator();
+
+    // Copy the AP startup code to a 16bit address
+    globalPageFrameAllocator.lockPhysicalPage(AP_STARTUP_ASM_ADDRESS);
+    globalPageFrameAllocator.lockPhysicalPage(AP_STARTUP_C_ENTRY_ADDRESS);
+    globalPageFrameAllocator.lockPhysicalPage((void*)AP_STARTUP_AP_RUNNING_COUNT_PTR);
+    globalPageFrameAllocator.lockPhysicalPage(AP_STARTUP_PAGE_TABLE_PTR_ADDRESS);
+
+    // Stack region top is 0x70000, 8 pages of 512 bytes of stack 
+    globalPageFrameAllocator.lockPhysicalPages(AP_STARTUP_STACK_POOL_BASE, AP_STARTUP_STACK_POOL_PAGE_COUNT);
+    
+    // Copy the startup assembly code to the real mode address
+    memcpy(AP_STARTUP_ASM_ADDRESS, (void*)__ap_startup_asm, PAGE_SIZE);
+    
+    // Copy the C entry point address to a lower physical address
+    uint64_t apStartupCEntryAddress = (uint64_t)apStartupEntryC;
+    memcpy(AP_STARTUP_C_ENTRY_ADDRESS, &apStartupCEntryAddress, sizeof(uint64_t));
+
+    // Copy the page table address to a lower physical address
+    // to setup paging during AP startup initialization sequence.
+    memcpy(AP_STARTUP_PAGE_TABLE_PTR_ADDRESS, paging::g_kernelRootPageTable, PAGE_SIZE);
+
+    // Configure startup flags
+    *AP_STARTUP_AP_RUNNING_COUNT_PTR = 0;
+    *AP_STARTUP_BSP_ID_PTR = 0;
+
+    // Get the BSP's Local APIC ID
+    __asm__ __volatile__ ("mov $1, %%eax; cpuid; shrl $24, %%ebx;" : "=b"(*AP_STARTUP_BSP_SPINLOCK_PTR) ::);
+}
+
+void initializeApCores() {
+    auto& acpiController = AcpiController::get();
+    Madt* apicTable = acpiController.getApicTable();
+
+    auto& sched = RRScheduler::get();
+
+    RUN_ELEVATED({
+        // Copy the necessary resources and data to the lower physical address
+        // accessible from the 16-bit real mode that APs are in at this point.
+        _prepareApStartupMemoryMappings();
+
+        // Acquire the BSP spinlock to manage the
+        // initialization sequence of AP cores.
+        _acquireApStartupSpinlockFlag();
+
+        // For each core, register it in the scheduler for SMP support
+        // and send the appropriate startup signal.
+        // *Note* starting at 1 because BSP_ID is 0
+        for (size_t cpu = 1; cpu < apicTable->getCpuCount(); cpu++) {
+            // Create a scheduler run queue for each detected core
+            sched.registerCpuCore(cpu);
+
+            // Get the APIC ID of the core
+            uint8_t apicid = apicTable->getLocalApicDescriptor(cpu).apicId;
+
+            // Send the startup IPIs
+            bootAndInitApCore(apicid);
+        }
+
+        // Let the AP cores continue on their own asynchronously
+        _releaseApStartupSpinlockFlag();
+    });
+}
+
+void bootAndInitApCore(uint8_t apicid) {
+    auto& lapic = Apic::getLocalApic();
+
+    lapic->sendIpi(apicid, 0x500);
+    msleep(20);
+
+    lapic->sendIpi(apicid, 0x600 | ((uint32_t)((uint64_t)AP_STARTUP_ASM_ADDRESS >> 12)));
+    msleep(20);
+}
+
+EXTERN_C __PRIVILEGED_CODE
+void apStartupEntryC(int apicid) {
+    // Allocate a clean 2k kernel stack
+    void* apKernelStack = zallocPages(2);
+    uint64_t apKernelStackTop = (uint64_t)apKernelStack + 2 * PAGE_SIZE;
+    asm volatile ("mov %0, %%rsp" :: "r"(apKernelStackTop));
+
+    // Setup the GDT
+    initializeAndInstallGDT(apicid, (void*)apKernelStackTop);
+
+    // Initialize the default root kernel swapper task (this thread).
+    g_kernelSwapperTasks[apicid].state = ProcessState::RUNNING;
+    g_kernelSwapperTasks[apicid].pid = apicid;
+    zeromem(&g_kernelSwapperTasks[apicid].context, sizeof(CpuContext));
+    g_kernelSwapperTasks[apicid].context.rflags |= 0x200;
+    g_kernelSwapperTasks[apicid].elevated = 0;
+
+    // Set the current task in the per cpu region
+    __per_cpu_data.__cpu[apicid].currentTask = &g_kernelSwapperTasks[apicid];
+    __per_cpu_data.__cpu[apicid].currentTask->cpu = apicid;
+
+    // Install the existing IDT and enable interrupts
+    loadIdtr();
+    enableInterrupts();
+
+    // Enable syscalls
+    enableSyscallInterface();
+
+    // Update cr3
+    paging::setCurrentTopLevelPageTable(paging::g_kernelRootPageTable);
+    
+    // Setup a clean 8k per-cpu stack
+    char* usermodeStack = (char*)zallocPages(8);
+    size_t usermodeStackSize = 8 * PAGE_SIZE;
+    size_t userStackTop = (uint64_t)(usermodeStack + usermodeStackSize);
+
+    __call_lowered_entry(apStartupEntryLowered, (void*)userStackTop);
+    while (1);
+}
+
+void apStartupEntryLowered() {
+    // Initialize core's LAPIC
+    Apic::initializeLocalApic();
+
+    // Calibrate apic timer tickrate to 100 milliseconds
+    KernelTimer::calibrateApicTimer(100);
+
+    // Start the kernel-wide APIC periodic timer
+    KernelTimer::startApicPeriodicTimer();
+
+    while (1);
+}
diff --git a/kernel/src/arch/x86/ap_startup.h b/kernel/src/arch/x86/ap_startup.h
new file mode 100644
index 0000000..0b46f8d
--- /dev/null
+++ b/kernel/src/arch/x86/ap_startup.h
@@ -0,0 +1,28 @@
+#ifndef AP_STARTUP_H
+#define AP_STARTUP_H
+#include "apic.h"
+#include "msr.h"
+
+/**
+ * @brief Boots and initializes all secondary Application Processor (AP) cores.
+ *
+ * This function handles the process of bringing up all secondary AP cores 
+ * in the system. It sends the necessary startup signals to each AP core and 
+ * initializes them to be part of the system's SMP (Symmetric Multiprocessing) 
+ * environment, making them ready to execute tasks.
+ */
+void initializeApCores();
+
+/**
+ * @brief Boots and initializes a secondary Application Processor (AP) core.
+ *
+ * This function handles the process of bringing up an AP core by sending the
+ * necessary startup signals and configuring its initial state. Once booted, the
+ * core is initialized and ready to handle tasks as part of the system's SMP (Symmetric
+ * Multiprocessing) environment.
+ *
+ * @param apicid The cpu ID of the AP core to be booted and initialized.
+ */
+void bootAndInitApCore(uint8_t apicid);
+
+#endif
diff --git a/kernel/src/arch/x86/asm/ap_startup.s b/kernel/src/arch/x86/asm/ap_startup.s
index 6680328..3493538 100644
--- a/kernel/src/arch/x86/asm/ap_startup.s
+++ b/kernel/src/arch/x86/asm/ap_startup.s
@@ -12,7 +12,7 @@
 .equ stack_base, 0x18000      # Base of the stack region
 .equ stack_size, 512          # Size of each stack (512 bytes)
 
-.section .text
+.section .ktext
 .code16
 __ap_startup_asm:
     cli
diff --git a/kernel/src/entry/kernel_entry.cpp b/kernel/src/entry/kernel_entry.cpp
index cd63022..6c510a5 100644
--- a/kernel/src/entry/kernel_entry.cpp
+++ b/kernel/src/entry/kernel_entry.cpp
@@ -14,12 +14,13 @@
 #include <arch/x86/apic_timer.h>
 #include <arch/x86/gsfsbase.h>
 #include <arch/x86/pat.h>
+#include <arch/x86/ap_startup.h>
 #include <sched/sched.h>
 #include <syscall/syscalls.h>
 #include <kelevate/kelevate.h>
 #include <acpi/acpi_controller.h>
-#include <kprint.h>
 #include <time/ktime.h>
+#include <kprint.h>
 
 #include "tests/kernel_entry_tests.h"
 
@@ -177,13 +178,8 @@ void _kuser_entry() {
     auto& sched = RRScheduler::get();
     sched.init();
 
-    // If multiple cores are present, register all of them in the scheduler
-    if (acpiController.hasApicTable()) {
-        auto apicTable = acpiController.getApicTable();
-        size_t cpuCount = apicTable->getCpuCount();
-
-        kuPrint("==== Detected %lli CPUs ====\n", cpuCount);
-    }
+    // Bring up all available processor cores
+    initializeApCores();
 
 #ifdef KE_TEST_MULTITHREADING
     ke_test_multithreading();
diff --git a/kernel/src/entry/tests/kernel_entry_ap_startup_test.cpp b/kernel/src/entry/tests/kernel_entry_ap_startup_test.cpp
index 6f9d4c7..35c8dda 100644
--- a/kernel/src/entry/tests/kernel_entry_ap_startup_test.cpp
+++ b/kernel/src/entry/tests/kernel_entry_ap_startup_test.cpp
@@ -1,150 +1,27 @@
 #include "kernel_entry_tests.h"
-#include <core/kprint.h>
-#include <kelevate/kelevate.h>
-#include <paging/page.h>
-#include <paging/page_frame_allocator.h>
-#include <paging/phys_addr_translation.h>
-#include <memory/kmemory.h>
-#include <time/ktime.h>
-#include <acpi/acpi_controller.h>
-#include <arch/x86/apic.h>
-#include <arch/x86/msr.h>
-#include <arch/x86/x86_cpu_control.h>
-#include <gdt/gdt.h>
-#include <interrupts/idt.h>
+#include <arch/x86/ap_startup.h>
+#include <kprint.h>
 #include <sched/sched.h>
-#include <sync.h>
-
-EXTERN_C void __ap_startup_asm();
-
-void __ap_startup_user_entry();
-
-DECLARE_SPINLOCK(__ap_startup_lock);
-
-EXTERN_C
-void __ap_startup(int apicid) {
-    // Allocate a clean 8k kernel stack
-    void* apKernelStack = zallocPages(2);
-    uint64_t apKernelStackTop = (uint64_t)apKernelStack + 2 * PAGE_SIZE;
-    asm volatile ("mov %0, %%rsp" :: "r"(apKernelStackTop));
-
-    // Setup the GDT
-    initializeAndInstallGDT(apicid, (void*)apKernelStackTop);
-
-    // Initialize the default root kernel swapper task (this thread).
-    g_kernelSwapperTasks[apicid].state = ProcessState::RUNNING;
-    g_kernelSwapperTasks[apicid].pid = apicid;
-    zeromem(&g_kernelSwapperTasks[apicid].context, sizeof(CpuContext));
-    g_kernelSwapperTasks[apicid].context.rflags |= 0x200;
-    g_kernelSwapperTasks[apicid].elevated = 0;
-
-    // Set the current task in the per cpu region
-    __per_cpu_data.__cpu[apicid].currentTask = &g_kernelSwapperTasks[apicid];
-    __per_cpu_data.__cpu[apicid].currentTask->cpu = apicid;
-
-    // Install the existing IDT and enable interrupts
-    loadIdtr();
-    enableInterrupts();
-
-    // Enable syscalls
-    enableSyscallInterface();
-
-    // Update cr3
-    paging::setCurrentTopLevelPageTable(paging::g_kernelRootPageTable);
-    
-    // Setup per-cpu stack info
-    char* usermodeStack = (char*)zallocPages(8);
-    size_t usermodeStackSize = 8 * PAGE_SIZE;
-    size_t userStackTop = (uint64_t)(usermodeStack + usermodeStackSize);
-
-    __call_lowered_entry(__ap_startup_user_entry, (void*)userStackTop);
-    while (1);
-}
+#include <time/ktime.h>
 
 void sayHelloCore() {
     kuPrint("Hello from core %i!\n", getCurrentCpuId());
     exitKernelThread();
 }
 
-void __ap_startup_user_entry() {
-    uint64_t rsp = 0;
-    asm volatile ("mov %%rsp, %0" : "=r"(rsp));
-
-    uint8_t cpu = getCurrentCpuId();
-
-    kuPrint("[CPU%i] stack: 0x%llx\n", cpu, rsp);
-
-    // Initialize core's LAPIC
-    Apic::initializeLocalApic();
-
-    // Calibrate apic timer tickrate to 100 milliseconds
-    KernelTimer::calibrateApicTimer(100);
-
-    // Start the kernel-wide APIC periodic timer
-    KernelTimer::startApicPeriodicTimer();
-
-    auto& sched = RRScheduler::get();
-    Task* task = createKernelTask(sayHelloCore);
-    sched.addTask(task);
-
-    while (1);
+void sayGoodbyeCore() {
+    kuPrint("Goodbye from core %i!\n", getCurrentCpuId());
+    exitKernelThread();
 }
 
 void ke_test_ap_startup() {
-    auto& globalPageFrameAllocator = paging::getGlobalPageFrameAllocator();
-    auto& acpiController = AcpiController::get();
-    Madt* apicTable = acpiController.getApicTable();
-
     auto& sched = RRScheduler::get();
 
-    RUN_ELEVATED({
-        void* __ap_startup_code_real_mode_address = (void*)0x8000;
-
-        // Copy the AP startup code to a 16bit address
-        globalPageFrameAllocator.lockPhysicalPage(__ap_startup_code_real_mode_address);
-        globalPageFrameAllocator.lockPhysicalPage((void*)0x9000);
-        globalPageFrameAllocator.lockPhysicalPage((void*)0x11000);
-        globalPageFrameAllocator.lockPhysicalPage((void*)0x15000);
-        // Stack region top is 0x70000, 8 pages of 512 bytes of stack 
-        globalPageFrameAllocator.lockPhysicalPages((void*)0x18000, 8);
-        
-        // Copy the startup assembly code to the real mode address
-        memcpy(__ap_startup_code_real_mode_address, (void*)__ap_startup_asm, PAGE_SIZE);
-        
-        // Map the C entry point to a lower physical address
-        uint64_t __ap_startup_c_entry_address = (uint64_t)__ap_startup;
-        memcpy((void*)0x9000, &__ap_startup_c_entry_address, sizeof(uint64_t));
-
-        kprint("Sending SIPI to AP cores...\n");
-
-        volatile uint8_t* aprunning_ptr = (volatile uint8_t*)0x11000;  // count how many APs have started
-        uint8_t* bspid_ptr = (uint8_t*)0x11008; // BSP id
-        uint8_t* bspdone_ptr = (uint8_t*)0x11010; // Spinlock flag
-
-        memcpy((void*)0x15000, paging::g_kernelRootPageTable, PAGE_SIZE);
-
-        *aprunning_ptr = 0;
-        *bspid_ptr = 0;
-        *bspdone_ptr = 0;
-
-        // get the BSP's Local APIC ID
-        __asm__ __volatile__ ("mov $1, %%eax; cpuid; shrl $24, %%ebx;": "=b"(*bspid_ptr) : :);
-
-        auto& lapic = Apic::getLocalApic();
-
-        for (size_t i = 1; i < apicTable->getCpuCount(); i++) {
-            // Create a scheduler run queue for each detected core
-            sched.registerCpuCore(i);
-
-            uint8_t apicid = apicTable->getLocalApicDescriptor(i).apicId;
-
-            lapic->sendIpi(apicid, 0x500);
-            msleep(20);
-
-            lapic->sendIpi(apicid, 0x600 | ((uint32_t)((uint64_t)__ap_startup_code_real_mode_address >> 12)));
-            msleep(20);
-        }
+    for (int i = 0; i < 10; i++) {
+        Task* helloTask = createKernelTask(sayHelloCore);
+        Task* goodbyeTask = createKernelTask(sayGoodbyeCore);
 
-        *bspdone_ptr = 1;
-    });
+        sched.addTask(helloTask);
+        sched.addTask(goodbyeTask);
+    }
 }
diff --git a/kernel/src/memory/kheap.cpp b/kernel/src/memory/kheap.cpp
index 2e2c374..40e1a46 100644
--- a/kernel/src/memory/kheap.cpp
+++ b/kernel/src/memory/kheap.cpp
@@ -1,6 +1,7 @@
 #include "kheap.h"
 #include "kmemory.h"
 #include <paging/page_frame_allocator.h>
+#include <sync.h>
 #include <kprint.h>
 
 #define MIN_HEAP_SEGMENT_CAPACITY 1
@@ -12,6 +13,8 @@
 
 DynamicMemoryAllocator g_kernelHeapAllocator;
 
+DECLARE_SPINLOCK(__kheap_lock);
+
 DynamicMemoryAllocator& DynamicMemoryAllocator::get() {
     return g_kernelHeapAllocator;
 }
@@ -36,16 +39,20 @@ void DynamicMemoryAllocator::init(uint64_t base, size_t size) {
 }
 
 void* DynamicMemoryAllocator::allocate(size_t size) {
+    acquireSpinlock(&__kheap_lock);
+
     size_t newSegmentSize = size + sizeof(HeapSegmentHeader);
 
     // + sizeof(HeapSegmentHeader) is to account for the splitting
     HeapSegmentHeader* segment = _findFreeSegment(newSegmentSize + sizeof(HeapSegmentHeader));
 
     if (!segment) {
+        releaseSpinlock(&__kheap_lock);
         return nullptr;
     }
 
     if (!_splitSegment(segment, newSegmentSize)) {
+        releaseSpinlock(&__kheap_lock);
         return nullptr;
     }
 
@@ -55,10 +62,13 @@ void* DynamicMemoryAllocator::allocate(size_t size) {
     // Return the usable memory after the segment header
     uint8_t* usableRegionStart = reinterpret_cast<uint8_t*>(segment) + sizeof(HeapSegmentHeader);
 
+    releaseSpinlock(&__kheap_lock);
     return static_cast<void*>(usableRegionStart);
 }
 
 void DynamicMemoryAllocator::free(void* ptr) {
+    acquireSpinlock(&__kheap_lock);
+
     HeapSegmentHeader* segment = reinterpret_cast<HeapSegmentHeader*>(
         reinterpret_cast<uint8_t*>(ptr) - sizeof(HeapSegmentHeader)
     );
@@ -66,6 +76,7 @@ void DynamicMemoryAllocator::free(void* ptr) {
     // Verify the given pointer to be a heap segment header
     if (memcmp(segment->magic, (void*)KERNEL_HEAP_SEGMENT_HDR_SIGNATURE, 7) != 0) {
         kuPrint("Invalid pointer provided to free()!\n");
+        releaseSpinlock(&__kheap_lock);
         return;
     }
 
@@ -81,6 +92,8 @@ void DynamicMemoryAllocator::free(void* ptr) {
     if (segment->prev && segment->prev->flags.free) {
         _mergeSegmentWithPrevious(segment);
     }
+
+    releaseSpinlock(&__kheap_lock);
 }
 
 void* DynamicMemoryAllocator::reallocate(void* ptr, size_t newSize) {
@@ -89,6 +102,8 @@ void* DynamicMemoryAllocator::reallocate(void* ptr, size_t newSize) {
         return allocate(newSize);
     }
 
+    acquireSpinlock(&__kheap_lock);
+
     HeapSegmentHeader* segment = reinterpret_cast<HeapSegmentHeader*>(
         reinterpret_cast<uint8_t*>(ptr) - sizeof(HeapSegmentHeader)
     );
@@ -96,6 +111,7 @@ void* DynamicMemoryAllocator::reallocate(void* ptr, size_t newSize) {
     // Verify the given pointer to be a heap segment header
     if (memcmp(segment->magic, (void*)KERNEL_HEAP_SEGMENT_HDR_SIGNATURE, 7) != 0) {
         kuPrint("Invalid pointer provided to realloc()!\n");
+        releaseSpinlock(&__kheap_lock);
         return nullptr;
     }
 
@@ -103,8 +119,13 @@ void* DynamicMemoryAllocator::reallocate(void* ptr, size_t newSize) {
     // potentially resize (shrink) it and return the same pointer.
     if (segment->size >= newSize + sizeof(HeapSegmentHeader)) {
         _splitSegment(segment, newSize + sizeof(HeapSegmentHeader));
+        releaseSpinlock(&__kheap_lock);
         return ptr;
     } else {
+        // Release the currently held lock because further `allocate`
+        // and `free` calls will attempt to acquire the lock themselves.
+        releaseSpinlock(&__kheap_lock);
+
         // Allocate new memory, and check if allocation was successful
         void* newPtr = allocate(newSize);
         if (!newPtr) {
diff --git a/kernel/src/sched/sched.cpp b/kernel/src/sched/sched.cpp
index b9b4d02..5b54cc8 100644
--- a/kernel/src/sched/sched.cpp
+++ b/kernel/src/sched/sched.cpp
@@ -3,10 +3,14 @@
 #include <paging/page.h>
 #include <gdt/gdt.h>
 #include <kelevate/kelevate.h>
+#include <sync.h>
 
 RRScheduler s_globalRRScheduler;
 Task g_kernelSwapperTasks[MAX_CPUS] = {};
 
+DECLARE_SPINLOCK(__sched_lock);
+DECLARE_SPINLOCK(__sched_load_balancing_lock);
+
 size_t g_availableTaskPid = 10;
 size_t _allocateTaskPid() {
     size_t pid = g_availableTaskPid++;
@@ -156,15 +160,27 @@ bool RRScheduler::addTask(Task* task, int cpu) {
         return false;
     }
 
+    acquireSpinlock(&__sched_lock);
+
     // Assign a cpu to the task
     task->cpu = cpu;
 
     auto& runQueue = m_runQueues[cpu];
-    return runQueue->addTask(task);
+    bool ret = runQueue->addTask(task);
+
+    releaseSpinlock(&__sched_lock);
+    return ret;
 }
 
 bool RRScheduler::addTask(Task* task) {
-    return addTask(task, getNextAvailableCpu());
+    acquireSpinlock(&__sched_load_balancing_lock);
+
+    int cpu = _getNextAvailableCpu();
+    bool ret = addTask(task, cpu);
+
+    releaseSpinlock(&__sched_load_balancing_lock);
+
+    return ret;
 }
 
 bool RRScheduler::removeTask(Task* task, int cpu) {
@@ -178,8 +194,13 @@ bool RRScheduler::removeTask(pid_t pid, int cpu) {
         return false;
     }
 
+    acquireSpinlock(&__sched_lock);
+
     auto& runQueue = m_runQueues[cpu];
-    return runQueue->removeTask(pid);
+    bool ret = runQueue->removeTask(pid);
+
+    releaseSpinlock(&__sched_lock);
+    return ret;
 }
 
 Task* RRScheduler::getCurrentTask(int cpu) {
@@ -189,8 +210,13 @@ Task* RRScheduler::getCurrentTask(int cpu) {
         return nullptr;
     }
 
+    acquireSpinlock(&__sched_lock);
+
     auto& runQueue = m_runQueues[cpu];
-    return runQueue->getCurrentTask(); 
+    Task* task = runQueue->getCurrentTask();
+
+    releaseSpinlock(&__sched_lock);
+    return task;
 }
 
 Task* RRScheduler::peekNextTask(int cpu) {
@@ -200,8 +226,13 @@ Task* RRScheduler::peekNextTask(int cpu) {
         return nullptr;
     }
 
+    acquireSpinlock(&__sched_lock);
+
     auto& runQueue = m_runQueues[cpu];
-    return runQueue->peekNextTask();
+    Task* task = runQueue->peekNextTask();
+
+    releaseSpinlock(&__sched_lock);
+    return task;
 }
 
 void RRScheduler::scheduleNextTask(int cpu) {
@@ -210,16 +241,20 @@ void RRScheduler::scheduleNextTask(int cpu) {
         asm volatile ("hlt");
         return;
     }
+
+    acquireSpinlock(&__sched_lock);
     
     auto& runQueue = m_runQueues[cpu];
     runQueue->scheduleNextTask();
+
+    releaseSpinlock(&__sched_lock);
 }
 
-size_t RRScheduler::getNextAvailableCpu() {
-    size_t cpu = 0;
+int RRScheduler::_getNextAvailableCpu() {
+    int cpu = 0;
     size_t leastTaskCount = m_runQueues[cpu]->size();
 
-    for (size_t i = 1; i < m_usableCpuCount; ++i) {
+    for (int i = 1; i < (int)m_usableCpuCount; ++i) {
         size_t cpuTaskCount = m_runQueues[i]->size();
         if (cpuTaskCount < leastTaskCount) {
             cpu = i;
diff --git a/kernel/src/sched/sched.h b/kernel/src/sched/sched.h
index 4f488bd..c0f6db3 100644
--- a/kernel/src/sched/sched.h
+++ b/kernel/src/sched/sched.h
@@ -89,16 +89,16 @@ class RRScheduler {
     // the next available task.
     void scheduleNextTask(int cpu);
 
-    // Calculates the next least loaded CPU
-    // core to schedule next task(s) on.
-    size_t getNextAvailableCpu();
-
 private:
     // Per-core task run queues
     kstl::vector<RoundRobinRunQueue*> m_runQueues;
 
     // Number of actual usable cpu cores
     size_t m_usableCpuCount;
+
+    // Calculates the next least loaded CPU
+    // core to schedule next task(s) on.
+    int _getNextAvailableCpu();
 };
 
 //