From b3a78994285b999ef1f6d2973a4f7a487a36d60a Mon Sep 17 00:00:00 2001
From: Chuandong Li <lichuandong.lcd@antgroup.com>
Date: Sun, 2 Jul 2023 23:32:33 +0800
Subject: [PATCH] Re-implement mutex with waiter queue

---
 framework/jinux-frame/src/sync/mutex.rs       | 61 +++++++++++++++----
 framework/jinux-frame/src/sync/wait.rs        |  7 ++-
 .../jinux-frame/src/vm/heap_allocator.rs      |  4 +-
 3 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/framework/jinux-frame/src/sync/mutex.rs b/framework/jinux-frame/src/sync/mutex.rs
index bb0377183..5e4217b71 100644
--- a/framework/jinux-frame/src/sync/mutex.rs
+++ b/framework/jinux-frame/src/sync/mutex.rs
@@ -1,30 +1,59 @@
-use super::spin::{SpinLock, SpinLockGuard};
+use super::WaitQueue;
+use core::cell::UnsafeCell;
 use core::ops::{Deref, DerefMut};
 
 use core::fmt;
+use core::sync::atomic::{AtomicBool, Ordering};
 
+/// A mutex with waitqueue.
 pub struct Mutex<T> {
-    inner: SpinLock<T>,
+    val: UnsafeCell<T>,
+    lock: AtomicBool,
+    queue: WaitQueue,
 }
 
 impl<T> Mutex<T> {
-    #[inline(always)]
-    pub const fn new(val: T) -> Self {
+    /// Create a new mutex.
+    pub fn new(val: T) -> Self {
         Self {
-            inner: SpinLock::new(val),
+            val: UnsafeCell::new(val),
+            lock: AtomicBool::new(false),
+            queue: WaitQueue::new(),
         }
     }
 
+    /// Acquire the mutex.
+    ///
+    /// This method runs in a block way until the mutex can be acquired.
     pub fn lock(&self) -> MutexGuard<T> {
-        MutexGuard {
-            lock: self.inner.lock(),
-        }
+        self.queue.wait_until(|| self.try_lock())
+    }
+
+    /// Try Acquire the mutex immedidately.
+    pub fn try_lock(&self) -> Option<MutexGuard<T>> {
+        self.acquire_lock().then(|| MutexGuard { mutex: &self })
+    }
+
+    /// Release the mutex and wake up one thread which is blocked on this mutex.
+    fn unlock(&self) {
+        self.release_lock();
+        self.queue.wake_one();
+    }
+
+    fn acquire_lock(&self) -> bool {
+        self.lock
+            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
+            .is_ok()
+    }
+
+    fn release_lock(&self) {
+        self.lock.store(false, Ordering::Release);
     }
 }
 
 impl<T: fmt::Debug> fmt::Debug for Mutex<T> {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        fmt::Debug::fmt(&self.inner, f)
+        fmt::Debug::fmt(&self.val, f)
     }
 }
 
@@ -32,20 +61,26 @@ unsafe impl<T: Send> Send for Mutex<T> {}
 unsafe impl<T: Send> Sync for Mutex<T> {}
 
 pub struct MutexGuard<'a, T> {
-    lock: SpinLockGuard<'a, T>,
+    mutex: &'a Mutex<T>,
 }
 
 impl<'a, T> Deref for MutexGuard<'a, T> {
     type Target = T;
 
-    fn deref(&self) -> &T {
-        self.lock.deref()
+    fn deref(&self) -> &Self::Target {
+        unsafe { &*self.mutex.val.get() }
     }
 }
 
 impl<'a, T> DerefMut for MutexGuard<'a, T> {
     fn deref_mut(&mut self) -> &mut Self::Target {
-        self.lock.deref_mut()
+        unsafe { &mut *self.mutex.val.get() }
+    }
+}
+
+impl<'a, T> Drop for MutexGuard<'a, T> {
+    fn drop(&mut self) {
+        self.mutex.unlock();
     }
 }
 
diff --git a/framework/jinux-frame/src/sync/wait.rs b/framework/jinux-frame/src/sync/wait.rs
index cdaa3290b..c34a45c3b 100644
--- a/framework/jinux-frame/src/sync/wait.rs
+++ b/framework/jinux-frame/src/sync/wait.rs
@@ -1,12 +1,11 @@
 use core::sync::atomic::{AtomicBool, Ordering};
 
+use super::SpinLock;
 use alloc::{collections::VecDeque, sync::Arc};
 use bitflags::bitflags;
 
 use crate::task::schedule;
 
-use super::SpinLock;
-
 /// A wait queue.
 ///
 /// One may wait on a wait queue to put its executing thread to sleep.
@@ -80,7 +79,9 @@ impl WaitQueue {
 
     /// removes all waiters that have finished wait
     fn finish_wait(&self) {
-        self.waiters.lock_irq_disabled().retain(|waiter| !waiter.is_finished())
+        self.waiters
+            .lock_irq_disabled()
+            .retain(|waiter| !waiter.is_finished())
     }
 }
 
diff --git a/framework/jinux-frame/src/vm/heap_allocator.rs b/framework/jinux-frame/src/vm/heap_allocator.rs
index 99e405737..8a48c7461 100644
--- a/framework/jinux-frame/src/vm/heap_allocator.rs
+++ b/framework/jinux-frame/src/vm/heap_allocator.rs
@@ -46,6 +46,8 @@ unsafe impl<const ORDER: usize> GlobalAlloc for LockedHeap<ORDER> {
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         debug_assert!(ptr as usize != 0);
-        self.0.lock_irq_disabled().dealloc(NonNull::new_unchecked(ptr), layout)
+        self.0
+            .lock_irq_disabled()
+            .dealloc(NonNull::new_unchecked(ptr), layout)
     }
 }