Repository: ocaml-multicore/picos
Pull request #44 (2024-03-04 14:20:38)

Allow to specify order in which triggers are signaled on completion
The order has an effect on scheduling so it makes to allow it to be specified.


Patch:
diff --git a/lib/picos/bootstrap.ml b/lib/picos/bootstrap.ml
index d93b671..d931167 100644
--- a/lib/picos/bootstrap.ml
+++ b/lib/picos/bootstrap.ml
@@ -34,11 +34,16 @@ module Computation = struct
   type 'a state =
     | Canceled of Exn_bt.t
     | Returned of 'a
-    | Continue of { balance : int; triggers : Trigger.t list }
+    | Continue of { balance_and_mode : int; triggers : Trigger.t list }
 
   type 'a t = 'a state Atomic.t
 
-  let create () = Atomic.make (Continue { balance = 0; triggers = [] })
+  let fifo_bit = 1
+  let one = 2
+
+  let create ?(mode : [ `FIFO | `LIFO ] = `FIFO) () =
+    let balance_and_mode = Bool.to_int (mode == `FIFO) in
+    Atomic.make (Continue { balance_and_mode; triggers = [] })
 
   let canceled t =
     match Atomic.get t with
@@ -50,13 +55,16 @@ module Computation = struct
         [detach].  This ensures that the [O(n)] lazy removal done by [gc] cannot
         cause starvation, because the only reason that CAS fails after [gc] is
         that someone else completed the [gc]. *)
-    let rec gc balance triggers = function
+    let rec gc balance_and_mode triggers = function
       | [] ->
-          let triggers = if balance <= 1 then triggers else List.rev triggers in
-          Continue { balance; triggers }
+          let triggers =
+            if balance_and_mode <= one + fifo_bit then triggers
+            else List.rev triggers
+          in
+          Continue { balance_and_mode; triggers }
       | r :: rs ->
-          if Trigger.is_signaled r then gc balance triggers rs
-          else gc (balance + 1) (r :: triggers) rs
+          if Trigger.is_signaled r then gc balance_and_mode triggers rs
+          else gc (balance_and_mode + one) (r :: triggers) rs
   end
 
   let rec try_attach t trigger backoff =
@@ -67,10 +75,14 @@ module Computation = struct
         (not (Trigger.is_signaled trigger))
         &&
         let after =
-          if 0 <= r.balance then
+          if fifo_bit <= r.balance_and_mode then
             Continue
-              { balance = r.balance + 1; triggers = trigger :: r.triggers }
-          else gc 1 [ trigger ] r.triggers
+              {
+                balance_and_mode = r.balance_and_mode + one;
+                triggers = trigger :: r.triggers;
+              }
+          else
+            gc (one + (r.balance_and_mode land fifo_bit)) [ trigger ] r.triggers
         in
         Atomic.compare_and_set t before after
         || try_attach t trigger (Backoff.once backoff)
@@ -82,8 +94,10 @@ module Computation = struct
     | Returned _ | Canceled _ -> ()
     | Continue r as before ->
         let after =
-          if 0 <= r.balance then Continue { r with balance = r.balance - 2 }
-          else gc 0 [] r.triggers
+          if fifo_bit <= r.balance_and_mode then
+            Continue
+              { r with balance_and_mode = r.balance_and_mode - (2 * one) }
+          else gc (r.balance_and_mode land fifo_bit) [] r.triggers
         in
         if not (Atomic.compare_and_set t before after) then
           detach t (Backoff.once backoff)
@@ -105,7 +119,10 @@ module Computation = struct
       | Returned _ | Canceled _ -> false
       | Continue r as before ->
           if Atomic.compare_and_set t before after then begin
-            List.iter Trigger.signal r.triggers;
+            List.iter Trigger.signal
+              (if r.balance_and_mode land fifo_bit = fifo_bit then
+                 List.rev r.triggers
+               else r.triggers);
             true
           end
           else try_terminate t after (Backoff.once backoff)
diff --git a/lib/picos/picos.mli b/lib/picos/picos.mli
index fc9b254..77fd9ab 100644
--- a/lib/picos/picos.mli
+++ b/lib/picos/picos.mli
@@ -567,8 +567,15 @@ module Computation : sig
   type !'a t
   (** Represents a cancelable computation. *)
 
