diff --git a/src/context/context.rs b/src/context/context.rs index 0aaae6382c..b5d3da0656 100644 --- a/src/context/context.rs +++ b/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)] diff --git a/src/context/mod.rs b/src/context/mod.rs index 2990d50e08..49543e38ee 100644 --- a/src/context/mod.rs +++ b/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; + } +} diff --git a/src/scheme/user.rs b/src/scheme/user.rs index d987fae528..7da89f6db3 100644 --- a/src/scheme/user.rs +++ b/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(¤t_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(¤t_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