fix(sync/wait_condition): deadlock in WaitCondition::wait

Instead of using a simple switch to determine if preemption is enabled
(`is_preemptable: bool`), a counter is used instead. This handles the
case where a function holding a `PreemptGuard` calls another function
that creates a new `PreemptGuard`.

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

src/context/context.rs

@@ -144,8 +144,10 @@ pub struct Context {
     pub egid: u32,
     pub pid: usize,
 
-    // Use PreemptGuard
-    pub(super) is_preemptable: bool,
+    // See [`PreemptGuard`]
+    //
+    // When > 0, preemption is disabled.
+    pub(super) preempt_locks: usize,
 }
 
 #[derive(Debug)]
@@ -201,12 +203,16 @@ impl Context {
             #[cfg(feature = "syscall_debug")]
             syscall_debug_info: crate::syscall::debug::SyscallDebugInfo::default(),
 
-            is_preemptable: true,
+            preempt_locks: 0,
         };
         cpu_stats::add_context();
         Ok(this)
     }
 
+    pub fn is_preemptable(&self) -> bool {
+        self.preempt_locks == 0
+    }
+
     /// Block the context, and return true if it was runnable before being blocked
     pub fn block(&mut self, reason: &'static str) -> bool {
         if self.status.is_runnable() {

src/context/mod.rs

@@ -190,7 +190,7 @@ pub struct PreemptGuard<'a> {
 
 impl<'a> PreemptGuard<'a> {
     pub fn new(context: &'a ContextLock, token: &'a mut CleanLockToken) -> PreemptGuard<'a> {
-        context.write(token.token()).is_preemptable = false;
+        context.write(token.token()).preempt_locks += 1;
         PreemptGuard { context, token }
     }
 
@@ -205,6 +205,6 @@ impl<'a> PreemptGuard<'a> {
 
 impl Drop for PreemptGuard<'_> {
     fn drop(&mut self) {
-        self.context.write(self.token.token()).is_preemptable = true;
+        self.context.write(self.token.token()).preempt_locks -= 1;
     }
 }

src/context/switch.rs

@@ -165,7 +165,7 @@ pub fn switch(token: &mut CleanLockToken) -> SwitchResult {
         // We are careful not to lock this context twice
         let prev_context_guard = unsafe { prev_context_lock.write_arc() };
 
-        if !prev_context_guard.is_preemptable {
+        if !prev_context_guard.is_preemptable() {
             return SwitchResult::AllContextsIdle;
         }
 

src/sync/wait_condition.rs

@@ -4,8 +4,8 @@ use alloc::{
 };
 
 use crate::{
-    context::{self, ContextLock},
-    sync::{CleanLockToken, Mutex, L1},
+    context::{self, ContextLock, PreemptGuard},
+    sync::{CleanLockToken, L1, Mutex},
 };
 
 #[derive(Debug)]
@@ -50,6 +50,12 @@ impl WaitCondition {
     pub fn wait<T>(&self, guard: T, reason: &'static str, token: &mut CleanLockToken) -> bool {
         let current_context_ref = context::current();
         {
+            // Avoid a context switch between blocking ourselves and adding
+            // ourselves to the wait list as otherwise we might miss a wakeup.
+            // We cannot add ourselves to the wait list first as that would lead
+            // 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()