Repository: theseus-os/Theseus (language: rust)

Issue:
Improve mutex performance
#576 provides a simple mutex implementation; ideally, we'd use a more performant one. Currently, `raw_mutex` only has a fast path (mutex is unlocked) and a slow path (mutex is locked - block current task). We should also have a middle path, which spins for a bit, hoping the mutex is unlocked soon.

Some useful links:
- How Linux does it: https://0xax.gitbooks.io/linux-insides/content/SyncPrim/linux-sync-4.html
- Another possible implementation: https://arxiv.org/pdf/1906.00490.pdf
- The concept of a backoff could be useful in some capacity: https://docs.rs/crossbeam/latest/crossbeam/utils/struct.Backoff.html


Discussion / hints:
#576 isn't merged. How can it the provide the simple mutex implementation? Is this issue still relevant?
@tsoutsman now that we have your new sync primitives, do you want to attempt adding the middle path? or just close this issue?

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index 7506503f23..32ccf3487d 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3712,6 +3712,7 @@ name = "sync_block"
 version = "0.1.0"
 dependencies = [
  "cpu_local_preemption",
+ "log",
  "mpmc_queue",
  "scheduler",
  "sync",
diff --git a/kernel/sync_block/Cargo.toml b/kernel/sync_block/Cargo.toml
index 567444d9cc..43ee4aa209 100644
--- a/kernel/sync_block/Cargo.toml
+++ b/kernel/sync_block/Cargo.toml
@@ -7,6 +7,7 @@ edition = "2021"
 [dependencies]
 cpu_local_preemption = { path = "../cpu_local_preemption" }
 mpmc_queue = { path = "../../libs/mpmc_queue" }
+log = "0.4.8"
 scheduler = { path = "../scheduler" }
 sync = { path = "../../libs/sync" }
 task = { path = "../task" }
diff --git a/kernel/sync_block/src/lib.rs b/kernel/sync_block/src/lib.rs
index 8838840063..7786c465aa 100644
--- a/kernel/sync_block/src/lib.rs
+++ b/kernel/sync_block/src/lib.rs
@@ -1,8 +1,9 @@
-#![feature(negative_impls)]
+#![feature(negative_impls, let_chains)]
 #![no_std]
 
 mod condvar;
 
+use core::sync::atomic::{AtomicUsize, Ordering};
 use sync::{spin, MutexFlavor, RwLockFlavor};
 use wait_queue::WaitQueue;
 
@@ -22,21 +23,39 @@ pub struct Block {}
 
 impl MutexFlavor for Block {
     #[allow(clippy::declare_interior_mutable_const)]
-    const INIT: Self::LockData = WaitQueue::new();
+    const INIT: Self::LockData = Self::LockData {
+        queue: WaitQueue::new(),
+        holder: AtomicUsize::new(0),
+    };
 
-    type LockData = WaitQueue;
+    type LockData = MutexData;
 
     type Guard = ();
 
     #[inline]
     fn try_lock<'a, T>(
         mutex: &'a spin::Mutex<T>,
-        _: &'a Self::LockData,
+        data: &'a Self::LockData,
     ) -> Option<(spin::MutexGuard<'a, T>, Self::Guard)>
     where
         T: ?Sized,
     {
-        mutex.try_lock().map(|guard| (guard, ()))
+        let guard = mutex.try_lock()?;
+        // There is an interleaving where we are here, the previous holder hasn't yet
+        // ran post_unlock, and another thread A acquires the holder_id for the middle
+        // path all at the same time. This will result in thread A waiting on the old
+        // holder of the lock. This isn't great, but it is still correct, as the thread
+        // would only do so for a short period of time. It does however mean that the
+        // middle path would be "useless" because as soon as the holder changes, thread
+        // A would enter the slow path.
+        //
+        // We can prevent this by using an atomic usize in the inner spin lock rather
+        // than an atomic bool. A non-zero value would represent the task ID of the
+        // holder, and a zero would represent the unlocked state. However, this
+        // would be very hard to integrate with the current sync API.
+        data.holder
+            .store(task::get_my_current_task_id(), Ordering::Release);
+        Some((guard, ()))
     }
 
     #[inline]
@@ -47,21 +66,57 @@ impl MutexFlavor for Block {
     where
         T: ?Sized,
     {
+        // Fast path
         // This must be a strong compare exchange, otherwise we could block ourselves
         // when the mutex is unlocked and never be unblocked.
         if let Some(guards) = Self::try_lock(mutex, data) {
-            guards
+            return guards;
+        }
+
+        // Middle path: Wait for one timeslice (see below) of the mutex holder to see if
+        // they release the lock.
+        let holder_id = data.holder.load(Ordering::Acquire);
+        if holder_id != 0 && let Some(holder_task) = task::get_task(holder_id).and_then(|task| task.upgrade()) {
+            // Hypothetically, if holder_task is running on another core and is perfectly in
+            // sync with us, we would only ever check if they are running when we are also
+            // running and so we wouldn't detect when their timeslice is over. However, the
+            // likelihood of this is infinitesimally small and the code, is still correct as
+            // once the lock is released the holder will still set data.holder to 0 and we
+            // will exit the loop.
+            while holder_task.is_running() && data.holder.load(Ordering::Acquire) == holder_id {
+                core::hint::spin_loop();
+            }
+            // Holder is either no longer running, or has released the lock.
+            // Either way we will try the fast path one more time before moving
+            // onto the slow path.
         } else {
-            data.wait_until(|| Self::try_lock(mutex, data))
+            // Unlikely case that another thread just acquired the lock, but hasn't yet set
+            // data.holder.
+            log::warn!("could not get holder task for mutex middle path");
         }
+
+        if let Some(guards) = Self::try_lock(mutex, data) {
+            return guards;
+        }
+
+        // Slow path
+        data.queue.wait_until(|| Self::try_lock(mutex, data))
     }
 
     #[inline]
     fn post_unlock(data: &Self::LockData) {
-        data.notify_one();
+        // See comments in try_lock and lock on why this is necessary.
+        data.holder.store(0, Ordering::Release);
+        data.queue.notify_one();
     }
 }
 
+#[doc(hidden)]
+pub struct MutexData {
+    queue: WaitQueue,
+    holder: AtomicUsize,
+}
+
 impl RwLockFlavor for Block {
     #[allow(clippy::declare_interior_mutable_const)]
     const INIT: Self::LockData = RwLockData {
diff --git a/kernel/task_struct/src/lib.rs b/kernel/task_struct/src/lib.rs
index b3931e4723..360ecbb501 100755
--- a/kernel/task_struct/src/lib.rs
+++ b/kernel/task_struct/src/lib.rs
@@ -307,8 +307,9 @@ impl Task {
         stack: Option<Stack>,
         states_to_inherit: InheritedStates,
     ) -> Result<Task, &'static str> {
-        /// The counter of task IDs
-        static TASKID_COUNTER: AtomicUsize = AtomicUsize::new(0);
+        /// The counter of task IDs. We start at `1` such that `0` can be used 
+        /// as a task ID that indicates the absence of a task, e.g., in sync primitives. 
+        static TASKID_COUNTER: AtomicUsize = AtomicUsize::new(1);
 
         let (mmi, namespace, env, app_crate) = states_to_inherit.into_tuple();
         let kstack = stack


Test patch: