Fix more synchronization.

src/context/context.rs

@@ -11,7 +11,7 @@ use crate::{
     memory::{Frame, RaiiFrame},
     paging::{RmmA, RmmArch},
     scheme::{CallerCtx, FileHandle, SchemeNamespace},
-    sync::WaitMap,
+    sync::WaitMap, percpu::PercpuBlock,
 };
 
 use crate::syscall::{
@@ -399,9 +399,26 @@ impl Context {
         &mut self,
         addr_space: Arc<AddrSpaceWrapper>,
     ) -> Option<Arc<AddrSpaceWrapper>> {
-        if self.id == super::context_id() {
+        if let Some(ref addrsp) = self.addr_space && Arc::ptr_eq(addrsp, &addr_space) {
+            return Some(addr_space);
+        }
+
+        let current_context_id = super::context_id();
+        let current_cpu_id = crate::cpu_id();
+
+        if self.id == current_context_id {
             unsafe {
-                addr_space.inner.read().table.utable.make_current();
+                let guard = addr_space.inner.read();
+
+                let prev_addrsp = self.addr_space.as_ref().map(|asp| asp.inner.read());
+
+                if let Some(prev_addrsp) = prev_addrsp.as_deref() {
+                    prev_addrsp.used_by.atomic_clear(current_cpu_id);
+                }
+
+                guard.used_by.atomic_set(current_cpu_id);
+                PercpuBlock::current().current_addrspace.set(Arc::as_ptr(&addr_space));
+                guard.table.utable.make_current();
             }
         }