Lock wait queue together when receive_into_user

src/sync/ordered.rs

@@ -271,6 +271,38 @@ impl<L: Level, T> Mutex<L, T> {
         })
     }
 
+    /// Arcquires the lock_token to replace older MutexGuard.
+    /// SAFETY: Caller must guarantee lock_token is coming from MutexWriteGuard::into_token() from the same lock.
+    ///         OR Caller must guarantee lock_token is coming from different lock, which can happen when two lock need to copy data each other.
+    pub unsafe fn relock<'a>(&'a self, lock_token: LockToken<'a, L>) -> MutexGuard<'a, L, T> {
+        let inner = {
+            #[cfg(feature = "busy_panic")]
+            let mut i = DEADLOCK_SPIN_CAP;
+            let my_percpu = PercpuBlock::current();
+
+            loop {
+                match self.inner.try_lock() {
+                    Some(inner) => break inner,
+                    None => {
+                        my_percpu.maybe_handle_tlb_shootdown();
+                        core::hint::spin_loop();
+                        #[cfg(feature = "busy_panic")]
+                        {
+                            i -= 1;
+                            if i == 0 {
+                                panic!("Deadlock at mutex may have triggered")
+                            }
+                        }
+                    }
+                }
+            }
+        };
+        MutexGuard {
+            inner,
+            lock_token: lock_token,
+        }
+    }
+
     /// Consumes this Mutex, returning the underlying data.
     pub fn into_inner(self) -> T {
         self.inner.into_inner()
@@ -299,6 +331,14 @@ impl<'a, L: Level, T: ?Sized + 'a> MutexGuard<'a, L, T> {
     pub fn into_split(self) -> (spin::MutexGuard<'a, T>, LockToken<'a, L>) {
         (self.inner, self.lock_token)
     }
+
+    /// Merge the guard from `into_split`
+    pub fn from_split(lock: spin::MutexGuard<'a, T>, token: LockToken<'a, L>) -> Self {
+        Self {
+            inner: lock,
+            lock_token: token,
+        }
+    }
 }
 
 impl<'a, L: Level, T: ?Sized + 'a> core::ops::Deref for MutexGuard<'a, L, T> {

src/sync/wait_queue.rs

@@ -2,7 +2,7 @@ use alloc::collections::VecDeque;
 use syscall::{EAGAIN, EINTR};
 
 use crate::{
-    sync::{CleanLockToken, LockToken, Mutex, WaitCondition, L1, L2},
+    sync::{CleanLockToken, LockToken, Mutex, MutexGuard, WaitCondition, L1, L2},
     syscall::{
         error::{Error, Result, EINVAL},
         usercopy::UserSliceWo,
@@ -36,38 +36,34 @@ impl<T> WaitQueue<T> {
         reason: &'static str,
         token: &mut CleanLockToken,
     ) -> Result<usize> {
+        let mut out_guard = self.outgoing.lock(token.token());
         loop {
-            let mut tmp_queue = VecDeque::new();
-            {
-                let mut out = self.outgoing.lock(token.token());
-                if !out.is_empty() {
-                    let (s1, s2) = out.as_slices();
-                    let s1_bytes = unsafe {
-                        core::slice::from_raw_parts(s1.as_ptr().cast::<u8>(), size_of_val(s1))
-                    };
-                    let s2_bytes = unsafe {
-                        core::slice::from_raw_parts(s2.as_ptr().cast::<u8>(), size_of_val(s2))
-                    };
+            let (mut outgoing, mut token) = out_guard.into_split();
+            if !outgoing.is_empty() {
+                let (s1, s2) = outgoing.as_slices();
+                let s1_bytes = unsafe {
+                    core::slice::from_raw_parts(s1.as_ptr().cast::<u8>(), size_of_val(s1))
+                };
+                let s2_bytes = unsafe {
+                    core::slice::from_raw_parts(s2.as_ptr().cast::<u8>(), size_of_val(s2))
+                };
 
-                    let mut bytes_copied = buf.copy_common_bytes_from_slice(s1_bytes)?;
+                let mut bytes_copied = buf.copy_common_bytes_from_slice(s1_bytes)?;
 
-                    if let Some(buf_for_s2) = buf.advance(s1_bytes.len()) {
-                        bytes_copied += buf_for_s2.copy_common_bytes_from_slice(s2_bytes)?;
-                    }
-
-                    let _ = out.drain(..bytes_copied / size_of::<T>());
-                    return Ok(bytes_copied);
+                if let Some(buf_for_s2) = buf.advance(s1_bytes.len()) {
+                    bytes_copied += buf_for_s2.copy_common_bytes_from_slice(s2_bytes)?;
                 }
 
-                // Act as tmp_queue.drain(..), but much faster
-                core::mem::swap(&mut *out, &mut tmp_queue);
+                let _ = outgoing.drain(..bytes_copied / size_of::<T>());
+                return Ok(bytes_copied);
             }
 
-            let incoming_guard = self.incoming.lock(token.token());
+            let incoming_guard = unsafe { self.incoming.relock(token.token()) };
             let (mut incoming, mut split_token) = incoming_guard.into_split();
 
             if incoming.is_empty() {
                 if block {
+                    drop(outgoing);
                     // SAFETY: Uses wait_inner because this inner is L2. It's guaranteed there's no other
                     // lock held at this point because clean token is provided from caller.
                     if !self
@@ -76,6 +72,7 @@ impl<T> WaitQueue<T> {
                     {
                         return Err(Error::new(EINTR));
                     }
+                    out_guard = unsafe { self.outgoing.relock(token) };
                     continue;
                 } else if buf.is_empty() {
                     return Ok(0);
@@ -87,15 +84,10 @@ impl<T> WaitQueue<T> {
                 }
             }
 
-            // Act as incoming.drain(..), but much faster
-            core::mem::swap(&mut *incoming, &mut tmp_queue);
+            core::mem::swap(&mut *incoming, &mut outgoing);
             drop(incoming);
 
-            {
-                let mut out = self.outgoing.lock(token.token());
-                // outgoing = incoming
-                *out = tmp_queue;
-            }
+            out_guard = MutexGuard::from_split(outgoing, token);
         }
     }