Use regular non-NMI IPIs for TLB shootdown.

src/arch/x86_64/gdt.rs

@@ -1,6 +1,7 @@
 //! Global descriptor table
 
-use core::{convert::TryInto, mem, sync::atomic::AtomicUsize, cell::Cell};
+use core::{convert::TryInto, mem, cell::{Cell, RefCell}};
+use core::sync::atomic::AtomicBool;
 
 use crate::{
     cpu_set::LogicalCpuId,
@@ -239,8 +240,9 @@ pub unsafe fn init_paging(stack_offset: usize, cpu_id: LogicalCpuId) {
     pcr.percpu = PercpuBlock {
         cpu_id,
         switch_internals: Default::default(),
-        nmi_flags_lock: AtomicUsize::new(0),
-        current_addrspace: Cell::new(core::ptr::null_mut()),
+        current_addrsp: RefCell::new(None),
+        old_addrsp_tmp: RefCell::new(None),
+        wants_tlb_shootdown: AtomicBool::new(false),
 
         #[cfg(feature = "profiling")]
         profiling: None,

src/arch/x86_64/interrupt/exception.rs

@@ -1,8 +1,5 @@
-use core::sync::atomic::Ordering;
-
 use x86::irq::PageFaultError;
 
-use crate::percpu::{PercpuBlock, NmiReasons};
 use crate::{
     interrupt::stack_trace, interrupt_error, interrupt_stack, memory::GenericPfFlags,
     paging::VirtualAddress, ptrace, syscall::flag::*,
@@ -43,27 +40,14 @@ interrupt_stack!(debug, @paranoid, |stack| {
 });
 
 interrupt_stack!(non_maskable, @paranoid, |stack| {
-    let percpu = PercpuBlock::current();
+    #[cfg(feature = "profiling")]
+    crate::profiling::nmi_handler(stack);
 
-    let reasons = NmiReasons::from_bits_retain(percpu.nmi_flags_lock.load(Ordering::Relaxed));
-
-    if let Some(reason) = reasons.iter().next() {
-        if reason == NmiReasons::TLB_SHOOTDOWN && let Some(addrsp) = percpu.current_addrspace.get().as_ref() {
-            rmm::PageFlushAll::<crate::paging::RmmA>::new().flush();
-            addrsp.tlb_ack.fetch_add(1, Ordering::Relaxed);
-        }
-
-        percpu.nmi_flags_lock.fetch_and(!reason.bits(), Ordering::Relaxed);
-    } else {
-        #[cfg(feature = "profiling")]
-        crate::profiling::nmi_handler(stack);
-
-        #[cfg(not(feature = "profiling"))]
-        {
-            // TODO: This will likely deadlock
-            println!("Non-maskable interrupt");
-            stack.dump();
-        }
+    #[cfg(not(feature = "profiling"))]
+    {
+        // TODO: This will likely deadlock
+        println!("Non-maskable interrupt");
+        stack.dump();
     }
 });
 

src/arch/x86_64/mod.rs

@@ -70,10 +70,3 @@ pub unsafe extern "C" fn arch_copy_to_user(dst: usize, src: usize, len: usize) -
 pub use arch_copy_to_user as arch_copy_from_user;
 
 pub use alternative::kfx_size;
-
-pub fn send_tlb_nmi(target: LogicalCpuId) {
-    // TODO: Logical to APIC ID translation
-    unsafe {
-        device::local_apic::LOCAL_APIC.ipi_nmi(target.get());
-    }
-}

src/arch/x86_shared/device/local_apic.rs

@@ -4,7 +4,7 @@ use core::{
 };
 use x86::msr::*;
 
-use crate::paging::{KernelMapper, PageFlags, PhysicalAddress, RmmA, RmmArch, VirtualAddress};
+use crate::{paging::{KernelMapper, PageFlags, PhysicalAddress, RmmA, RmmArch, VirtualAddress}, ipi::IpiKind};
 
 use crate::arch::cpuid::cpuid;
 
@@ -140,12 +140,12 @@ impl LocalApic {
         }
     }
 
-    pub fn ipi(&mut self, apic_id: usize) {
-        let mut icr = 0x4040;
+    pub fn ipi(&mut self, apic_id: u32, kind: IpiKind) {
+        let mut icr = 0x40 | kind as u64;
         if self.x2 {
-            icr |= (apic_id as u64) << 32;
+            icr |= u64::from(apic_id) << 32;
         } else {
-            icr |= (apic_id as u64) << 56;
+            icr |= u64::from(apic_id) << 56;
         }
         self.set_icr(icr);
     }

src/arch/x86_shared/interrupt/ipi.rs

@@ -1,5 +1,8 @@
+use core::sync::atomic::Ordering;
+
 use x86::tlb;
 
+use crate::percpu::PercpuBlock;
 use crate::{context, device::local_apic::LOCAL_APIC};
 
 interrupt!(wakeup, || {
@@ -7,9 +10,9 @@ interrupt!(wakeup, || {
 });
 
 interrupt!(tlb, || {
-    LOCAL_APIC.eoi();
+    tlb_shootdown_handler();
 
-    tlb::flush_all();
+    LOCAL_APIC.eoi();
 });
 
 interrupt!(switch, || {
@@ -24,3 +27,21 @@ interrupt!(pit, || {
     // Switch after a sufficient amount of time since the last switch.
     context::switch::tick();
 });
+
+unsafe fn tlb_shootdown_handler() {
+    let pcpu = PercpuBlock::current();
+
+    if pcpu.wants_tlb_shootdown.swap(false, Ordering::Relaxed) == false {
+        // Spurious TLB IPI, could have been manually triggered after the IPI was sent.
+        return;
+    }
+
+    tlb::flush_all();
+
+    {
+        let addrsp = pcpu.current_addrsp.borrow();
+        if let Some(ref addrsp) = &*addrsp {
+            addrsp.tlb_ack.fetch_add(1, Ordering::Release);
+        }
+    }
+}

src/arch/x86_shared/ipi.rs

@@ -37,3 +37,19 @@ pub fn ipi(kind: IpiKind, target: IpiTarget) {
     let icr = (target as u64) << 18 | 1 << 14 | (kind as u64);
     unsafe { LOCAL_APIC.set_icr(icr) };
 }
+use crate::cpu_set::LogicalCpuId;
+
+#[cfg(feature = "multi_core")]
+#[inline(always)]
+pub fn ipi_single(kind: IpiKind, target: LogicalCpuId) {
+    use crate::device::local_apic::LOCAL_APIC;
+
+    unsafe {
+        // TODO: Distinguish between logical and physical CPU IDs
+        LOCAL_APIC.ipi(target.get(), kind);
+    }
+}
+
+#[cfg(not(feature = "multi_core"))]
+#[inline(always)]
+pub fn ipi_single(_kind: IpiKind, _target: LogicalCpuId) {}

src/context/arch/x86_64.rs

@@ -10,7 +10,7 @@ use crate::{
     interrupt::handler::ScratchRegisters,
     paging::{RmmA, RmmArch, TableKind},
     pop_scratch, push_scratch,
-    syscall::FloatRegisters,
+    syscall::FloatRegisters, percpu::PercpuBlock,
 };
 
 use core::mem::offset_of;
@@ -148,7 +148,6 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
     if let Some(ref stack) = next.kstack {
         crate::gdt::set_tss_stack(pcr, stack.as_ptr() as usize + stack.len());
     }
-    let percpu_block = &*core::ptr::addr_of!((*pcr).percpu);
 
     core::arch::asm!(
         alternative2!(
@@ -224,49 +223,9 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
         );
     }
 
-    let this_cpu = crate::cpu_id();
-
-    match next.addr_space {
-        // Since Arc essentially just wraps a pointer, in this case a regular pointer (as opposed
-        // to dyn or slice fat pointers), and NonNull optimization exists, map_or will hopefully be
-        // optimized down to checking prev and next pointers, as next cannot be null.
-        Some(ref next_space) => {
-            if prev
-                .addr_space
-                .as_ref()
-                .map_or(true, |prev_space| !Arc::ptr_eq(prev_space, next_space))
-            {
-                let prev_space_guard = prev.addr_space.as_ref().map(|asp| asp.inner.read());
-                // This lock needs to be held, because if the address space is write-locked by some
-                // context that is e.g. unmapping memory, we either need to switch after its
-                // changes have been made, or it needs to know this context is potentially using
-                // this address space, at that time.
-                let next_space_guard = next_space.inner.read();
-
-                if let Some(prev_space_guard) = prev_space_guard.as_deref() {
-                    prev_space_guard.used_by.atomic_clear(this_cpu);
-                }
-
-                percpu_block.current_addrspace.set(Arc::as_ptr(&next_space));
-                next_space_guard.used_by.atomic_set(this_cpu);
-                next_space_guard.table.utable.make_current();
-            }
-        }
-        None => {
-            // The next context is kernel-only, so switch to the page table without any user
-            // mappings.
-            if let Some(ref prev_space) = prev.addr_space {
-                let prev_space = prev_space.inner.read();
-                prev_space.used_by.atomic_clear(this_cpu);
-
-                // Do this while the lock is held, to prevent deadlocks.
-                percpu_block.current_addrspace.set(core::ptr::null());
-            } else {
-                percpu_block.current_addrspace.set(core::ptr::null());
-            }
-            RmmA::set_table(TableKind::User, empty_cr3());
-        }
-    }
+    let percpu = PercpuBlock::current();
+    *percpu.current_addrsp.borrow_mut() = next.addr_space.clone();
+    *percpu.old_addrsp_tmp.borrow_mut() = prev.addr_space.clone();
 
     switch_to_inner(&mut prev.arch, &mut next.arch)
 }
@@ -374,3 +333,29 @@ unsafe extern "C" fn signal_handler_wrapper() {
         options(noreturn)
     );
 }
+pub unsafe fn switch_arch_hook() {
+    let percpu = PercpuBlock::current();
+
+    let prev_addrsp = percpu.old_addrsp_tmp.borrow();
+    let next_addrsp = percpu.current_addrsp.borrow();
+
+    let retain_pgtbl = match (&*prev_addrsp, &*next_addrsp) {
+        (Some(ref p), Some(ref n)) => Arc::ptr_eq(p, n),
+        (Some(_), None) | (None, Some(_)) => false,
+        (None, None) => true,
+    };
+    if retain_pgtbl {
+        // If we are not switching to a different address space, we can simply return early.
+    }
+    if let Some(ref prev_addrsp) = &*prev_addrsp {
+        prev_addrsp.inner.read().used_by.atomic_clear(percpu.cpu_id);
+    }
+
+    if let Some(next_addrsp) = &*next_addrsp {
+        let next = next_addrsp.inner.read();
+        next.used_by.atomic_set(percpu.cpu_id);
+        next.table.utable.make_current();
+    } else {
+        RmmA::set_table(TableKind::User, empty_cr3());
+    }
+}

src/context/context.rs

@@ -399,26 +399,25 @@ impl Context {
         &mut self,
         addr_space: Arc<AddrSpaceWrapper>,
     ) -> Option<Arc<AddrSpaceWrapper>> {
-        if let Some(ref addrsp) = self.addr_space && Arc::ptr_eq(addrsp, &addr_space) {
+        if let Some(ref old) = self.addr_space && Arc::ptr_eq(old, &addr_space) {
             return Some(addr_space);
-        }
+        };
 
-        let current_context_id = super::context_id();
-        let current_cpu_id = crate::cpu_id();
+        if self.id == super::context_id() {
+            let this_percpu = PercpuBlock::current();
+
+            if let Some(ref prev_addrsp) = self.addr_space {
+                assert!(Arc::ptr_eq(&this_percpu.current_addrsp.borrow().as_ref().unwrap(), prev_addrsp));
+                prev_addrsp.inner.read().used_by.atomic_clear(this_percpu.cpu_id);
+            }
+
+            *this_percpu.current_addrsp.borrow_mut() = Some(Arc::clone(&addr_space));
+
+            let new_addrsp = addr_space.inner.read();
+            new_addrsp.used_by.atomic_set(this_percpu.cpu_id);
 
-        if self.id == current_context_id {
             unsafe {
-                let guard = addr_space.inner.read();
-
-                let prev_addrsp = self.addr_space.as_ref().map(|asp| asp.inner.read());
-
-                if let Some(prev_addrsp) = prev_addrsp.as_deref() {
-                    prev_addrsp.used_by.atomic_clear(current_cpu_id);
-                }
-
-                guard.used_by.atomic_set(current_cpu_id);
-                PercpuBlock::current().current_addrspace.set(Arc::as_ptr(&addr_space));
-                guard.table.utable.make_current();
+                new_addrsp.table.utable.make_current();
             }
         }
 

src/context/list.rs

@@ -4,7 +4,7 @@ use core::{iter, mem};
 use spinning_top::RwSpinlock;
 
 use super::context::{Context, ContextId};
-use super::memory::{AddrSpaceWrapper, AddrSpace};
+use super::memory::AddrSpaceWrapper;
 use crate::syscall::error::{Error, Result, EAGAIN};
 
 /// Context list type

src/context/switch.rs

@@ -111,6 +111,7 @@ pub unsafe extern "C" fn switch_finish_hook() {
         // TODO: unreachable_unchecked()?
         crate::arch::stop::emergency_reset();
     }
+    super::arch::switch_arch_hook();
     arch::CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst);
 }
 

src/percpu.rs

@@ -1,5 +1,7 @@
-use core::cell::Cell;
-use core::sync::atomic::{AtomicUsize, AtomicPtr, Ordering};
+use core::cell::{Cell, RefCell};
+use core::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
+
+use alloc::sync::Arc;
 
 use crate::context::memory::AddrSpaceWrapper;
 use crate::cpu_set::MAX_CPU_COUNT;
@@ -13,14 +15,9 @@ pub struct PercpuBlock {
     /// Context management
     pub switch_internals: ContextSwitchPercpu,
 
-    // TODO: This lock can probably be relaxed further, but verify correctness first.
-
-    // The NMI lock. Can be set by any CPU, but can only be cleared by the CPU this percpu block
-    // refers to. Multiple flags can be set simultaneously, but NMI senders need to wait for the
-    // flag to be cleared if it was already set.
-    pub nmi_flags_lock: AtomicUsize,
-
-    pub current_addrspace: Cell<*const AddrSpaceWrapper>,
+    pub current_addrsp: RefCell<Option<Arc<AddrSpaceWrapper>>>,
+    pub old_addrsp_tmp: RefCell<Option<Arc<AddrSpaceWrapper>>>,
+    pub wants_tlb_shootdown: AtomicBool,
 
     // TODO: Put mailbox queues here, e.g. for TLB shootdown? Just be sure to 128-byte align it
     // first to avoid cache invalidation.
@@ -37,41 +34,34 @@ pub unsafe fn init_tlb_shootdown(id: LogicalCpuId, block: *mut PercpuBlock) {
 
 // PercpuBlock::current() is implemented somewhere in the arch-specific modules
 
-bitflags::bitflags! {
-    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
-    pub struct NmiReasons: usize {
-        const TLB_SHOOTDOWN = 1;
-
-        // TODO: Profiling code wakes all CPUs up, so use a global and percpu ack counter for that.
-        //
-        //#[cfg(feature = "profiling")]
-        //const PROFILING = 1 << 32;
-    }
-}
-
 #[cfg(not(feature = "multi_core"))]
 pub fn shootdown_tlb_ipi(_target: Option<LogicalCpuId>) {}
 
 #[cfg(feature = "multi_core")]
 pub fn shootdown_tlb_ipi(target: Option<LogicalCpuId>) {
     if let Some(target) = target {
+        let my_percpublock = PercpuBlock::current();
+        assert_ne!(target, my_percpublock.cpu_id);
+
         let Some(percpublock) = (unsafe { ALL_PERCPU_BLOCKS[target.get() as usize].load(Ordering::Acquire).as_ref() }) else {
             log::warn!("Trying to TLB shootdown a CPU that doesn't exist or isn't initialized.");
             return;
         };
-        let bit = NmiReasons::TLB_SHOOTDOWN.bits();
-
-        while percpublock.nmi_flags_lock.fetch_or(bit, Ordering::Release) & bit == bit {
+        while percpublock.wants_tlb_shootdown.swap(true, Ordering::Release) == true {
             // Load is faster than CAS or on x86, LOCK BTS
-            while percpublock.nmi_flags_lock.load(Ordering::Relaxed) & bit == bit {
+            while percpublock.wants_tlb_shootdown.load(Ordering::Relaxed) == true {
+                if my_percpublock.wants_tlb_shootdown.swap(false, Ordering::Release) == true {
+                    log::warn!("oops");
+                }
                 core::hint::spin_loop();
             }
         }
 
-        crate::arch::send_tlb_nmi(target);
+        crate::ipi::ipi_single(crate::ipi::IpiKind::Tlb, target);
     } else {
         for id in 0..crate::cpu_count() {
-            // TODO: Optimize: use global counter and percpu ack counters.
+            // TODO: Optimize: use global counter and percpu ack counters, send IPI using
+            // destination shorthand "all CPUs".
             shootdown_tlb_ipi(Some(LogicalCpuId::new(id)));
         }
     }