Solve some lock ordering

src/context/memory.rs

@@ -317,7 +317,6 @@ impl AddrSpaceWrapper {
         unpin: bool,
         token: &mut CleanLockToken,
     ) -> Result<Vec<UnmapResult>> {
-        let mut token = token.token();
         let mut guard = self.acquire_write(token.downgrade());
         let guard = &mut *guard;
 
@@ -666,7 +665,9 @@ impl AddrSpace {
             }
 
             // Remove irrelevant region
-            let unmap_result = grant.unmap(this_mapper, this_flusher);
+            // TODO: Lock ordering violation
+            let mut token = unsafe { CleanLockToken::new() };
+            let unmap_result = grant.unmap(this_mapper, this_flusher, &mut token);
 
             // Notify scheme that holds grant
             if unmap_result.file_desc.is_some() {
@@ -765,7 +766,7 @@ impl AddrSpace {
             // longer arc-rwlock wrapped, it cannot be referenced `External`ly by borrowing grants,
             // so it should suffice to iterate over PageInfos and decrement and maybe deallocate
             // the underlying pages (and send some funmaps).
-            let res = { grant.unmap(&mut self.table.utable, &mut NopFlusher) };
+            let res = { grant.unmap(&mut self.table.utable, &mut NopFlusher, token) };
 
             let _ = res.unmap(token);
         }
@@ -1964,6 +1965,7 @@ impl Grant {
         mut self,
         mapper: &mut PageMapper,
         flusher: &mut impl GenericFlusher,
+        token: &mut CleanLockToken,
     ) -> UnmapResult {
         assert!(self.info.mapped);
         assert!(!self.info.is_pinned());
@@ -1974,9 +1976,6 @@ impl Grant {
             ..
         } = self.info.provider
         {
-            // TODO: Lock ordering violation
-            let mut token = unsafe { CleanLockToken::new() };
-            let mut token = token.token();
             let mut guard = address_space.acquire_write(token.downgrade());
 
             for (_, grant) in guard

src/scheme/proc.rs

@@ -441,7 +441,7 @@ impl KernelScheme for ProcScheme {
                         arg1,
                     },
             } => {
-                let old_ctx = try_stop_context(context, token, |context, _| {
+                let old_ctx = try_stop_context(context, token, |context, token| {
                     let regs = context.regs_mut().ok_or(Error::new(EBADFD))?;
                     regs.set_instr_pointer(new_ip);
                     regs.set_stack_pointer(new_sp);
@@ -452,9 +452,7 @@ impl KernelScheme for ProcScheme {
                     ))]
                     regs.set_arg1(arg1);
 
-                    // TODO: Lock ordering violation
-                    let mut token = unsafe { CleanLockToken::new() };
-                    Ok(context.set_addr_space(Some(new), token.downgrade()))
+                    Ok(context.set_addr_space(Some(new), token))
                 })?;
                 if let Some(old_ctx) = old_ctx
                     && let Some(addrspace) = Arc::into_inner(old_ctx)

src/scheme/user.rs

@@ -1340,12 +1340,15 @@ impl<const READ: bool, const WRITE: bool> CaptureGuard<READ, WRITE> {
             dst.copy_from_slice(&src.buf()[..dst.len()])?;
         }
         let unpin = true;
-        if let Some(ref addrsp) = self.addrsp
-            && !self.span.is_empty()
-        {
-            let res = addrsp.munmap(self.span, unpin, token)?;
-            for r in res {
-                let _ = r.unmap(token);
+        if let Some(addrsp) = self.addrsp.take() {
+            if !self.span.is_empty() {
+                let res = addrsp.munmap(self.span, unpin, token)?;
+                for r in res {
+                    let _ = r.unmap(token);
+                }
+            }
+            if let Some(addrsp) = Arc::into_inner(addrsp) {
+                addrsp.into_drop(token);
             }
         }
 
@@ -1353,11 +1356,6 @@ impl<const READ: bool, const WRITE: bool> CaptureGuard<READ, WRITE> {
     }
     pub fn release(mut self, token: &mut CleanLockToken) -> Result<()> {
         self.release_inner(token)?;
-        if let Some(addrsp) = self.addrsp.take()
-            && let Some(addrsp) = Arc::into_inner(addrsp)
-        {
-            addrsp.into_drop(token);
-        }
         if let Some(src) = self.head.src.take() {
             src.into_drop(token);
         }

src/syscall/process.rs

@@ -36,13 +36,11 @@ pub fn exit_this_context(excp: Option<syscall::Exception>, token: &mut CleanLock
     let context_lock = context::current();
     {
         let mut context = context_lock.write(token.token());
-        // let (context, mut token) = context.token_split();
+        let (context, mut token) = context.token_split();
         close_files = Arc::try_unwrap(mem::take(&mut context.files))
             .map_or_else(|_| FdTbl::new(), RwLock::into_inner);
-        // TODO: Lock ordering violation
-        let mut token = unsafe { CleanLockToken::new() };
         addrspace_opt = context
-            .set_addr_space(None, token.downgrade())
+            .set_addr_space(None, token)
             .and_then(|a| Arc::try_unwrap(a).ok());
         drop(mem::replace(&mut context.syscall_head, SyscallFrame::Dummy));
         drop(mem::replace(&mut context.syscall_tail, SyscallFrame::Dummy));