-  val create : unit -> 'a t
-  (** [create ()] creates a new computation in the running state. *)
+  val create : ?mode:[ `FIFO | `LIFO ] -> unit -> 'a t
+  (** [create ()] creates a new computation in the running state.
+
+      The optional [mode] specifies the order in which {{!Trigger} triggers}
+      {{!try_attach} attached} to the computation will be {{!Trigger.signal}
+      signaled} after the computation has been completed.  [`FIFO] ordering may
+      reduce latency of IO bound computations and is the default.  [`LIFO] may
+      improve thruput of CPU bound computations and be preferable on a
+      work-stealing scheduler, for example. *)
 
   val finished : unit t
   (** [finished] is a constant finished computation. *)


Test patch:
diff --git a/test/lib/elements/promise.ml b/test/lib/elements/promise.ml
index fefd8ae..aae5933 100644
--- a/test/lib/elements/promise.ml
+++ b/test/lib/elements/promise.ml
@@ -98,7 +98,7 @@ let any xs =
   Fiber.spawn ~forbid:false y [ main ];
   y
 
-let create = Computation.create
+let create () = Computation.create ()
 let try_return_to = Computation.try_return
 let try_reify_to t thunk = Computation.try_capture t thunk ()
 let publish = of_computation
diff --git a/test/lib/elements/sleep.ml b/test/lib/elements/sleep.ml
index 282d6ad..5d8977c 100644
--- a/test/lib/elements/sleep.ml
+++ b/test/lib/elements/sleep.ml
@@ -3,7 +3,7 @@ open Picos
 let exit_bt = Exn_bt.get_callstack 0 Exit
 
 let sleepf seconds =
