memory: Fix SMP safety and perf issues

Author: marv7000

kernel/src/arch/x86_64/consts.rs

@@ -282,3 +282,4 @@ pub const IDT_SX: u8 = 0x1E;
 pub const IDT_IPI_PANIC: u8 = 0xFF;
 pub const IDT_IPI_RESCHED: u8 = 0xFE;
 pub const IDT_IPI_SHOOTDOWN: u8 = 0xFD;
+pub const IDT_TIMER: u8 = 0xFC;

kernel/src/arch/x86_64/irq.rs

@@ -269,15 +269,20 @@ unsafe extern "C" fn idt_handler(context: *mut Context) {
         consts::IDT_IPI_PANIC => {
             halt();
         }
-        consts::IDT_IPI_RESCHED => {
-            clock::handle_tick();
-            unsafe { crate::arch::sched::preempt_disable() };
-            let should_reschedule = unsafe { crate::arch::sched::preempt_enable() };
+        consts::IDT_TIMER => {
+            if CpuData::get().id == 0 {
+                clock::handle_tick();
+            }
+            let should_reschedule = CpuData::get().scheduler.tick();
             LAPIC.get().eoi();
             if should_reschedule {
-                CpuData::get().scheduler.reschedule();
+                CpuData::get().scheduler.request_reschedule();
             }
         }
+        consts::IDT_IPI_RESCHED => {
+            LAPIC.get().eoi();
+            CpuData::get().scheduler.handle_remote_reschedule(from_user);
+        }
         // Any other ISR is an IRQ with a dynamic handler.
         _ => {
             match &mut IRQ_LINES.get().lock()[isr as usize - 0x20] {

kernel/src/arch/x86_64/system/apic.rs

@@ -1,7 +1,7 @@
 use crate::{
     arch::x86_64::{
         asm,
-        consts::{self, IDT_IPI_RESCHED},
+        consts::{self, IDT_TIMER},
         irq::IRQ_LINES,
     },
     clock,
@@ -151,7 +151,7 @@ impl LocalApic {
         );
 
         // Finally, run the periodic timer interrupt.
-        lapic.write_reg(lapic_regs::LVT_TR, IDT_IPI_RESCHED as u64 | 0x20000);
+        lapic.write_reg(lapic_regs::LVT_TR, IDT_TIMER as u64 | 0x20000);
         lapic.write_reg(lapic_regs::DCR, 3);
         lapic.write_reg(
             lapic_regs::ICR_TIMER,

kernel/src/arch/x86_64/system/smp.rs

@@ -322,7 +322,6 @@ fn start_ap(temp_cr3: u32, id: u32) {
 
     unsafe {
         KernelAlloc::dealloc(mem, 1);
-        KernelAlloc::dealloc_bytes(stack_mem, KERNEL_STACK_SIZE);
     }
 }
 

kernel/src/memory/bump.rs

@@ -1,8 +1,8 @@
 //! Early memory setup and allocator.
 
 use super::{
-    PhysAddr,
     pmm::{AllocFlags, PageAllocator},
+    PhysAddr,
 };
 use crate::arch::{self};
 use alloc::alloc::AllocError;
@@ -25,7 +25,7 @@ impl PageAllocator for BumpAllocator {
             let mem = PhysAddr(BUMP_CURRENT.fetch_add(bytes, Ordering::Relaxed));
 
             if !flags.contains(AllocFlags::NoZero) {
-                unsafe { (mem.as_hhdm() as *mut u8).write_bytes(0, bytes) };
+                mem.zero_hhdm(bytes);
             }
             return Ok(mem);
         }
@@ -38,4 +38,10 @@ impl PageAllocator for BumpAllocator {
             "The bump allocator is not supposed to free anything. Remove this .dealloc()"
         )
     }
+
+    unsafe fn dealloc_bytes(_addr: PhysAddr, _bytes: usize) {
+        unimplemented!(
+            "The bump allocator is not supposed to free anything. Remove this .dealloc_bytes()"
+        )
+    }
 }

kernel/src/memory/mod.rs

@@ -21,7 +21,7 @@ use crate::{
     arch::{self, virt::get_page_size},
     boot::{BootInfo, PhysMemoryUsage},
     memory::virt::{allocator::VirtualAllocator, mmu::PageTable},
-    util::{align_down, align_up, mutex::spin::SpinMutex},
+    util::{align_down, align_up, divide_up, mutex::spin::SpinMutex},
 };
 use alloc::sync::Arc;
 use bump::BumpAllocator;
@@ -44,6 +44,10 @@ impl PhysAddr {
     pub fn as_hhdm<T>(self) -> *mut T {
         VirtAddr(self.0 + HHDM_START.get().0).as_ptr()
     }
+
+    pub(crate) fn zero_hhdm(self, len: usize) {
+        unsafe { ptr::write_bytes(self.as_hhdm::<u8>(), 0, len) };
+    }
 }
 
 /// Represents a virtual address. It can't be directly read from or written to.
@@ -325,36 +329,28 @@ pub fn MEMORY_STAGE() {
         arch::virt::get_hhdm_base().value()
     );
 
-    // We record metadata for every single page of available memory in a large array.
-    // This array is contiguous in virtual memory, but is sparsely populated.
-    // Only those array entries which represent usable memory are mapped.
+    // We record metadata for every single physical page in a large array.
+    // This array is contiguous in virtual memory so it can be represented as
+    // an ordinary Rust slice by the physical allocator.
 
     // The offset where we start mapping the page array.
     log!("Highest physical address is {:#018x}", highest_phys.0);
     let page_base = arch::virt::get_pfndb_base();
-    let page_length = highest_phys.0 / size_of::<Page>();
     let page_size = get_page_size();
-
-    memory_map
-        .iter()
-        .filter(|x| x.length != 0)
-        .for_each(|entry| {
-            let length = align_up((entry.length / page_size) * size_of::<Page>(), page_size);
-            let virt = align_down(
-                (page_base + entry.address.0 / page_size * size_of::<Page>()).value(),
-                page_size,
-            );
-
-            for page in (0..=length).step_by(page_size) {
-                table
-                    .map_single::<BumpAllocator>(
-                        (virt + page).into(),
-                        BumpAllocator::alloc(1, AllocFlags::empty()).unwrap(),
-                        VmFlags::Read | VmFlags::Write,
-                    )
-                    .unwrap();
-            }
-        });
+    let page_length = divide_up(highest_phys.0, page_size);
+    let page_array_len = align_up(page_length * size_of::<Page>(), page_size);
+
+    // The physical allocator keeps the PFN database as a Rust slice, which
+    // requires the whole range to be mapped and backed by memory.
+    for page in (0..page_array_len).step_by(page_size) {
+        table
+            .map_single::<BumpAllocator>(
+                (page_base.value() + page).into(),
+                BumpAllocator::alloc(1, AllocFlags::empty()).unwrap(),
+                VmFlags::Read | VmFlags::Write,
+            )
+            .unwrap();
+    }
 
     log!(
         "Initalized page array region at {:#018x}",
@@ -386,7 +382,8 @@ pub fn MEMORY_STAGE() {
     // ----------------------------------------
 
     // Initialize the physical memory allocator.
-    pmm::init(&memory_map, (page_base.as_ptr(), page_length));
+    let page_db = unsafe { core::slice::from_raw_parts_mut(page_base.as_ptr(), page_length) };
+    pmm::init(&memory_map, page_db);
 
     // Save the page table.
     unsafe { virt::KERNEL_PAGE_TABLE.init(Arc::new(table)) };