Always check shootdown reqs when locking addrsp.

src/arch/x86_shared/interrupt/ipi.rs

@@ -10,7 +10,7 @@ interrupt!(wakeup, || {
 });
 
 interrupt!(tlb, || {
-    tlb_shootdown_handler();
+    PercpuBlock::current().maybe_handle_tlb_shootdown();
 
     LOCAL_APIC.eoi();
 });
@@ -27,21 +27,3 @@ 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/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)]

src/debugger.rs

@@ -244,9 +244,9 @@ pub unsafe fn debugger(target_id: Option<crate::context::ContextId>) {
 
         // Switch to context page table to ensure syscall debug and stack dump will work
         if let Some(ref space) = context.addr_space {
-            let was_new = spaces.insert(space.inner.read().table.utable.table().phys().data());
-            RmmA::set_table(TableKind::User, space.inner.read().table.utable.table().phys());
-            check_consistency(&mut space.inner.write(), was_new, &mut tree);
+            let was_new = spaces.insert(space.acquire_read().table.utable.table().phys().data());
+            RmmA::set_table(TableKind::User, space.acquire_read().table.utable.table().phys());
+            check_consistency(&mut space.acquire_write(), was_new, &mut tree);
         }
 
         println!("status: {:?}", context.status);
@@ -260,7 +260,7 @@ pub unsafe fn debugger(target_id: Option<crate::context::ContextId>) {
             );
         }
         if let Some(ref addr_space) = context.addr_space {
-            let addr_space = addr_space.inner.read();
+            let addr_space = addr_space.acquire_read();
             if !addr_space.grants.is_empty() {
                 println!("grants:");
                 for (base, info) in addr_space.grants.iter() {
@@ -285,8 +285,7 @@ pub unsafe fn debugger(target_id: Option<crate::context::ContextId>) {
             for _ in 0..64 {
                 if context.addr_space.as_ref().map_or(false, |space| {
                     space
-                        .inner
-                        .read()
+                        .acquire_read()
                         .table
                         .utable
                         .translate(crate::paging::VirtualAddress::new(rsp))

src/percpu.rs

@@ -50,9 +50,7 @@ pub fn shootdown_tlb_ipi(target: Option<LogicalCpuId>) {
         while percpublock.wants_tlb_shootdown.swap(true, Ordering::Release) == true {
             // Load is faster than CAS or on x86, LOCK BTS
             while percpublock.wants_tlb_shootdown.load(Ordering::Relaxed) == true {
-                if my_percpublock.wants_tlb_shootdown.swap(false, Ordering::Release) == true {
-                    log::warn!("oops");
-                }
+                my_percpublock.maybe_handle_tlb_shootdown();
                 core::hint::spin_loop();
             }
         }
@@ -66,3 +64,22 @@ pub fn shootdown_tlb_ipi(target: Option<LogicalCpuId>) {
         }
     }
 }
+impl PercpuBlock {
+    pub fn maybe_handle_tlb_shootdown(&self) {
+        if self.wants_tlb_shootdown.swap(false, Ordering::Relaxed) == false {
+            return;
+        }
+
+        // TODO: Finer-grained flush
+        unsafe {
+            x86::tlb::flush_all();
+        }
+
+        {
+            let addrsp = self.current_addrsp.borrow();
+            if let Some(ref addrsp) = &*addrsp {
+                addrsp.tlb_ack.fetch_add(1, Ordering::Release);
+            }
+        }
+    }
+}

src/scheme/memory.rs

@@ -83,7 +83,7 @@ impl MemoryScheme {
             return Err(Error::new(EOPNOTSUPP));
         }
 
-        let page = addr_space.inner.write().mmap(
+        let page = addr_space.acquire_write().mmap(
             &addr_space,
             (map.address != 0).then_some(span.base),
             page_count,
@@ -134,8 +134,7 @@ impl MemoryScheme {
 
         let current_addrsp = AddrSpace::current()?;
 
-        let base_page = current_addrsp.inner
-            .write()
+        let base_page = current_addrsp.acquire_write()
             .mmap_anywhere(
                 &current_addrsp,
                 page_count,

src/scheme/proc.rs

@@ -717,7 +717,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
 
                 let requested_dst_base = (map.address != 0).then_some(requested_dst_page);
 
-                let mut src_addr_space = addrspace.inner.write();
+                let mut src_addr_space = addrspace.acquire_write();
 
                 let src_page_count = NonZeroUsize::new(src_span.count).ok_or(Error::new(EINVAL))?;
 
@@ -734,7 +734,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
                         &mut notify_files,
                     )?
                 } else {
-                    let mut dst_addrsp_guard = dst_addr_space.inner.write();
+                    let mut dst_addrsp_guard = dst_addr_space.acquire_write();
                     dst_addrsp_guard.mmap(
                         &dst_addr_space,
                         requested_dst_base,
@@ -901,7 +901,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
 
                 for (dst, (grant_base, grant_info)) in dst
                     .iter_mut()
-                    .zip(addrspace.inner.read().grants.iter().skip(orig_offset))
+                    .zip(addrspace.acquire_read().grants.iter().skip(orig_offset))
                 {
                     *dst = GrantDesc {
                         base: grant_base.start_address().data(),
@@ -981,7 +981,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
                 read_from(buf, &data.buf, &mut data.offset)
             }
             Operation::MmapMinAddr(ref addrspace) => {
-                buf.write_usize(addrspace.inner.read().mmap_min)?;
+                buf.write_usize(addrspace.acquire_read().mmap_min)?;
                 Ok(mem::size_of::<usize>())
             }
             Operation::SchedAffinity => {
@@ -1300,7 +1300,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
                 if val % PAGE_SIZE != 0 || val > crate::USER_END_OFFSET {
                     return Err(Error::new(EINVAL));
                 }
-                addrspace.inner.write().mmap_min = val;
+                addrspace.acquire_write().mmap_min = val;
                 Ok(mem::size_of::<usize>())
             }
             Operation::SchedAffinity => {
@@ -1457,7 +1457,7 @@ impl<const FULL: bool> KernelScheme for ProcScheme<FULL> {
                         let page = Page::containing_address(VirtualAddress::new(page_addr));
 
                         match addrspace
-                            .inner.read()
+                            .acquire_read()
                             .grants
                             .contains(page)
                             .ok_or(Error::new(EINVAL))?

src/scheme/sys/context.rs

@@ -32,7 +32,7 @@ pub fn resource() -> Result<Vec<u8>> {
             // TODO: All user programs must have some grant in order for executable memory to even
             // exist, but is this a good indicator of whether it is user or kernel?
             stat_string.push(if let Ok(addr_space) = context.addr_space() {
-                if addr_space.inner.read().grants.is_empty() {
+                if addr_space.acquire_read().grants.is_empty() {
                     'K'
                 } else {
                     'U'
@@ -84,7 +84,7 @@ pub fn resource() -> Result<Vec<u8>> {
                 memory += kstack.len();
             }
             if let Ok(addr_space) = context.addr_space() {
-                for (_base, info) in addr_space.inner.read().grants.iter() {
+                for (_base, info) in addr_space.acquire_read().grants.iter() {
                     // TODO: method
                     if matches!(info.provider, context::memory::Provider::Allocated { .. }) {
                         memory += info.page_count() * PAGE_SIZE;

src/scheme/user.rs

@@ -242,7 +242,7 @@ impl UserInner {
         tail.buf_mut()[..buf.len()].copy_from_slice(buf);
 
         let is_pinned = true;
-        let dst_page = dst_addr_space.inner.write().mmap_anywhere(
+        let dst_page = dst_addr_space.acquire_write().mmap_anywhere(
             &dst_addr_space,
             ONE,
             PROT_READ,
@@ -343,7 +343,7 @@ impl UserInner {
             .split_at(core::cmp::min(align_offset, user_buf.len()))
             .expect("split must succeed");
 
-        let mut dst_space = dst_space_lock.inner.write();
+        let mut dst_space = dst_space_lock.acquire_write();
 
         let free_span = dst_space
             .grants
@@ -437,7 +437,7 @@ impl UserInner {
 
                     Ok(Grant::borrow(
                         Arc::clone(&cur_space_lock),
-                        &mut *cur_space_lock.inner.write(),
+                        &mut *cur_space_lock.acquire_write(),
                         first_middle_src_page,
                         dst_page,
                         middle_page_count.get(),
@@ -671,8 +671,7 @@ impl UserInner {
                     let context = context.upgrade().ok_or(Error::new(ESRCH))?;
 
                     let (frame, _) = AddrSpace::current()?
-                        .inner
-                        .read()
+                        .acquire_read()
                         .table
                         .utable
                         .translate(base_addr)
@@ -841,7 +840,7 @@ impl UserInner {
                 BorrowedFmapSource {
                     src_base: Page::containing_address(VirtualAddress::new(base_addr)),
                     addr_space_lock,
-                    addr_space_guard: addr_space_lock.inner.write(),
+                    addr_space_guard: addr_space_lock.acquire_write(),
                     mode: if map.flags.contains(MapFlags::MAP_SHARED) {
                         MmapMode::Shared
                     } else {
@@ -854,7 +853,7 @@ impl UserInner {
 
         let page_count_nz = NonZeroUsize::new(page_count).expect("already validated map.size != 0");
         let mut notify_files = Vec::new();
-        let dst_base = dst_addr_space.inner.write().mmap(
+        let dst_base = dst_addr_space.acquire_write().mmap(
             &dst_addr_space,
             dst_base,
             page_count_nz,

src/syscall/driver.rs

@@ -52,7 +52,7 @@ pub fn virttophys(virtual_address: usize) -> Result<usize> {
     enforce_root()?;
 
     let addr_space = Arc::clone(context::current()?.read().addr_space()?);
-    let addr_space = addr_space.inner.read();
+    let addr_space = addr_space.acquire_read();
 
     match addr_space
         .table

src/syscall/futex.rs

@@ -53,7 +53,7 @@ pub fn futex(addr: usize, op: usize, val: usize, val2: usize, addr2: usize) -> R
 
     // Keep the address space locked so we can safely read from the physical address. Unlock it
     // before context switching.
-    let addr_space_guard = addr_space_lock.inner.read();
+    let addr_space_guard = addr_space_lock.acquire_read();
 
     let target_physaddr =
         validate_and_translate_virt(&*addr_space_guard, VirtualAddress::new(addr))

src/syscall/process.rs

@@ -587,7 +587,7 @@ pub unsafe fn usermode_bootstrap(bootstrap: &Bootstrap) -> ! {
         let page_count = NonZeroUsize::new(bootstrap.page_count)
             .expect("bootstrap contained no pages!");
 
-        let _base_page = addr_space.inner.write().mmap(&addr_space, Some(base), page_count, flags, &mut Vec::new(), |page, flags, mapper, flusher| {
+        let _base_page = addr_space.acquire_write().mmap(&addr_space, Some(base), page_count, flags, &mut Vec::new(), |page, flags, mapper, flusher| {
             let shared = false;
             Ok(Grant::zeroed(PageSpan::new(page, bootstrap.page_count), flags, mapper, flusher, shared)?)
         });