Optimize MPMC unbounded queue with status bits

Author: art-w

bench/bench_spsc_queue.ml

@@ -2,7 +2,7 @@ let item_count = 2_000_000
 
 module type QUEUE = sig
   type 'a t
-  val make : unit -> 'a t
+  val make : unit -> int t
   val push : 'a t -> 'a -> unit
   val pop : 'a t -> 'a option
   val name : string
@@ -55,7 +55,7 @@ end)
 
 module Mpmc_queue = Bench (struct
   include Lockfree.Mpmc_queue
-  let make () = make ()
+  let make () = make ~dummy:(-1) ()
   let name = "mpmc-queue"
 end)
 

bench/mpmc_queue.ml

@@ -1,17 +1,17 @@
 module type QUEUE = sig
   type 'a t
-  val make : unit -> 'a t
+  val make : unit -> int t
   val push : 'a t -> 'a -> unit
   val pop : 'a t -> 'a option
   val name : string
 end
 
 module Bench (Q : QUEUE) = struct
 
-  let num_of_elements = ref 500_000
+  let num_of_elements = ref 2_100_000
   let num_of_pushers = ref 4
   let num_of_takers = ref 4
-  let num_of_iterations = ref 10
+  let num_of_iterations = ref 20
 
   let taker queue num_of_elements () =
     let i = ref 0 in
@@ -101,14 +101,15 @@ module Bench (Q : QUEUE) = struct
 
   let bench : (unit -> _) list =
     [
+      benchmark ~takers:1 ~pushers:1;
       benchmark ~takers:4 ~pushers:4;
-      benchmark ~takers:1 ~pushers:8;
-      benchmark ~takers:8 ~pushers:1;
+      benchmark ~takers:1 ~pushers:7;
+      benchmark ~takers:7 ~pushers:1;
     ]
 end
 
 module Relaxed = Bench (struct
-  let name = "mpmc-relaxed"
+  let name = "mpmc-relaxed-fad"
   module Q = Lockfree.Mpmc_relaxed_queue
   include Q.Not_lockfree
   type 'a t = 'a Q.t
@@ -132,7 +133,7 @@ end)
 module Unbounded = Bench (struct
   let name = "mpmc-unbounded"
   include Lockfree.Mpmc_queue
-  let make () = make ()
+  let make () = make ~dummy:(-1) ()
 end)
 
 let bench = Relaxed.bench @ Relaxed_cas.bench @ Unbounded.bench

src/mpmc_queue.ml

@@ -1,13 +1,14 @@
-let default_capacity = 4096
-let spinlock_iterations = 16
+module Array = struct
+  include Array
+  let get = unsafe_get
+  let set = unsafe_set
+end
 
-type 'a cell =
-  | Empty
-  | Tombstone
-  | Value of 'a
+let default_capacity = 512
 
 type 'a s =
