More SMP fixes

kernel/src/main.rs

@@ -8,7 +8,7 @@ use core::{
     arch::global_asm,
     panic::PanicInfo,
     ptr::null_mut,
-    sync::atomic::{AtomicU64, Ordering},
+    sync::atomic::{AtomicBool, AtomicU64, Ordering},
 };
 
 use acpi::AcpiTables;
@@ -29,7 +29,7 @@ use sys::{
     acpica_osl::AE_OK,
     early_acpi::EarlyACPIHandler,
     hpet::setup_hpet,
-    lapic::{get_current_lapic_id, setup_lapic_timer, start_aps},
+    lapic::{get_current_lapic_id, setup_lapic_timer, smp_broadcast_panic, start_aps},
     madt::parse_madt,
     pic::disable_pic,
     sync::LOCKS_HELD,
@@ -57,6 +57,7 @@ static LOADER_STRUCT: Mutex<LoaderStruct> = Mutex::new(LoaderStruct {
         stride: 0,
     },
 });
+pub static BROADCASTED_PANIC: AtomicBool = AtomicBool::new(false);
 
 global_asm!(include_str!("cpu/boot.s"), options(att_syntax));
 global_asm!(include_str!("cpu/trampoline.s"), options(att_syntax));
@@ -106,10 +107,13 @@ fn main() {
 #[panic_handler]
 fn panic(info: &PanicInfo) -> ! {
     cli();
+    if !BROADCASTED_PANIC.load(Ordering::SeqCst) {
+        smp_broadcast_panic();
         LOCKS_HELD[get_current_lapic_id()].fetch_add(1, Ordering::SeqCst);
         error!("{}", info);
         let str = format!("{}", info);
         display_print(&str);
+    }
     loop {
         hlt();
     }

kernel/src/sys/lapic.rs

@@ -1,19 +1,25 @@
 use core::{
-    arch::asm,
     ptr::{copy, null_mut},
     sync::atomic::{AtomicBool, AtomicPtr, AtomicUsize, Ordering},
 };
 
 use alloc::{vec, vec::Vec};
-use kernel_common::{instructions::pause, paging::PageTable};
+use kernel_common::{
+    instructions::{get_rflags, pause},
+    paging::PageTable,
+};
 
-use crate::cpu::{
+use crate::{
+    cpu::{
         isr::{ISR_INVALIDATE_TLB, ISR_SCHEDULER},
         paging::{map_physical, map_range, unmap, virt_to_phys, CURRENT_PML4},
+    },
+    BROADCASTED_PANIC,
 };
 
 use super::{
     hpet::sleep,
+    sync::Spinlock,
     task::{ALL_APS_STARTED, STACK_SIZE, STARTING_AP_ID},
 };
 
@@ -37,6 +43,7 @@ const REGISTER_TIMER_LVT: usize = 0xc8;
 const REGISTER_TIMER_INITIAL_COUNT: usize = 0xe0;
 const REGISTER_TIMER_CURRENT_COUNT: usize = 0xe4;
 const REGISTER_TIMER_DIVIDE: usize = 0xf8;
+const IPI_NMI: u32 = 0x400;
 const IPI_INIT: u32 = 0x500;
 const IPI_STARTUP: u32 = 0x600;
 
@@ -50,6 +57,7 @@ static LAPICS: [LAPIC; 256] = [const {
 }; 256];
 static TICKS_PER_MS: AtomicUsize = AtomicUsize::new(0);
 pub static NEXT_LAPIC_ID: AtomicUsize = AtomicUsize::new(0);
+static INVALIDATE_TLB_LOCK: Spinlock = Spinlock::new();
 
 pub fn send_eoi() {
     let address = ADDRESS.load(Ordering::SeqCst);
@@ -60,10 +68,7 @@ pub fn send_eoi() {
 pub fn get_current_lapic_id() -> usize {
     let address = ADDRESS.load(Ordering::SeqCst);
     let lapic_id = unsafe { address.add(REGISTER_ID).read_volatile() as usize >> 24 };
-    let rflags: u64;
-    unsafe {
-        asm!("pushfq; pop {};", out(reg) rflags);
-    }
+    let rflags = get_rflags();
     assert!(rflags & (1 << 9) == 0);
     lapic_id
 }
@@ -164,11 +169,28 @@ pub fn smp_invalidate_tlb() {
     if !ALL_APS_STARTED.load(Ordering::SeqCst) {
         return;
     }
+    INVALIDATE_TLB_LOCK.lock();
+    let current_lapic_id = get_current_lapic_id();
     for i in 0..NEXT_LAPIC_ID.load(Ordering::SeqCst) {
         let lapic_id = LAPICS[i].lapic_id.load(Ordering::SeqCst);
-        if lapic_id == BSP_LAPIC_ID.load(Ordering::SeqCst) {
+        if lapic_id == current_lapic_id {
             continue;
         }
         send_ipi(lapic_id, ISR_INVALIDATE_TLB as u32);
     }
+    INVALIDATE_TLB_LOCK.unlock();
+}
+pub fn smp_broadcast_panic() {
+    BROADCASTED_PANIC.store(true, Ordering::SeqCst);
+    if !ALL_APS_STARTED.load(Ordering::SeqCst) {
+        return;
+    }
+    let current_lapic_id = get_current_lapic_id();
+    for i in 0..NEXT_LAPIC_ID.load(Ordering::SeqCst) {
+        let lapic_id = LAPICS[i].lapic_id.load(Ordering::SeqCst);
+        if lapic_id == current_lapic_id {
+            continue;
+        }
+        send_ipi(lapic_id, IPI_NMI);
+    }
 }

kernel/src/sys/task.rs

@@ -1,11 +1,8 @@
-use core::{
-    arch::asm,
-    sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering},
-};
+use core::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, AtomicUsize, Ordering};
 
 use alloc::{sync::Arc, vec, vec::Vec};
 use hashbrown::HashMap;
-use kernel_common::instructions::{cli, hlt, pause, sti};
+use kernel_common::instructions::{cli, get_rflags, hlt, pause, sti};
 use spin::{Lazy, Mutex};
 
 use crate::{
@@ -156,10 +153,7 @@ fn idle_main() {
     }
 }
 pub fn setup_multitasking() -> ! {
-    let mut rflags;
-    unsafe {
-        asm!("pushf; pop {0:r}", out(reg) rflags);
-    }
+    let rflags = get_rflags();
     RFLAGS.store(rflags, core::sync::atomic::Ordering::SeqCst);
     for _ in 0..NEXT_LAPIC_ID.load(Ordering::SeqCst) {
         create_idle_task();

lib/kernel-common/src/instructions.rs

@@ -20,3 +20,10 @@ pub fn hlt() {
         asm!("hlt");
     }
 }
+pub fn get_rflags() -> u64 {
+    let rflags;
+    unsafe {
+        asm!("pushf; pop {}", out(reg) rflags);
+    }
+    rflags
+}