Add a poisonable Barrier to picos_std.sync

lib/picos_std.sync/barrier.ml

@@ -0,0 +1,50 @@
+open Picos_std_awaitable
+
+type t = int Awaitable.t
+
+let parties_bits = (Sys.int_size - 2) / 2
+let max_parties = (1 lsl parties_bits) - 1
+let parties_mask = max_parties
+let sense_bit = 1 lsl parties_bits
+let awaiting_shift = 1 + 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 lor (parties lsl awaiting_shift))
+
+exception Poisoned
+
+let await t =
+  let before = Awaitable.fetch_and_add t (-awaiting_one) - awaiting_one in
+  let after_sense = lnot before land sense_bit in
+  if before < awaiting_one then
+    let after =
+      let parties = before land parties_mask in
+      parties lor after_sense lor (parties lsl awaiting_shift)
+    in
+    if Awaitable.compare_and_set t (before land lnot poisoned_bit) after then
+      Awaitable.broadcast t
+    else
+      let _ : int = Awaitable.fetch_and_add t awaiting_one in
+      raise Poisoned
+  else
+    let state = ref before in
+    while !state land sense_bit <> after_sense do
+      Awaitable.await t !state;
+      state := Awaitable.get t
+    done;
+    if 0 <= !state then () else raise Poisoned
+
+let rec poison t =
+  let before = Awaitable.get t in
+  if 0 < before then
+    let after =
+      let parties = before land parties_mask in
+      let after_sense = lnot before land sense_bit in
+      parties lor after_sense lor (parties lsl awaiting_shift) lor poisoned_bit
+    in
+    if Awaitable.compare_and_set t before after then Awaitable.broadcast t
+    else poison t

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,28 @@ module Sem : sig
       {{!poison} poisoned}. *)
 end
 
+module Barrier : sig
+  (** *)
+
+  type t
+  (** *)
+
+  val max_parties : int
+  (** *)
+
+  val create : ?padded:bool -> int -> t
+  (** *)
+
+  exception Poisoned
+  (** *)
+
+  val await : t -> unit
+  (** *)
+
+  val poison : t -> unit
+  (** *)
+end
+
 module Lazy : sig
   (** A lazy suspension.
 

test/dune

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

test/test_sync.ml

@@ -607,6 +607,24 @@ 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 () =
   try
     [
@@ -664,6 +682,7 @@ 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 ]);
       ( "Non-cancelable ops",
         [ Alcotest.test_case "are not canceled" `Quick test_non_cancelable_ops ]
       );