Optimize mutex and semaphore for speed on low contention

Author: polytypic

lib/picos_std.sync/mutex.ml

@@ -9,18 +9,21 @@ type _ tdt =
   | Nothing : [> `Nothing ] tdt
 
 type state =
+  | Locked
   | Unlocked
-  | Locked of { fiber : Fiber.Maybe.t; waiters : [ `Entry ] tdt Q.t }
+  | Queued of { fiber : Fiber.Maybe.t; waiters : [ `Entry ] tdt Q.t }
 
 type t = state Atomic.t
 
 let create ?padded () = Multicore_magic.copy_as ?padded @@ Atomic.make Unlocked
-let locked_nothing = Locked { fiber = Fiber.Maybe.nothing; waiters = T Zero }
 
 let rec unlock_as owner t backoff =
   match Atomic.get t with
   | Unlocked -> unlocked ()
-  | Locked r as before ->
+  | Locked as before ->
+      if not (Atomic.compare_and_set t before Unlocked) then
+        unlock_as owner t (Backoff.once backoff)
+  | Queued r as before ->
       if Fiber.Maybe.equal r.fiber owner then
         match r.waiters with
         | T Zero ->
@@ -29,42 +32,72 @@ let rec unlock_as owner t backoff =
         | T (One _ as q) ->
             let (Entry { trigger; fiber }) = Q.head q in
             let waiters = Q.tail q in
-            let after = Locked { fiber; waiters } in
+            let after = Queued { fiber; waiters } in
             if Atomic.compare_and_set t before after then Trigger.signal trigger
             else unlock_as owner t (Backoff.once backoff)
       else not_owner ()
 
 let[@inline] unlock ?checked t =
-  let owner = Fiber.Maybe.current_if checked in
-  unlock_as owner t Backoff.default
+  match checked with
+  | Some false ->
+      if
+        (* The unlock operation will mutate the atomic location and will be
+           sequentially consistent.  The fenceless get potentially allows us to
+           avoid performing a failed mutation attempt causing cache coherency
+           traffic and fenceless get here performs better on ARM. *)
+        Multicore_magic.fenceless_get t != Locked
+        || not (Atomic.compare_and_set t Locked Unlocked)
+      then unlock_as Fiber.Maybe.nothing t Backoff.default
+  | None | Some true ->
+      let owner = Fiber.Maybe.of_fiber (Fiber.current ()) in
+      unlock_as owner t Backoff.default
 
 let rec cleanup_as (Entry entry_r as entry : [ `Entry ] tdt) t backoff =
   (* We have been canceled.  If we are the owner, we must unlock the mutex.
      Otherwise we must remove our entry from the queue. *)
   match Atomic.get t with
-  | Locked r as before -> begin
+  | Queued r as before -> begin
       match r.waiters with
       | T Zero -> unlock_as entry_r.fiber t backoff
       | T (One _ as q) ->
           let waiters = Q.remove q entry in
           if r.waiters == waiters then unlock_as entry_r.fiber t backoff
           else
-            let after = Locked { fiber = r.fiber; waiters } in
+            let after = Queued { fiber = r.fiber; waiters } in
             if not (Atomic.compare_and_set t before after) then
               cleanup_as entry t (Backoff.once backoff)
     end
+  | Locked -> unlock_as entry_r.fiber t backoff
   | Unlocked -> unlocked ()
 
 let rec lock_as fiber t entry backoff =
   match Atomic.get t with
   | Unlocked as before ->
       let after =
-        if fiber == Fiber.Maybe.nothing then locked_nothing
-        else Locked { fiber; waiters = T Zero }
+        if fiber == Fiber.Maybe.nothing then Locked
+        else Queued { fiber; waiters = T Zero }
       in
       if not (Atomic.compare_and_set t before after) then
         lock_as fiber t entry (Backoff.once backoff)
-  | Locked r as before ->
+  | Locked as before ->
+      let (Entry entry_r as entry : [ `Entry ] tdt) =
+        match entry with
+        | Nothing ->
+            let trigger = Trigger.create () in
+            Entry { trigger; fiber }
+        | Entry _ as entry -> entry
+      in
+      let waiters = Q.singleton entry in
+      let after = Queued { fiber = Fiber.Maybe.nothing; waiters } in
+      if Atomic.compare_and_set t before after then begin
+        match Trigger.await entry_r.trigger with
+        | None -> ()
+        | Some (exn, bt) ->
+            cleanup_as entry t Backoff.default;
+            Printexc.raise_with_backtrace exn bt
+      end
+      else lock_as fiber t entry (Backoff.once backoff)
+  | Queued r as before ->
       if Fiber.Maybe.unequal r.fiber fiber then
         let (Entry entry_r as entry : [ `Entry ] tdt) =
           match entry with
