Add a poisonable Barrier to picos_std.sync

Author: polytypic

lib/picos_std.sync/barrier.ml

@@ -0,0 +1,74 @@
+open Picos_std_awaitable
+
+type t = int Awaitable.t
+
+let sense_bit = 1
+let parties_shift = 1
+let parties_bits = (Sys.int_size - 2) / 2
+let max_parties = (1 lsl parties_bits) - 1
+let parties_mask = max_parties lsl parties_shift
+let awaiting_shift = parties_shift + parties_bits
+let awaiting_one = 1 lsl awaiting_shift
+let poisoned_bit = Int.min_int
+
+let create ?padded parties =
+  if parties <= 0 || max_parties < parties then
+    invalid_arg "invalid number of parties";
+  Awaitable.make ?padded
+    ((parties lsl parties_shift) lor (parties lsl awaiting_shift))
+
+exception Poisoned
+
+let rec poison t =
+  let before = Awaitable.get t in
+  if 0 < before then
+    let after =
+      let parties_shifted = before land parties_mask in
+      let after_sense = lnot before land sense_bit in
+      parties_shifted lor after_sense
+      lor (max_parties lsl awaiting_shift)
+      lor poisoned_bit
+    in
+    if Awaitable.compare_and_set t before after then Awaitable.broadcast t
+    else poison t
+
+let await t =
+  let[@inline never] poison_and_raise t =
+    poison t;
+    raise Poisoned
+  in
+  let prior = Awaitable.fetch_and_add t (-awaiting_one) in
+  if awaiting_one < prior then begin
+    let before = prior - awaiting_one in
+    let after_sense = prior land sense_bit lxor sense_bit in
+    if before < awaiting_one then
+      let after =
+        let parties_shifted = before land parties_mask in
+        parties_shifted lor after_sense
+        lor (parties_shifted lsl (awaiting_shift - parties_shift))
+      in
+      if Awaitable.compare_and_set t before after then Awaitable.broadcast t
+      else
+        (* The barrier is being misused?  Poison the barrier. *)
+        poison_and_raise t
+    else
+      let state = ref before in
+      match
+        while !state land sense_bit <> after_sense do
+          Awaitable.await t !state;
+          state := Awaitable.get t
+        done
+      with
+      | () -> if 0 <= !state then () else raise Poisoned
+      | exception exn ->
+          let bt = Printexc.get_raw_backtrace () in
+          poison t;
+          Printexc.raise_with_backtrace exn bt
+  end
+  else begin
+    (* The barrier was poisoned.  Undo the decrement. *)
+    Awaitable.fetch_and_add t awaiting_one |> ignore;
+    poison_and_raise t
+  end
+
+let parties t = (Awaitable.get t land parties_mask) lsr parties_shift

lib/picos_std.sync/picos_std_sync.ml

@@ -6,5 +6,6 @@ module Rwlock = Rwlock
 module Sem = Sem
 module Lazy = Lazy
 module Latch = Latch
+module Barrier = Barrier
 module Ivar = Ivar
 module Stream = Stream

lib/picos_std.sync/picos_std_sync.mli

@@ -481,6 +481,42 @@ module Sem : sig
       {{!poison} poisoned}. *)
 end
 
