Always check shootdown reqs when locking addrsp.

src/context/arch/x86_64.rs

@@ -348,11 +348,11 @@ pub unsafe fn switch_arch_hook() {
         // 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);
+        prev_addrsp.acquire_read().used_by.atomic_clear(percpu.cpu_id);
     }
 
     if let Some(next_addrsp) = &*next_addrsp {
-        let next = next_addrsp.inner.read();
+        let next = next_addrsp.acquire_read();
         next.used_by.atomic_set(percpu.cpu_id);
         next.table.utable.make_current();
     } else {

src/context/context.rs

@@ -408,12 +408,12 @@ impl Context {
 
             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);
+                prev_addrsp.acquire_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();
+            let new_addrsp = addr_space.acquire_read();
             new_addrsp.used_by.atomic_set(this_percpu.cpu_id);
 
             unsafe {

src/context/memory.rs

@@ -2,7 +2,7 @@ use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
 use core::{cmp, fmt::Debug, num::NonZeroUsize, sync::atomic::{Ordering, AtomicU32}};
 use hashbrown::HashMap;
 use rmm::{Arch as _, PageFlush};
-use spin::{RwLock, RwLockUpgradableGuard, RwLockWriteGuard};
+use spin::{RwLock, RwLockUpgradableGuard, RwLockWriteGuard, RwLockReadGuard};
 use syscall::{error::*, flag::MapFlags, GrantFlags, MunmapFlags};
 
 use crate::{
@@ -14,7 +14,7 @@ use crate::{
     paging::{
         Page, PageFlags, PageMapper, RmmA, TableKind, VirtualAddress,
     },
-    scheme::{self, KernelSchemes}, cpu_set::LogicalCpuSet,
+    scheme::{self, KernelSchemes}, cpu_set::LogicalCpuSet, percpu::PercpuBlock,
 };
 
 use super::{context::HardBlockedReason, file::FileDescription};
@@ -75,7 +75,7 @@ impl UnmapResult {
 
 #[derive(Debug)]
 pub struct AddrSpaceWrapper {
-    pub inner: RwLock<AddrSpace>,
+    inner: RwLock<AddrSpace>,
     pub tlb_ack: AtomicU32,
 }
 impl AddrSpaceWrapper {
@@ -85,6 +85,32 @@ impl AddrSpaceWrapper {
             tlb_ack: AtomicU32::new(0),
         }).map_err(|_| Error::new(ENOMEM))
     }
+    pub fn acquire_read(&self) -> RwLockReadGuard<'_, AddrSpace> {
+        let my_percpu = PercpuBlock::current();
+
+        loop {
+            match self.inner.try_read() {
+                Some(g) => return g,
+                None => {
+                    my_percpu.maybe_handle_tlb_shootdown();
+                    core::hint::spin_loop();
+                }
+            }
+        }
+    }
+    pub fn acquire_write(&self) -> RwLockWriteGuard<'_, AddrSpace> {
+        let my_percpu = PercpuBlock::current();
+
+        loop {
+            match self.inner.try_write() {
+                Some(g) => return g,
+                None => {
+                    my_percpu.maybe_handle_tlb_shootdown();
+                    core::hint::spin_loop();
+                }
+            }
+        }
+    }
 }
 
 #[derive(Debug)]