Add ordered lock for WaitQueue mutex

src/context/mod.rs

@@ -11,8 +11,7 @@ use crate::{
     paging::{RmmA, RmmArch, TableKind},
     percpu::PercpuBlock,
     sync::{
-        ArcRwLockWriteGuard, CleanLockToken, LockToken, RwLock, RwLockReadGuard, RwLockWriteGuard,
-        L0, L1, L2,
+        ArcRwLockWriteGuard, CleanLockToken, L0, L1, L4, LockToken, RwLock, RwLockReadGuard, RwLockWriteGuard
     },
     syscall::error::Result,
 };
@@ -23,8 +22,8 @@ pub use self::{
     switch::switch,
 };
 
-pub type ContextLock = RwLock<L2, Context>;
-pub type ArcContextLockWriteGuard = ArcRwLockWriteGuard<L2, Context>;
+pub type ContextLock = RwLock<L4, Context>;
+pub type ArcContextLockWriteGuard = ArcRwLockWriteGuard<L4, Context>;
 
 #[cfg(target_arch = "aarch64")]
 #[path = "arch/aarch64.rs"]

src/context/signal.rs

@@ -78,7 +78,7 @@ pub fn excp_handler(excp: syscall::Exception) {
 
     let context = current.write(token.token());
 
-    let Some(eh) = context.sig.as_ref().and_then(|s| s.excp_handler) else {
+    let Some(_eh) = context.sig.as_ref().and_then(|s| s.excp_handler) else {
         // TODO: Let procmgr print this?
         info!(
             "UNHANDLED EXCEPTION, CPU {}, PID {}, NAME {}, CONTEXT {current:p}",

src/event.rs

@@ -33,8 +33,8 @@ impl EventQueue {
         }
     }
 
-    pub fn is_currently_empty(&self) -> bool {
-        self.queue.is_currently_empty()
+    pub fn is_currently_empty(&self, token: &mut CleanLockToken) -> bool {
+        self.queue.is_currently_empty(token)
     }
 
     pub fn read(&self, buf: UserSliceWo, block: bool, token: &mut CleanLockToken) -> Result<usize> {

src/scheme/event.rs

@@ -119,7 +119,7 @@ impl KernelScheme for EventScheme {
             // It is always possible to write events
             ready |= EventFlags::EVENT_WRITE;
         }
-        if flags.contains(EventFlags::EVENT_READ) && !queue.is_currently_empty() {
+        if flags.contains(EventFlags::EVENT_READ) && !queue.is_currently_empty(token) {
             // It is possible to read if queue is not empty
             ready |= EventFlags::EVENT_READ;
         }

src/scheme/mod.rs

@@ -305,8 +305,8 @@ impl KernelScheme for SchemeList {
         token: &mut CleanLockToken,
     ) -> Result<EventFlags> {
         match self.get_user_inner(id, token) {
-            Some(inner) => inner.fevent(flags),
-            _ => Err(Error::new(EBADF)),
+            Some(inner) => inner.fevent(flags, token),
+            _ => return Err(Error::new(EBADF)),
         }
     }
 

src/scheme/sys/mod.rs

@@ -5,7 +5,7 @@
 
 use ::syscall::{
     dirent::{DirEntry, DirentBuf, DirentKind},
-    EACCES, EBADFD, EINVAL, EIO, EISDIR, ENOTDIR, EPERM,
+    EACCES, EINVAL, EIO, EISDIR, ENOTDIR, EPERM,
 };
 use alloc::{sync::Arc, vec::Vec};
 use core::{

src/scheme/user.rs

@@ -971,10 +971,10 @@ impl UserInner {
         Ok(())
     }
 
-    pub fn fevent(&self, flags: EventFlags) -> Result<EventFlags> {
+    pub fn fevent(&self, flags: EventFlags, token: &mut CleanLockToken) -> Result<EventFlags> {
         // TODO: Should the root scheme also suppress events if `flags` does not contain
         // `EVENT_READ`?
-        Ok(if self.todo.is_currently_empty() {
+        Ok(if self.todo.is_currently_empty(token) {
             EventFlags::empty()
         } else {
             EventFlags::EVENT_READ.intersection(flags)

src/sync/wait_condition.rs

@@ -5,12 +5,12 @@ use alloc::{
 
 use crate::{
     context::{self, ContextLock, PreemptGuard},
-    sync::{CleanLockToken, Mutex, L1},
+    sync::{CleanLockToken, LockToken, Lower, Mutex, L1, L2, L3},
 };
 
 #[derive(Debug)]
 pub struct WaitCondition {
-    contexts: Mutex<L1, Vec<Weak<ContextLock>>>,
+    contexts: Mutex<L3, Vec<Weak<ContextLock>>>,
 }
 
 impl WaitCondition {
@@ -46,7 +46,8 @@ impl WaitCondition {
         len
     }
 
-    // Wait until notified. Unlocks guard when blocking is ready. Returns false if resumed by a signal or the notify_signal function
+    /// Wait until notified. Unlocks guard when blocking is ready. Returns false if resumed by a signal or the notify_signal function.
+    /// Wrapper to wait_setup -> drop(guard) -> context::switch -> wait_cleanup without currently holding lock for guard.
     pub fn wait<T>(&self, guard: T, reason: &'static str, token: &mut CleanLockToken) -> bool {
         let current_context_ref = context::current();
         {
@@ -56,36 +57,68 @@ impl WaitCondition {
             // to deadlock if we were woken up immediately.
             let mut preempt = PreemptGuard::new(&current_context_ref, token);
             let token = preempt.token();
-            {
-                let mut context = current_context_ref.write(token.token());
-                if let Some((control, pctl, _)) = context.sigcontrol()
-                    && control.currently_pending_unblocked(pctl) != 0
-                {
-                    return false;
-                }
-                context.block(reason);
+            if !self.wait_setup(&current_context_ref, reason, token.token()) {
+                return false;
             }
 
-            self.contexts
-                .lock(token.token())
-                .push(Arc::downgrade(&current_context_ref));
-
             drop(guard);
         }
 
         context::switch(token);
 
-        let mut waited = true;
+        self.wait_cleanup(&current_context_ref, token.token())
+    }
 
+    /// Enqueues the context and sets it to blocked.
+    /// Returns true if successfully blocked, false if a signal is pending.
+    pub fn wait_setup<'a, LP>(
+        &self,
+        current_context_ref: &Arc<ContextLock>,
+        reason: &'static str,
+        mut lock_token: LockToken<'a, LP>,
+    ) -> bool
+    where
+        LP: Lower<L2>,
+    {
         {
-            let mut contexts = self.contexts.lock(token.token());
-
-            if let Some(index) = contexts
-                .iter()
-                .position(|c| Weak::as_ptr(c) == Arc::as_ptr(&current_context_ref))
+            let mut context = current_context_ref.write(LockToken::downgraded(lock_token.token()));
+            if let Some((control, pctl, _)) = context.sigcontrol()
+                && control.currently_pending_unblocked(pctl) != 0
             {
-                contexts.remove(index);
+                return false;
+            }
+            context.block(reason);
+        }
+
+        self.contexts
+            .lock(LockToken::downgraded(lock_token))
+            .push(Arc::downgrade(current_context_ref));
+
+        true
+    }
+
+    /// Cleans up the wait list after waking up.
+    /// Returns true if we actually waited, false if we were removed by signal/notify_signal.
+    pub fn wait_cleanup<'a, LP>(
+        &self,
+        current_context_ref: &Arc<ContextLock>,
+        lock_token: LockToken<'a, LP>,
+    ) -> bool
+    where
+        LP: Lower<L1>,
+    {
+        let mut waited = true;
+        let mut contexts = self.contexts.lock(LockToken::downgraded(lock_token));
+
+        // TODO: retain
+        let mut i = 0;
+        while i < contexts.len() {
+            if Weak::as_ptr(&contexts[i]) == Arc::as_ptr(current_context_ref) {
+                contexts.remove(i);
                 waited = false;
+                break;
+            } else {
+                i += 1;
             }
         }
 

src/sync/wait_queue.rs

@@ -1,9 +1,9 @@
 use alloc::collections::VecDeque;
-use spin::Mutex;
 use syscall::{EAGAIN, EINTR};
 
 use crate::{
-    sync::{CleanLockToken, WaitCondition},
+    context::{self, PreemptGuard},
+    sync::{CleanLockToken, Mutex, WaitCondition, L1},
     syscall::{
         error::{Error, Result, EINVAL},
         usercopy::UserSliceWo,
@@ -12,7 +12,7 @@ use crate::{
 
 #[derive(Debug)]
 pub struct WaitQueue<T> {
-    pub inner: Mutex<VecDeque<T>>,
+    pub inner: Mutex<L1, VecDeque<T>>,
     pub condition: WaitCondition,
 }
 
@@ -23,8 +23,8 @@ impl<T> WaitQueue<T> {
             condition: WaitCondition::new(),
         }
     }
-    pub fn is_currently_empty(&self) -> bool {
-        self.inner.lock().is_empty()
+    pub fn is_currently_empty(&self, token: &mut CleanLockToken) -> bool {
+        self.inner.lock(token.token()).is_empty()
     }
 
     pub fn receive(
@@ -33,8 +33,12 @@ impl<T> WaitQueue<T> {
         reason: &'static str,
         token: &mut CleanLockToken,
     ) -> Result<T> {
+        let current_context_ref = context::current();
+
         loop {
-            let mut inner = self.inner.lock();
+            let mut preempt = PreemptGuard::new(&current_context_ref, token);
+
+            let mut inner = self.inner.lock(preempt.token().token());
 
             match inner.pop_front() {
                 Some(t) => {
@@ -42,9 +46,23 @@ impl<T> WaitQueue<T> {
                 }
                 _ => {
                     if block {
-                        if !self.condition.wait(inner, reason, token) {
+                        let (_, mut inner_token) = inner.token_split();
+                        if !self.condition.wait_setup(
+                            &current_context_ref,
+                            reason,
+                            inner_token.token(),
+                        ) {
                             return Err(Error::new(EINTR));
                         }
+
+                        drop(inner);
+                        drop(preempt);
+
+                        context::switch(token);
+
+                        self.condition
+                            .wait_cleanup(&current_context_ref, token.token());
+
                         continue;
                     } else {
                         return Err(Error::new(EAGAIN));
@@ -53,22 +71,34 @@ impl<T> WaitQueue<T> {
             }
         }
     }
-
-    pub fn receive_into_user(
+pub fn receive_into_user(
         &self,
         buf: UserSliceWo,
         block: bool,
         reason: &'static str,
         token: &mut CleanLockToken,
     ) -> Result<usize> {
+        let current_context_ref = context::current();
+
         loop {
-            let mut inner = self.inner.lock();
+            let mut preempt = PreemptGuard::new(&current_context_ref, token);
+            
+            let mut inner = self.inner.lock(preempt.token().token());
 
             if inner.is_empty() {
                 if block {
-                    if !self.condition.wait(inner, reason, token) {
+                    let (_, mut inner_token) = inner.token_split();
+                    if !self.condition.wait_setup(&current_context_ref, reason, inner_token.token()) {
                         return Err(Error::new(EINTR));
                     }
+
+                    drop(inner);
+                    drop(preempt);
+
+                    context::switch(token);
+
+                    self.condition.wait_cleanup(&current_context_ref, token.token());
+                    
                     continue;
                 } else if buf.is_empty() {
                     return Ok(0);
@@ -102,7 +132,7 @@ impl<T> WaitQueue<T> {
 
     pub fn send(&self, value: T, token: &mut CleanLockToken) -> usize {
         let len = {
-            let mut inner = self.inner.lock();
+            let mut inner = self.inner.lock(token.token());
             inner.push_back(value);
             inner.len()
         };