-  let sleep = Computation.create () in
+  let sleep = Computation.create ~mode:`LIFO () in
   Computation.cancel_after ~seconds sleep exit_bt;
   let trigger = Trigger.create () in
   if Computation.try_attach sleep trigger then
diff --git a/test/test_picos.ml b/test/test_picos.ml
index 4d73b2b..ee83e18 100644
--- a/test/test_picos.ml
+++ b/test/test_picos.ml
@@ -110,48 +110,55 @@ let test_cancel_after () =
     (Exn_bt.get_callstack 0 Not_found);
   Computation.await computation
 
-let test_computation_completion_signals_triggers_in_lifo_order () =
-  let state = Random.State.make_self_init () in
-  let num_non_trivial = ref 0 in
-  for _ = 1 to 10 do
-    let computation = Computation.create () in
-    let signals = ref [] in
-    let triggers = ref [] in
-    let counter = ref 0 in
-    let attach_one () =
-      let trigger = Trigger.create () in
-      triggers := trigger :: !triggers;
-      let i = !counter in
-      counter := i + 1;
-      assert (Computation.try_attach computation trigger);
-      assert (
-        Trigger.on_signal trigger () () (fun _ _ _ -> signals := i :: !signals))
-    in
-    let detach_one () =
-      let n = List.length !triggers in
-      if 0 < n then begin
-        let bits = Random.State.bits state in
-        let i = bits mod n in
-        let trigger = List.nth !triggers i in
-        triggers := List.filter (( != ) trigger) !triggers;
-        Computation.detach computation trigger
-      end
-    in
-    for _ = 1 to 10 do
-      for _ = 1 to 10 do
-        let bits = Random.State.bits state in
-        if bits land 3 <= 2 then attach_one () else detach_one ()
-      done;
-      for _ = 1 to List.length !triggers / 3 do
-        detach_one ()
-      done
-    done;
-    if List.length !triggers >= 2 then incr num_non_trivial;
-    signals := [];
-    Computation.finish computation;
-    assert (!signals = List.sort Int.compare !signals)
-  done;
-  assert (0 < !num_non_trivial)
+let test_computation_completion_signals_triggers_in_order () =
+  [ `FIFO; `LIFO ]
+  |> List.iter @@ fun mode ->
+     let state = Random.State.make_self_init () in
+     let num_non_trivial = ref 0 in
+     for _ = 1 to 10 do
+       let computation = Computation.create ~mode () in
+       let signals = ref [] in
+       let triggers = ref [] in
+       let counter = ref 0 in
+       let attach_one () =
+         let trigger = Trigger.create () in
+         triggers := trigger :: !triggers;
+         let i = !counter in
+         counter := i + 1;
+         assert (Computation.try_attach computation trigger);
+         assert (
+           Trigger.on_signal trigger () () (fun _ _ _ ->
+               signals := i :: !signals))
+       in
+       let detach_one () =
+         let n = List.length !triggers in
+         if 0 < n then begin
+           let bits = Random.State.bits state in
+           let i = bits mod n in
+           let trigger = List.nth !triggers i in
+           triggers := List.filter (( != ) trigger) !triggers;
+           Computation.detach computation trigger
+         end
+       in
+       for _ = 1 to 10 do
+         for _ = 1 to 10 do
+           let bits = Random.State.bits state in
+           if bits land 3 <= 2 then attach_one () else detach_one ()
+         done;
+         for _ = 1 to List.length !triggers / 3 do
+           detach_one ()
+         done
+       done;
+       if List.length !triggers >= 2 then incr num_non_trivial;
+       signals := [];
+       Computation.finish computation;
+       let expected =
+         List.sort Int.compare !signals
+         |> if mode = `FIFO then List.rev else Fun.id
+       in
+       assert (!signals = expected)
+     done;
+     assert (0 < !num_non_trivial)
 
 let () =
   [
@@ -163,10 +170,10 @@ let () =
     ( "Thread cancelation",
       [ Alcotest.test_case "" `Quick test_thread_cancelation ] );
     ("Cancel after", [ Alcotest.test_case "" `Quick test_cancel_after ]);
-    ( "Computation signals in LIFO order",
+    ( "Computation signals in order",
       [
         Alcotest.test_case "" `Quick
-          test_computation_completion_signals_triggers_in_lifo_order;
+          test_computation_completion_signals_triggers_in_order;
       ] );
   ]
   |> Alcotest.run "Picos"
diff --git a/test/test_picos_dscheck.ml b/test/test_picos_dscheck.ml
index 964a8a0..dee2b52 100644
--- a/test/test_picos_dscheck.ml
+++ b/test/test_picos_dscheck.ml
@@ -43,81 +43,88 @@ let test_trigger_contract () =
 
 (** This tries to cover much of the public contract of [Computation]s. *)
 let test_computation_contract () =
-  let attached_total = ref 0 and unattached_total = ref 0 in
-  let () =
-    Atomic.trace @@ fun () ->
-    let computation = Computation.create () in
-    let returns = ref 0 and cancels = ref 0 in
-    let () =
-      Atomic.spawn @@ fun () ->
-      if Computation.try_return computation 101 then incr returns
-    in
-    let () =
-      Atomic.spawn @@ fun () ->
-      if Computation.try_cancel computation (Exn_bt.get_callstack 1 Exit) then
-        incr cancels
-    in
-    let triggers = Array.init 2 @@ fun _ -> Trigger.create () in
-    let attached = ref 0 and unattached = ref 0 in
-    let () =
-      triggers
-      |> Array.iter @@ fun trigger ->
+  [ `FIFO; `LIFO ]
+  |> List.iter @@ fun mode ->
+     let attached_total = ref 0 and unattached_total = ref 0 in
+     let () =
+       Atomic.trace @@ fun () ->
+       let computation = Computation.create ~mode () in
+       let returns = ref 0 and cancels = ref 0 in
+       let () =
          Atomic.spawn @@ fun () ->
-         if Computation.try_attach computation trigger then incr attached
-         else incr unattached
-    in
-    Atomic.final @@ fun () ->
-    Atomic.check @@ fun () ->
-    attached_total += !attached;
-    unattached_total += !unattached;
-    begin
-      match Computation.peek computation with
-      | Some (Ok 101) when !returns = 1 && !cancels = 0 -> true
-      | Some (Error { exn = Exit; _ }) when !returns = 0 && !cancels = 1 -> true
-      | _ -> false
-    end
-    && !attached + !unattached = Array.length triggers
-    && !attached
-       = sum_as
-           (fun trigger -> Bool.to_int (Trigger.is_signaled trigger))
-           triggers
-  in
-  [ attached_total; unattached_total ]
-  |> List.iter @@ fun total -> if !total = 0 then Alcotest.fail "uncovered case"
+         if Computation.try_return computation 101 then incr returns
+       in
+       let () =
+         Atomic.spawn @@ fun () ->
+         if Computation.try_cancel computation (Exn_bt.get_callstack 1 Exit)
+         then incr cancels
+       in
+       let triggers = Array.init 2 @@ fun _ -> Trigger.create () in
+       let attached = ref 0 and unattached = ref 0 in
+       let () =
+         triggers
+         |> Array.iter @@ fun trigger ->
+            Atomic.spawn @@ fun () ->
+            if Computation.try_attach computation trigger then incr attached
+            else incr unattached
+       in
+       Atomic.final @@ fun () ->
+       Atomic.check @@ fun () ->
+       attached_total += !attached;
+       unattached_total += !unattached;
+       begin
+         match Computation.peek computation with
+         | Some (Ok 101) when !returns = 1 && !cancels = 0 -> true
+         | Some (Error { exn = Exit; _ }) when !returns = 0 && !cancels = 1 ->
+             true
+         | _ -> false
+       end
+       && !attached + !unattached = Array.length triggers
+       && !attached
+          = sum_as
+              (fun trigger -> Bool.to_int (Trigger.is_signaled trigger))
+              triggers
+     in
+     [ attached_total; unattached_total ]
+     |> List.iter @@ fun total ->
+        if !total = 0 then Alcotest.fail "uncovered case"
 
 (** This covers the contract of [Computation] to remove detached triggers.
 
     Testing this through the public API would require relying on GC
     statistics. *)
 let test_computation_removes_triggers () =
-  Atomic.trace @@ fun () ->
-  let computation = Computation.create () in
-  let triggers = Array.init 4 @@ fun _ -> Trigger.create () in
-  let () =
-    triggers
-    |> Array.iter @@ fun trigger ->
-       Atomic.spawn @@ fun () ->
-       Atomic.check (fun () -> Computation.try_attach computation trigger);
-       Computation.detach computation trigger
-  in
-  Atomic.final @@ fun () ->
-  Atomic.check @@ fun () ->
-  Array.for_all Trigger.is_signaled triggers
-  &&
-  match Atomic.get computation with
-  | Canceled _ | Returned _ -> false
-  | Continue { balance; triggers } ->
-      balance <= 0
-      && List.length triggers <= 2
-      &&
-      let trigger = Trigger.create () in
-      Computation.try_attach computation trigger
-      && begin
-           match Atomic.get computation with
-           | Canceled _ | Returned _ -> false
-           | Continue { balance; triggers } ->
-               balance = 1 && triggers = [ trigger ]
-         end
+  [ `FIFO; `LIFO ]
+  |> List.iter @@ fun mode ->
+     Atomic.trace @@ fun () ->
+     let computation = Computation.create ~mode () in
+     let triggers = Array.init 4 @@ fun _ -> Trigger.create () in
+     let () =
+       triggers
+       |> Array.iter @@ fun trigger ->
+          Atomic.spawn @@ fun () ->
+          Atomic.check (fun () -> Computation.try_attach computation trigger);
+          Computation.detach computation trigger
+     in
+     Atomic.final @@ fun () ->
+     Atomic.check @@ fun () ->
+     Array.for_all Trigger.is_signaled triggers
+     &&
+     match Atomic.get computation with
+     | Canceled _ | Returned _ -> false
+     | Continue { balance_and_mode; triggers } ->
+         balance_and_mode <= Computation.fifo_bit
+         && List.length triggers <= 2
+         &&
+         let trigger = Trigger.create () in
+         Computation.try_attach computation trigger
+         && begin
+              match Atomic.get computation with
+              | Canceled _ | Returned _ -> false
+              | Continue { balance_and_mode; triggers } ->
+                  balance_and_mode <= Computation.one + Computation.fifo_bit
+                  && triggers = [ trigger ]
+            end
 
 let () =
   Alcotest.run "Picos DSCheck"