Rendezvous and demonstration on Spsc_queue API

bench/bench_spsc_queue.ml

@@ -1,24 +1,20 @@
 open Lockfree
 
 let item_count = 2_000_000
-
-let rec try_until_success f =
-  try f () with Spsc_queue.Full -> try_until_success f
+let rdv = Rendezvous.semaphore_unit
 
 let run () =
-  let queue = Spsc_queue.create ~size_exponent:3 in
+  let queue = Spsc_queue.create ~size_exponent:10 in
   let pusher =
     Domain.spawn (fun () ->
         let start_time = Unix.gettimeofday () in
         for i = 1 to item_count do
-          try_until_success (fun () -> Spsc_queue.push queue i)
+          Spsc_queue.push ~rdv queue i
         done;
         start_time)
   in
-  for _ = 1 to item_count do
-    while Option.is_none (Spsc_queue.pop queue) do
-      ()
-    done
+  for i = 1 to item_count do
+    assert (Spsc_queue.pop ~rdv queue = i)
   done;
   let end_time = Unix.gettimeofday () in
   let start_time = Domain.join pusher in

src/lockfree.ml

@@ -25,6 +25,7 @@
 Copyright (c) 2017, Nicolas ASSOUAD <nicolas.assouad@ens.fr>
 ########
 *)
+module Rendezvous = Rendezvous
 module Ws_deque = Ws_deque
 module Spsc_queue = Spsc_queue
 module Mpsc_queue = Mpsc_queue

src/lockfree.mli

@@ -30,6 +30,7 @@ Copyright (c) 2017, Nicolas ASSOUAD <nicolas.assouad@ens.fr>
 ########
 *)
 
+module Rendezvous = Rendezvous
 module Ws_deque = Ws_deque
 module Spsc_queue = Spsc_queue
 module Mpsc_queue = Mpsc_queue

src/rendezvous.ml

@@ -0,0 +1,67 @@
+type 'a t = unit -> ('a -> unit) * (unit -> 'a)
+
+let fail () = failwith "Rendezvous: double send"
+
+let semaphore () =
+  let r = ref None in
+  let s = Semaphore.Binary.make false in
+  let send v =
+    r := Some v;
+    Semaphore.Binary.release s
+  and recv () =
+    Semaphore.Binary.acquire s;
+    Option.get !r
+  in
+  (send, recv)
+
+let semaphore_unit () =
+  let s = Semaphore.Binary.make false in
+  let send () = Semaphore.Binary.release s
+  and recv () = Semaphore.Binary.acquire s in
+  (send, recv)
+
+let backoff ?min_wait ?max_wait () =
+  let r = Atomic.make None in
+  let backoff = Backoff.create ?min_wait ?max_wait () in
+  let rec send v = if not (Atomic.compare_and_set r None (Some v)) then fail ()
+  and recv () =
+    match Atomic.get r with
+    | Some v -> v
+    | None ->
+        Backoff.once backoff;
+        recv ()
+  in
+  (send, recv)
+
+let backoff_unit ?min_wait ?max_wait () =
+  let r = Atomic.make false in
+  let backoff = Backoff.create ?min_wait ?max_wait () in
+  let rec send () = if not (Atomic.compare_and_set r false true) then fail ()
+  and recv () =
+    if not (Atomic.get r) then (
+      Backoff.once backoff;
+      recv ())
+  in
+  (send, recv)
+
+let spinlock () =
+  let r = Atomic.make None in
+  let rec send v = if not (Atomic.compare_and_set r None (Some v)) then fail ()
+  and recv () =
+    match Atomic.get r with
+    | Some v -> v
+    | None ->
+        Domain.cpu_relax ();
+        recv ()
+  in
+  (send, recv)
+
+let spinlock_unit () =
+  let r = Atomic.make false in
+  let rec send () = if not (Atomic.compare_and_set r false true) then fail ()
+  and recv () =
+    if not (Atomic.get r) then (
+      Domain.cpu_relax ();
+      recv ())
+  in
+  (send, recv)

src/rendezvous.mli

@@ -0,0 +1,43 @@
+(** One-shot synchronization between two domains.
+
+When using a lockfree datastructure with a library providing fibers, like
+[domainslib] or [eio], the rendezvous should be implemented in terms of the
+blocking primitives provided by the scheduler. *)
+
+type 'a t = unit -> ('a -> unit) * (unit -> 'a)
+(** The type to create a rendezvous allowing two domains to synchronize
+    on a future ['a] value.
+
+    Given a rendezvous [rdv : 'a t], typical usage follows:
+
+    {[
+      let send, recv = rdv () in
+      Atomic.set release send; (* publish [send] somehere accessible to the other domain *)
+      let v = recv () in (* block until [send v] has been called and returns [v] *)
+    }]
+
+    The function [send] and [recv] must be called at most once.
+*)
+
+val semaphore : 'a t
+(** [semaphore] is a rendezvous that uses a semaphore to suspend the domain.
+    Recommended when using raw domains without fibers.  *)
+
+val semaphore_unit : unit t
+(** Same as [semaphore], specialized for [unit]. *)
+
+val backoff : ?min_wait:int -> ?max_wait:int -> 'a t
+(** [backoff] is a rendezvous that locks by spinning, performing a {! Backoff}
+    on each failure to make progress.
+    Recommended when using raw domains without fibers, if the operation is
+    expected to complete soon. *)
+
+val backoff_unit : ?min_wait:int -> ?max_wait:int -> unit t
+(** Same as [backoff], specialized for [unit]. *)
+
+val spinlock : 'a t
+(** [spinlock] is a rendezvous that locks by spinning, requiring no syscalls.
+    Not recommended as it can have terrible performances. *)
+
+val spinlock_unit : unit t
+(** Same as [spinlock], specialized for [unit]. *)

src/spsc_queue.ml

@@ -22,10 +22,11 @@
  *)
 
 type 'a t = {
-  array : 'a Option.t Array.t;
+  array : 'a option array;
   tail : int Atomic.t;
   head : int Atomic.t;
   mask : int;
+  lock : (unit -> unit) option Atomic.t;
 }
 
 exception Full
@@ -36,29 +37,61 @@ let create ~size_exponent =
     head = Atomic.make 0;
     tail = Atomic.make 0;
     mask = size - 1;
-    array = Array.init size (fun _ -> None);
+    array = Array.make size None;
+    lock = Atomic.make None;
   }
 
