misc(scheme/user): use preempt guard

Signed-off-by: Anhad Singh <andypython@protonmail.com>

src/context/context.rs

@@ -144,7 +144,8 @@ pub struct Context {
     pub egid: u32,
     pub pid: usize,
 
-    pub is_preemptable: bool,
+    // Use PreemptGuard
+    pub(super) is_preemptable: bool,
 }
 
 #[derive(Debug)]

src/context/mod.rs

@@ -174,3 +174,37 @@ pub fn spawn(
     }
     Ok(context_lock)
 }
+
+/// A guard that disables preemption for a context while it is alive.
+///
+/// This guard is used to ensure that a sequence of operations is atomic with respect to preemption.
+/// It automatically re-enables preemption when dropped.
+///
+/// Because the guard must hold a mutable reference to the `CleanLockToken` to re-enable preemption
+/// in `Drop`, it consumes the token. The `token()` method allows re-borrowing the token for use
+/// within the guard's scope.
+pub struct PreemptGuard<'a> {
+    context: &'a ContextLock,
+    token: &'a mut CleanLockToken,
+}
+
+impl<'a> PreemptGuard<'a> {
+    pub fn new(context: &'a ContextLock, token: &'a mut CleanLockToken) -> PreemptGuard<'a> {
+        context.write(token.token()).is_preemptable = false;
+        PreemptGuard { context, token }
+    }
+
+    /// Get a mutable reference to the underlying `CleanLockToken`.
+    ///
+    /// This is necessary because the `PreemptGuard` owns the mutable reference to the token
+    /// (to use it in `Drop`), so we cannot use the original `token` variable while the guard exists.
+    pub fn token(&mut self) -> &mut CleanLockToken {
+        self.token
+    }
+}
+
+impl Drop for PreemptGuard<'_> {
+    fn drop(&mut self) {
+        self.context.write(self.token.token()).is_preemptable = true;
+    }
+}

src/scheme/user.rs

@@ -27,7 +27,7 @@ use crate::{
             AddrSpace, AddrSpaceWrapper, BorrowedFmapSource, Grant, GrantFileRef, MmapMode,
             PageSpan, DANGLING,
         },
-        BorrowedHtBuf, ContextLock, Status,
+        BorrowedHtBuf, ContextLock, PreemptGuard, Status,
     },
     event,
     memory::Frame,
@@ -311,7 +311,8 @@ impl UserInner {
             // need to ensure that the following operations are atomic as
             // otherwise the process will be blocked forever.
             let current_context = context::current();
-            current_context.write(token.token()).is_preemptable = false;
+            let mut preempt = PreemptGuard::new(&current_context, token);
+            let token = preempt.token();
             current_context
                 .write(token.token())
                 .block("UserInner::call");
@@ -331,7 +332,6 @@ impl UserInner {
             self.todo.send(sqe, token);
 
             event::trigger(self.root_id, self.handle_id, EVENT_READ);
-            current_context.write(token.token()).is_preemptable = true;
         }
 
         loop {
@@ -418,7 +418,8 @@ impl UserInner {
                             // We do not want to preempt between sending the
                             // cancellation and blocking again where we might
                             // miss a wakeup.
-                            current_context.write(token.token()).is_preemptable = false;
+                            let mut preempt = PreemptGuard::new(&current_context, token);
+                            let token = preempt.token();
 
                             if canceling {
                                 self.todo.send(
@@ -447,7 +448,6 @@ impl UserInner {
                                 .write(token.token())
                                 .block("UserInner::call (spurious wakeup)");
                             drop(states);
-                            current_context.write(token.token()).is_preemptable = true;
                         }
 
                         // invalid state