From 716f147cd4edc6033818a0b580892f6f88e16f46 Mon Sep 17 00:00:00 2001
From: 4lDO2 <4lDO2@protonmail.com>
Date: Mon, 4 Mar 2024 16:35:54 +0100
Subject: [PATCH] Fix synchronization so `acid tlb` succeeds!

---
 src/arch/x86_64/gdt.rs     |  2 +-
 src/context/arch/x86_64.rs | 18 ++++++++----------
 src/context/switch.rs      |  1 +
 src/percpu.rs              | 13 +++----------
 4 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/src/arch/x86_64/gdt.rs b/src/arch/x86_64/gdt.rs
index f3b0e85cc0..15b4c662ca 100644
--- a/src/arch/x86_64/gdt.rs
+++ b/src/arch/x86_64/gdt.rs
@@ -241,7 +241,7 @@ pub unsafe fn init_paging(stack_offset: usize, cpu_id: LogicalCpuId) {
         cpu_id,
         switch_internals: Default::default(),
         current_addrsp: RefCell::new(None),
-        old_addrsp_tmp: RefCell::new(None),
+        new_addrsp_tmp: Cell::new(None),
         wants_tlb_shootdown: AtomicBool::new(false),
 
         #[cfg(feature = "profiling")]
diff --git a/src/context/arch/x86_64.rs b/src/context/arch/x86_64.rs
index 164988cdfa..1ad7968581 100644
--- a/src/context/arch/x86_64.rs
+++ b/src/context/arch/x86_64.rs
@@ -223,9 +223,7 @@ pub unsafe fn switch_to(prev: &mut super::Context, next: &mut super::Context) {
         );
     }
 
-    let percpu = PercpuBlock::current();
-    *percpu.current_addrsp.borrow_mut() = next.addr_space.clone();
-    *percpu.old_addrsp_tmp.borrow_mut() = prev.addr_space.clone();
+    PercpuBlock::current().new_addrsp_tmp.set(next.addr_space.clone());
 
     switch_to_inner(&mut prev.arch, &mut next.arch)
 }
@@ -336,10 +334,10 @@ unsafe extern "C" fn signal_handler_wrapper() {
 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 cur_addrsp = percpu.current_addrsp.borrow();
+    let next_addrsp = percpu.new_addrsp_tmp.take();
 
-    let retain_pgtbl = match (&*prev_addrsp, &*next_addrsp) {
+    let retain_pgtbl = match (&*cur_addrsp, &next_addrsp) {
         (Some(ref p), Some(ref n)) => Arc::ptr_eq(p, n),
         (Some(_), None) | (None, Some(_)) => false,
         (None, None) => true,
@@ -347,17 +345,17 @@ pub unsafe fn switch_arch_hook() {
     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 {
+    if let Some(ref prev_addrsp) = &*cur_addrsp {
         prev_addrsp.acquire_read().used_by.atomic_clear(percpu.cpu_id);
     }
 
-    drop(prev_addrsp);
+    drop(cur_addrsp);
 
     // Tell future TLB shootdown handlers that old_addrsp_tmp is no longer the current address
     // space.
-    *percpu.old_addrsp_tmp.borrow_mut() = None;
+    *percpu.current_addrsp.borrow_mut() = next_addrsp;
 
-    if let Some(next_addrsp) = &*next_addrsp {
+    if let Some(next_addrsp) = &*percpu.current_addrsp.borrow() {
         let next = next_addrsp.acquire_read();
 
         next.used_by.atomic_set(percpu.cpu_id);
diff --git a/src/context/switch.rs b/src/context/switch.rs
index 9f2e3240f5..340cc477fe 100644
--- a/src/context/switch.rs
+++ b/src/context/switch.rs
@@ -133,6 +133,7 @@ pub unsafe fn switch() -> bool {
         .is_err()
     {
         interrupt::pause();
+        percpu.maybe_handle_tlb_shootdown();
     }
 
     let cpu_id = crate::cpu_id();
diff --git a/src/percpu.rs b/src/percpu.rs
index 8c915f0842..1df0e1fff8 100644
--- a/src/percpu.rs
+++ b/src/percpu.rs
@@ -16,7 +16,7 @@ pub struct PercpuBlock {
     pub switch_internals: ContextSwitchPercpu,
 
     pub current_addrsp: RefCell<Option<Arc<AddrSpaceWrapper>>>,
-    pub old_addrsp_tmp: RefCell<Option<Arc<AddrSpaceWrapper>>>,
+    pub new_addrsp_tmp: Cell<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
@@ -75,15 +75,8 @@ impl PercpuBlock {
             x86::tlb::flush_all();
         }
 
-        {
-            let old = self.old_addrsp_tmp.borrow();
-            let addrsp = self.current_addrsp.borrow();
-
-            if let Some(ref old) = &*old {
-                old.tlb_ack.fetch_add(1, Ordering::Release);
-            } else if let Some(ref addrsp) = &*addrsp {
-                addrsp.tlb_ack.fetch_add(1, Ordering::Release);
-            }
+        if let Some(ref addrsp) = &*self.current_addrsp.borrow() {
+            addrsp.tlb_ack.fetch_add(1, Ordering::Release);
         }
     }
 }