-let push { array; head; tail; mask; _ } element =
+let unlock lock =
+  if Atomic.get lock <> None then
+    match Atomic.exchange lock None with
+    | Some release -> release ()
+    | None -> ()
+
+let try_push { array; head; tail; mask; lock } element =
   let size = mask + 1 in
-  let head_val = Atomic.get head in
   let tail_val = Atomic.get tail in
-  if head_val + size == tail_val then raise Full
+  let head_val = Atomic.get head in
+  if head_val + size = tail_val then false
   else (
     Array.set array (tail_val land mask) (Some element);
-    Atomic.set tail (tail_val + 1))
+    Atomic.set tail (tail_val + 1);
+    unlock lock;
+    true)
 
-let pop { array; head; tail; mask; _ } =
+let try_pop { array; head; tail; mask; lock } =
   let head_val = Atomic.get head in
   let tail_val = Atomic.get tail in
-  if head_val == tail_val then None
+  if head_val = tail_val then None
   else
     let index = head_val land mask in
     let v = Array.get array index in
     (* allow gc to collect it *)
     Array.set array index None;
     Atomic.set head (head_val + 1);
     assert (Option.is_some v);
+    unlock lock;
     v
 
 let size { head; tail; _ } = Atomic.get tail - Atomic.get head
+
+let wait ~rdv t expected_size =
+  let release, wait = rdv () in
+  let some_release = Some release in
+  if Atomic.compare_and_set t.lock None some_release then
+    if size t = expected_size then wait ()
+    else
+      let (_ : bool) = Atomic.compare_and_set t.lock some_release None in
+      ()
+  else unlock t.lock
+
+let rec push ~rdv t element =
+  if not (try_push t element) then (
+    wait ~rdv t (t.mask + 1);
+    push ~rdv t element)
+
+let rec pop ~rdv t =
+  match try_pop t with
+  | Some v -> v
+  | None ->
+      wait ~rdv t 0;
+      pop ~rdv t

src/spsc_queue.mli

@@ -8,15 +8,20 @@ val create : size_exponent:int -> 'a t
 (** [create ~size_exponent:int] creates a empty queue of size
    [2^size_exponent].  *)
 
-val push : 'a t -> 'a -> unit
-(** [push q v] pushes [v] at the back of the queue. 
-
-   @raise [Full] if the queue is full.
-*)
+val try_push : 'a t -> 'a -> bool
+(** [try_push q v] attempts to push [v] at the back of the queue
+    and returns [true] on success, or [false] if the queue is full. *)
 
