Fix synchronization so acid tlb succeeds!

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")]

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);

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();

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);
         }
     }
 }