@@ -74,7 +107,7 @@ let rec lock_as fiber t entry backoff =
           | Entry _ as entry -> entry
         in
         let waiters = Q.add r.waiters entry in
-        let after = Locked { fiber = r.fiber; waiters } in
+        let after = Queued { fiber = r.fiber; waiters } in
         if Atomic.compare_and_set t before after then begin
           match Trigger.await entry_r.trigger with
           | None -> ()
@@ -86,15 +119,27 @@ let rec lock_as fiber t entry backoff =
       else owner ()
 
 let[@inline] lock ?checked t =
-  let fiber = Fiber.Maybe.current_and_check_if checked in
-  lock_as fiber t Nothing Backoff.default
+  match checked with
+  | Some false ->
+      if
+        (* The lock operation will mutate the atomic location and will be
+           sequentially consistent.  The fenceless get potentially allows us to
+           avoid performing a failed mutation attempt causing cache coherency
+           traffic and fenceless get here performs better on ARM. *)
+        Multicore_magic.fenceless_get t != Unlocked
+        || not (Atomic.compare_and_set t Unlocked Locked)
+      then lock_as Fiber.Maybe.nothing t Nothing Backoff.default
+  | None | Some true ->
+      let fiber = Fiber.current () in
+      Fiber.check fiber;
+      lock_as (Fiber.Maybe.of_fiber fiber) t Nothing Backoff.default
 
 let try_lock ?checked t =
   let fiber = Fiber.Maybe.current_and_check_if checked in
   Atomic.get t == Unlocked
   && Atomic.compare_and_set t Unlocked
-       (if fiber == Fiber.Maybe.nothing then locked_nothing
-        else Locked { fiber; waiters = T Zero })
+       (if fiber == Fiber.Maybe.nothing then Locked
+        else Queued { fiber; waiters = T Zero })
 
 let protect ?checked t body =
   let fiber = Fiber.Maybe.current_and_check_if checked in

lib/picos_std.sync/picos_std_sync.mli

@@ -525,4 +525,12 @@ end
     false sharing}.  For relatively long lived objects this can improve
     performance and make performance more stable at the cost of using more
     memory.  It is not recommended to use [~padded:true] for short lived
-    objects. *)
+    objects.
+
+    The primitives provided by this library are generally optimized for low
+    contention scenariors and size.  Generally speaking, for best performance
+    and scalability, you should try to avoid high contention scenarios by
+    architecting your program to distribute processing such that sequential
+    bottlenecks are avoided.  If high contention is unavoidable then other
+    communication and synchronization primitive implementations may provide
+    better performance. *)

lib/picos_std.sync/semaphore.ml

@@ -11,7 +11,10 @@ module Counting = struct
     Atomic.make (Obj.repr count) |> Multicore_magic.copy_as ?padded
 
   let rec release t backoff =
-    let before = Atomic.get t in
+    (* The release operation will mutate the atomic location and will be
+       sequentially consistent.  The fenceless get here performs better on
+       ARM. *)
+    let before = Multicore_magic.fenceless_get t in
     if Obj.is_int before then begin
       let count = Obj.obj before in
       if count < count + 1 then begin
@@ -39,7 +42,10 @@ module Counting = struct
       then cleanup t trigger (Backoff.once backoff)
 
   let rec acquire t backoff =
-    let before = Atomic.get t in
+    (* The acquire operation will mutate the atomic location and will be
+       sequentially consistent.  The fenceless get here performs better on
+       ARM. *)
+    let before = Multicore_magic.fenceless_get t in
     if Obj.is_int before then
       let count = Obj.obj before in
       if 0 < count then begin
@@ -96,7 +102,10 @@ module Binary = struct
   let make ?padded initial = Counting.make ?padded (Bool.to_int initial)
 
   let rec release t backoff =
-    let before = Atomic.get t in
+    (* The release operation will mutate the atomic location and will be
+       sequentially consistent.  The fenceless get here performs better on
+       ARM. *)
+    let before = Multicore_magic.fenceless_get t in
     if Obj.is_int before then begin
       let count = Obj.obj before in
       if count = 0 then