-val pop : 'a t -> 'a option
-(** [pop q] removes element from head of the queue, if any. This method can be used by
-  at most 1 thread at the time. *)
+val push : rdv:unit Rendezvous.t -> 'a t -> 'a -> unit
+(** [push ~rdv q v] pushes [v] at the back of the queue, potentially
+    blocking using [rdv] while the queue is full. *)
+
+val try_pop : 'a t -> 'a option
+(** [try_pop q] removes an element from the head of the queue, if any. *)
+
+val pop : rdv:unit Rendezvous.t -> 'a t -> 'a
+(** [pop q] removes an element from the head of the queue, potentially
+    blocking using [rdv] while the queue is empty. *)
 
 val size : 'a t -> int
 (** [size] returns the size of the queue. This method linearizes only when called

test/spsc_queue/qcheck_spsc_queue.ml

@@ -7,7 +7,7 @@ let pop_n_times q n =
   let rec loop count acc =
     if count = 0 then acc
     else
-      let v = Spsc_queue.pop q in
+      let v = Spsc_queue.try_pop q in
       Domain.cpu_relax ();
       loop (count - 1) (v :: acc)
   in
@@ -33,12 +33,7 @@ let tests =
 
           (* Sequential pushed : not Full exception should be
              raised. *)
-          let not_full_queue =
-            try
-              List.iter (Spsc_queue.push q) l;
-              true
-            with Spsc_queue.Full -> false
-          in
+          let not_full_queue = List.for_all (Spsc_queue.try_push q) l in
 
           (* Consumer domain pops *)
           let consumer = Domain.spawn (fun () -> pop_n_times q npop) in
@@ -63,7 +58,7 @@ let tests =
 
           (* Initialization *)
           let q = Spsc_queue.create ~size_exponent in
-          List.iter (Spsc_queue.push q) l;
+          List.iter (fun x -> assert (Spsc_queue.try_push q x)) l;
 
           (* Consumer pops *)
           let sema = Semaphore.Binary.make false in
@@ -82,7 +77,7 @@ let tests =
                 (* Main domain pushes.*)
                 List.iter
                   (fun elt ->
-                    Spsc_queue.push q elt;
+                    assert (Spsc_queue.try_push q elt);
                     Domain.cpu_relax ())
                   l')
           in
@@ -103,12 +98,7 @@ let tests =
 
           (* Initialization *)
           let q = Spsc_queue.create ~size_exponent in
-          let is_full =
-            try
-              List.iter (Spsc_queue.push q) l;
-              false
-            with Spsc_queue.Full -> true
-          in
+          let is_full = not (List.for_all (Spsc_queue.try_push q) l) in
 
           (* Property *)
           (List.length l > size_max && is_full)

test/spsc_queue/spsc_queue_dscheck.ml

@@ -4,21 +4,21 @@ let create_test ~shift_by () =
 
   (* shift the queue, that helps testing overlap handling *)
   for _ = 1 to shift_by do
-    Spsc_queue.push queue (-1);
-    assert (Option.is_some (Spsc_queue.pop queue))
+    assert (Spsc_queue.try_push queue (-1));
+    assert (Option.is_some (Spsc_queue.try_pop queue))
   done;
 
   (* enqueuer *)
   Atomic.spawn (fun () ->
       for i = 1 to items_count do
-        Spsc_queue.push queue i
+        assert (Spsc_queue.try_push queue i)
       done);
 
   (* dequeuer *)
   let dequeued = ref 0 in
   Atomic.spawn (fun () ->
       for _ = 1 to items_count + 1 do
-        match Spsc_queue.pop queue with
+        match Spsc_queue.try_pop queue with
         | None -> ()
         | Some v ->
             assert (v = !dequeued + 1);
@@ -34,17 +34,17 @@ let with_trace ?(shift_by = 0) f () = Atomic.trace (fun () -> f ~shift_by ())
 let size_linearizes_with_1_thr () =
   Atomic.trace (fun () ->
       let queue = Spsc_queue.create ~size_exponent:4 in
-      Spsc_queue.push queue (-1);
-      Spsc_queue.push queue (-1);
+      assert (Spsc_queue.try_push queue (-1));
+      assert (Spsc_queue.try_push queue (-1));
 
       Atomic.spawn (fun () ->
           for _ = 1 to 4 do
-            Spsc_queue.push queue (-1)
+            assert (Spsc_queue.try_push queue (-1))
           done);
 
       let size = ref 0 in
       Atomic.spawn (fun () ->
-          assert (Option.is_some (Spsc_queue.pop queue));
+          assert (Option.is_some (Spsc_queue.try_pop queue));
           size := Spsc_queue.size queue);
 
       Atomic.final (fun () -> Atomic.check (fun () -> 1 <= !size && !size <= 5)))