From bb30530ea09448a7e9b1b23adec7e95128dfd99b Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Fri, 9 Feb 2024 14:15:20 +0100
Subject: [PATCH] Track which CPUs are using any given AddrSpace.

---
 src/context/arch/x86_64.rs | 27 ++++++++++++++++++---------
 src/context/memory.rs      |  4 +++-
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs
index 8d5f9d3827..e6bdc4966d 100644
--- a/src/context/arch/x86_64.rs
+++ b/src/context/arch/x86_64.rs
@@ -221,6 +221,8 @@ 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
@@ -231,18 +233,25 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
                 .as_ref()
                 .map_or(true, |prev_space| !Arc::ptr_eq(prev_space, next_space))
             {
-                // Suppose we have two sibling threads A and B. A runs on CPU 0 and B on CPU 1. A
-                // recently called yield and is now here about to switch back. Meanwhile, B is
-                // currently creating a new mapping in their shared address space, for example a
-                // message on a channel.
-                //
-                // Unless we acquire this lock, it may be possible that the TLB will not contain new
-                // entries. While this can be caught and corrected in a page fault handler, this is not
-                // true when entries are removed from a page table!
-                next_space.read().table.utable.make_current();
+                if let Some(ref prev_space) = prev.addr_space {
+                    prev_space.read().used_by.atomic_clear(this_cpu);
+                }
+                // 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.read();
+
+                next_space.used_by.atomic_set(this_cpu);
+                next_space.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.read().used_by.atomic_clear(this_cpu);
+            }
             RmmA::set_table(TableKind::User, empty_cr3());
         }
     }
diff --git a/src/context/memory.rs b/src/context/memory.rs
index 93035af01b..22795dc840 100644
--- a/src/context/memory.rs
+++ b/src/context/memory.rs
@@ -15,7 +15,7 @@ use crate::{
         mapper::{Flusher, InactiveFlusher, PageFlushAll},
         Page, PageFlags, PageMapper, RmmA, TableKind, VirtualAddress,
     },
-    scheme::{self, KernelSchemes},
+    scheme::{self, KernelSchemes}, cpu_set::LogicalCpuSet,
 };
 
 use super::{context::HardBlockedReason, file::FileDescription};
@@ -82,6 +82,7 @@ pub fn new_addrspace() -> Result<Arc<RwLock<AddrSpace>>> {
 pub struct AddrSpace {
     pub table: Table,
     pub grants: UserGrants,
+    pub used_by: LogicalCpuSet,
     /// Lowest offset for mmap invocations where the user has not already specified the offset
     /// (using MAP_FIXED/MAP_FIXED_NOREPLACE). Cf. Linux's `/proc/sys/vm/mmap_min_addr`, but with
     /// the exception that we have a memory safe kernel which doesn't have to protect itself
@@ -188,6 +189,7 @@ impl AddrSpace {
             grants: UserGrants::new(),
             table: setup_new_utable()?,
             mmap_min: MMAP_MIN_DEFAULT,
+            used_by: LogicalCpuSet::empty(),
         })
     }
     pub fn is_current(&self) -> bool {