+module Barrier : sig
+  (** A poisonable barrier. *)
+
+  type t
+  (** Represents a poisonable barrier. *)
+
+  val max_parties : int
+  (** Maximum number of participants that a barrier can be configured with. *)
+
+  val create : ?padded:bool -> int -> t
+  (** [create parties] creates a new barrier for given number of parties.
+
+      @raise Invalid_argument
+        in case the given number of [parties] is less than [1] or greater than
+        [max_parties]. *)
+
+  val parties : t -> int
+  (** [paries barrier] returns the number of parties the barrier was
+      {{!create} created} with. *)
+
+  exception Poisoned
+  (** Exception raised by {!await} in case the barrier has become poisoned. *)
+
+  val await : t -> unit
+  (** [await barrier] awaits until the configured number of {!parties} has
+      called [await] on the barrier and returns, or raises {!Poisoned} in case
+      the barrier has become poisoned.
+
+      ℹ️ If the await is canceled, then [await] will {!poison} the barrier before
+      reraising the cancelation exception. *)
+
+  val poison : t -> unit
+  (** [poison barrier] marks the barrier as poisoned. Concurrent and subsequent
+      calls of {!await} will raise the {!Poisoned} exception. *)
+end
+
 module Lazy : sig
   (** A lazy suspension.
 

test/dune

@@ -51,6 +51,8 @@
    (run %{test} -- "^Lazy$" 0)
    (run %{test} -- "^Lazy$" 1)
    (run %{test} -- "^Event$" 0)
+   (run %{test} -- "^Barrier$" 0)
+   (run %{test} -- "^Barrier$" 1)
    (run %{test} -- "^Non-cancelable ops$" 0)
    ;;
    )))

test/test_sync.ml

@@ -607,6 +607,73 @@ end
 
 module Rwlock_is_a_submodule_of_Lock : module type of Lock = Rwlock
 
+let test_barrier_basics () =
+  for n = 1 to 4 do
+    let barrier = Barrier.create n in
+    Test_scheduler.run ~max_domains:(n + 1) @@ fun () ->
+    Flock.join_after @@ fun () ->
+    let n_outside = Atomic.make n in
+    for _ = 1 to n do
+      Flock.fork @@ fun () ->
+      for _ = 1 to 5 do
+        Atomic.decr n_outside;
+        Barrier.await barrier;
+        assert (Atomic.get n_outside = 0);
+        Barrier.await barrier;
+        Atomic.incr n_outside
+      done
+    done
+  done
+
+let test_barrier_poisoning () =
+  for n = 2 to 5 do
+    match
+      Test_scheduler.run ~max_domains:n @@ fun () ->
+      let barrier = Barrier.create n in
+      Flock.join_after @@ fun () ->
+      for _ = 2 to n do
+        Flock.fork @@ fun () ->
+        Barrier.await barrier;
+        assert false
+      done;
+      Barrier.poison barrier
+    with
+    | () -> assert false
+    | exception Barrier.Poisoned -> ()
+    | exception Control.Errors exns ->
+        assert (
+          List.for_all
+            (function Barrier.Poisoned, _ -> true | _ -> false)
+            exns);
+        assert (List.length exns <= n - 1)
+  done;
+  for n = 3 to 6 do
+    match
+      Test_scheduler.run ~max_domains:n @@ fun () ->
+      let barrier = Barrier.create n in
+      Flock.join_after @@ fun () ->
+      let to_be_terminated =
+        Flock.fork_as_promise @@ fun () ->
+        Barrier.await barrier;
+        assert false
+      in
+      for _ = 3 to n do
+        Flock.fork @@ fun () ->
+        Barrier.await barrier;
+        assert false
+      done;
+      Promise.terminate to_be_terminated
+    with
+    | () -> assert false
+    | exception Barrier.Poisoned -> ()
+    | exception Control.Errors exns ->
+        assert (
+          List.for_all
+            (function Barrier.Poisoned, _ -> true | _ -> false)
+            exns);
+        assert (List.length exns <= n - 2)
+  done
+
 let () =
   try
     [
@@ -664,6 +731,11 @@ let () =
           Alcotest.test_case "cancelation" `Quick test_lazy_cancelation;
         ] );
       ("Event", [ Alcotest.test_case "basics" `Quick test_event_basics ]);
+      ( "Barrier",
+        [
+          Alcotest.test_case "basics" `Quick test_barrier_basics;
+          Alcotest.test_case "poisoning" `Quick test_barrier_poisoning;
+        ] );
       ( "Non-cancelable ops",
         [ Alcotest.test_case "are not canceled" `Quick test_non_cancelable_ops ]
       );