-  { buffer : 'a cell Atomic.t array
+  { status : int Atomic.t array
+  ; buffer : 'a array
   ; head : int Atomic.t
   ; tail : int Atomic.t
   ; rest : 'a s option Atomic.t
@@ -16,36 +17,40 @@ type 'a s =
 type 'a t =
   { first : 'a s Atomic.t
   ; last : 'a s Atomic.t
+  ; dummy : 'a
   }
 
-let make_s ~capacity () =
+let pack_size = Sys.int_size / 2
+
+let make_s ~capacity ~dummy =
   { head = Atomic.make 0
   ; tail = Atomic.make (-1)
-  ; buffer = Array.init capacity (fun _ -> Atomic.make Empty)
+  ; buffer = Array.make capacity dummy
+  ; status = Array.init (1 + capacity / pack_size) (fun _ -> Atomic.make 0)
   ; rest = Atomic.make None
   }
 
-let make ?(capacity = default_capacity) () =
-  let s = make_s ~capacity () in
+let make ?(capacity = default_capacity) ~dummy () =
+  let s = make_s ~capacity ~dummy in
   { first = Atomic.make s
   ; last = Atomic.make s
+  ; dummy
   }
 
 let rec gift_rest t some_s =
-  if Atomic.compare_and_set t.rest None some_s
-  then ()
-  else follow_rest t some_s
+  if not (Atomic.compare_and_set t.rest None some_s)
+  then follow_rest t some_s
 
 and follow_rest t some_s =
   match Atomic.get t.rest with
   | None -> gift_rest t some_s
   | Some t -> follow_rest t some_s
 
-let force_rest t =
+let force_rest ~dummy t =
   match Atomic.get t.rest with
   | Some s -> s
   | None ->
-      let s = make_s ~capacity:(Array.length t.buffer) () in
+      let s = make_s ~capacity:(Array.length t.buffer) ~dummy in
       let some_s = Some s in
       if Atomic.compare_and_set t.rest None some_s
       then s
@@ -55,44 +60,45 @@ let force_rest t =
               gift_rest rest some_s ;
               rest
 
-let rec push_s t x =
+let mark t i =
+  let status = t.status.(i / pack_size) in
+  let shift = 2 * (i mod pack_size) in
+  let status = Atomic.fetch_and_add status (1 lsl shift) in
+  (status lsr shift) land 1 = 0
+
+let rec push_s ~dummy t x =
   let i = Atomic.fetch_and_add t.tail 1 in
   if i < 0
-  then (let _ = force_rest t in push_s t x)
+  then (let _ = force_rest ~dummy t in push_s ~dummy t x)
   else if i >= Array.length t.buffer
   then false
   else begin
-    let cell = Array.unsafe_get t.buffer i in
-    match Atomic.get cell with
-    | Empty ->
-        if Atomic.compare_and_set cell Empty (Value x)
-        then true
-        else begin
-          assert (Atomic.get cell = Tombstone) ;
-          push_s t x
-        end
-    | Tombstone ->
-        push_s t x
-    | Value _ -> assert false
+    t.buffer.(i) <- x ;
+    if mark t i
+    then true
+    else begin
+      t.buffer.(i) <- dummy ;
+      let hd = Atomic.get t.head in
+      let _ : bool = Atomic.compare_and_set t.tail (i + 1) (hd + 1) in
+      push_s ~dummy t x
+    end
   end
 
-let rec push t x =
-  let last = Atomic.get t.last in
-  if push_s last x
-  then ()
-  else begin
-    let rest = force_rest last in
-    let _ : bool = Atomic.compare_and_set t.last last rest in
+let rec push ({ last ; dummy ; _ } as t) x =
+  let last_s = Atomic.get last in
+  if not (push_s ~dummy last_s x)
+  then begin
+    let rest = force_rest ~dummy last_s in
+    let _ : bool = Atomic.compare_and_set last last_s rest in
     push t x
   end
 
-
 type 'a pop_result =
   | Is_empty
   | Wait_for_it
   | Pop of 'a
 
-let rec pop_s t =
+let rec pop_s ~dummy t =
   let current_head = Atomic.get t.head in
   if current_head >= Array.length t.buffer
   then Is_empty
@@ -102,38 +108,17 @@ let rec pop_s t =
     let i = Atomic.fetch_and_add t.head 1 in
     if i >= Array.length t.buffer
     then Is_empty
-    else
-      let cell = Array.unsafe_get t.buffer i in
-      if i >= Atomic.get t.tail
-      then tombstone t cell
-      else spinlock ~retry:spinlock_iterations t cell
-
-and tombstone t cell =
-  if Atomic.compare_and_set cell Empty Tombstone
-  then pop_s t
-  else begin match Atomic.get cell with
-       | (Value v) as old ->
-           assert (Atomic.compare_and_set cell old Tombstone) ;
-           Pop v
-       | _ -> assert false
-       end
-
-and spinlock ~retry t cell =
-  match Atomic.get cell with
-  | (Value v) as old ->
-      assert (Atomic.compare_and_set cell old Tombstone) ;
+    else if mark t i
+    then pop_s ~dummy t
+    else begin
+      let v = t.buffer.(i) in
+      t.buffer.(i) <- dummy ;
       Pop v
-  | Empty when retry <= 0 ->
-      tombstone t cell
-  | Empty ->
-      Domain.cpu_relax () ;
-      spinlock ~retry:(retry - 1) t cell
-  | Tombstone ->
-      assert false
+    end
 
 let rec pop t =
   let first = Atomic.get t.first in
-  match pop_s first with
+  match pop_s ~dummy:t.dummy first with
   | Pop v -> Some v
   | Wait_for_it -> None
   | Is_empty ->

src/mpmc_queue.mli

@@ -3,13 +3,15 @@
 type 'a t
 (** A queue of items of type ['a]. *)
 
-val make : ?capacity:int -> unit -> 'a t
-(** [make ()] creates a new empty queue.
+val make : ?capacity:int -> dummy:'a -> unit -> 'a t
+(** [make ~dummy ()] creates a new empty queue.
 
-    The optional parameter [?capacity] defaults to 4096 and is used to size the
-    internal buffers of the queue: Choosing a small number lower the pause
-    durations caused by allocations, but a larger capacity provides overall
-    faster operations. *)
+    - The [dummy] element is a placeholder for ['a] values.
+    - The optional parameter [?capacity] defaults to 512 and is used to size the
+      internal buffers of the queue: Choosing a small number lower the pause
+      durations caused by allocations, but a larger capacity can provide overall
+      faster operations.
+*)
 
 val push : 'a t -> 'a -> unit
 (** [push t x] adds [x] to the tail of the queue. *)

test/mpmc_queue/lin_mpmc_queue.ml

@@ -3,7 +3,7 @@ module Q = Lockfree.Mpmc_queue
 module Test = struct
   type t = int Q.t
 
-  let init () = Q.make ~capacity:1 ()
+  let init () = Q.make ~capacity:1 ~dummy:(-1) ()
 
   let cleanup _ = ()