Fix more synchronization.

src/arch/x86_64/gdt.rs

@@ -245,6 +245,8 @@ pub unsafe fn init_paging(stack_offset: usize, cpu_id: LogicalCpuId) {
         #[cfg(feature = "profiling")]
         profiling: None,
     };
+
+    crate::percpu::init_tlb_shootdown(cpu_id, &mut pcr.percpu);
 }
 #[derive(Copy, Clone, Debug)]
 #[repr(packed)]

src/arch/x86_shared/device/local_apic.rs

@@ -149,7 +149,6 @@ impl LocalApic {
         }
         self.set_icr(icr);
     }
-    // Not used just yet, but allows triggering an NMI to another processor.
     pub fn ipi_nmi(&mut self, apic_id: u32) {
         let shift = if self.x2 { 32 } else { 56 };
         self.set_icr((u64::from(apic_id) << shift) | (1 << 14) | (0b100 << 8));

src/context/arch/x86_64.rs

@@ -231,31 +231,38 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
         // 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) => {
-            percpu_block.current_addrspace.set(Arc::as_ptr(next_space));
-
             if prev
                 .addr_space
                 .as_ref()
                 .map_or(true, |prev_space| !Arc::ptr_eq(prev_space, next_space))
             {
-                if let Some(ref prev_space) = prev.addr_space {
-                    prev_space.inner.read().used_by.atomic_clear(this_cpu);
-                }
+                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 = next_space.inner.read();
+                let next_space_guard = next_space.inner.read();
 
-                next_space.used_by.atomic_set(this_cpu);
-                next_space.table.utable.make_current();
+                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 {
-                prev_space.inner.read().used_by.atomic_clear(this_cpu);
+                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());
         }

src/context/context.rs

@@ -11,7 +11,7 @@ use crate::{
     memory::{Frame, RaiiFrame},
     paging::{RmmA, RmmArch},
     scheme::{CallerCtx, FileHandle, SchemeNamespace},
-    sync::WaitMap,
+    sync::WaitMap, percpu::PercpuBlock,
 };
 
 use crate::syscall::{
@@ -399,9 +399,26 @@ impl Context {
         &mut self,
         addr_space: Arc<AddrSpaceWrapper>,
     ) -> Option<Arc<AddrSpaceWrapper>> {
-        if self.id == super::context_id() {
+        if let Some(ref addrsp) = self.addr_space && Arc::ptr_eq(addrsp, &addr_space) {
+            return Some(addr_space);
+        }
+
+        let current_context_id = super::context_id();
+        let current_cpu_id = crate::cpu_id();
+
+        if self.id == current_context_id {
             unsafe {
-                addr_space.inner.read().table.utable.make_current();
+                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();
             }
         }
 

src/percpu.rs

@@ -31,6 +31,10 @@ pub struct PercpuBlock {
 const NULL: AtomicPtr<PercpuBlock> = AtomicPtr::new(core::ptr::null_mut());
 static ALL_PERCPU_BLOCKS: [AtomicPtr<PercpuBlock>; MAX_CPU_COUNT as usize] = [NULL; MAX_CPU_COUNT as usize];
 
+pub unsafe fn init_tlb_shootdown(id: LogicalCpuId, block: *mut PercpuBlock) {
+    ALL_PERCPU_BLOCKS[id.get() as usize].store(block, Ordering::Release)
+}
+
 // PercpuBlock::current() is implemented somewhere in the arch-specific modules
 
 bitflags::bitflags! {
@@ -52,6 +56,7 @@ pub fn shootdown_tlb_ipi(_target: Option<LogicalCpuId>) {}
 pub fn shootdown_tlb_ipi(target: Option<LogicalCpuId>) {
     if let Some(target) = target {
         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();