Repository: ocaml-multicore/picos
Pull request #186 (2024-07-22 16:40:01)

Add `Control.terminate_after` and `Bundle.terminate_after`
This PR adds `Control.terminate_after ~seconds thunk` to run the given thunk with a timeout on the current fiber and `Bundle.terminate_after ~seconds bundle` to similarly terminate a bundle.


Patch:
diff --git a/lib/picos_structured/bundle.ml b/lib/picos_structured/bundle.ml
index 4fb2187..2c6f237 100644
--- a/lib/picos_structured/bundle.ml
+++ b/lib/picos_structured/bundle.ml
@@ -8,8 +8,11 @@ type t = {
 }
 
 let terminate ?callstack t =
-  let terminate_bt = Control.terminate_bt ?callstack () in
-  Computation.cancel t.bundle terminate_bt
+  Computation.cancel t.bundle (Control.terminate_bt ?callstack ())
+
+let terminate_after ?callstack t ~seconds =
+  Computation.cancel_after t.bundle ~seconds
+    (Control.terminate_bt ?callstack ())
 
 let error ?callstack t (exn_bt : Exn_bt.t) =
   if exn_bt.Exn_bt.exn != Control.Terminate then begin
diff --git a/lib/picos_structured/control.ml b/lib/picos_structured/control.ml
index b7786df..c889542 100644
--- a/lib/picos_structured/control.ml
+++ b/lib/picos_structured/control.ml
@@ -63,3 +63,27 @@ let block () =
   | Some exn_bt -> Exn_bt.raise exn_bt
 
 let protect thunk = Fiber.forbid (Fiber.current ()) thunk
+
+let terminate_after ?callstack ~seconds thunk =
+  let into = Computation.create ~mode:`LIFO () in
+  let fiber = Fiber.current () in
+  let (Packed from as packed) = Fiber.get_computation fiber in
+  let canceler = Computation.attach_canceler ~from ~into in
+  Fiber.set_computation fiber (Packed into);
+  match
+    Computation.cancel_after into ~seconds (terminate_bt ?callstack ());
+    thunk ()
+  with
+  | result ->
+      Computation.finish into;
+      let (Packed from) = packed in
+      Computation.detach from canceler;
+      Fiber.set_computation fiber packed;
+      result
+  | exception exn ->
+      let exn_bt = Exn_bt.get exn in
+      Computation.cancel into exn_bt;
+      let (Packed from) = packed in
+      Computation.detach from canceler;
+      Fiber.set_computation fiber packed;
+      Exn_bt.raise exn_bt
diff --git a/lib/picos_structured/picos_structured.mli b/lib/picos_structured/picos_structured.mli
index 2348ae7..481bf50 100644
--- a/lib/picos_structured/picos_structured.mli
+++ b/lib/picos_structured/picos_structured.mli
@@ -147,6 +147,23 @@ module Control : sig
 
       @raise Invalid_argument in case propagation of cancelation has been
         {{!protect} forbidden}. *)
+
+  val terminate_after : ?callstack:int -> seconds:float -> (unit -> 'a) -> 'a
+  (** [terminate_after ~seconds thunk] arranges to terminate the execution of
+      [thunk] on the current fiber after the specified timeout in [seconds].
+
+      The optional [callstack] argument specifies the number of callstack
+      entries to capture with the {{!Control.Terminate} [Terminate]} exception.
+      The default is [0].
+
+      As an example, [terminate_after] could be implemented using {!Bundle} as
+      follows:
+      {[
+        let terminate_after ?callstack ~seconds thunk =
+          Bundle.join_after @@ fun bundle ->
+          Bundle.terminate_after ?callstack ~seconds bundle;
+          thunk ()
+      ]} *)
 end
 
 module Promise : sig
@@ -198,7 +215,7 @@ module Promise : sig
   val terminate_after : ?callstack:int -> 'a t -> seconds:float -> unit
   (** [terminate_after ~seconds promise] arranges to terminate the [promise] by
       canceling it with the {{!Control.Terminate} [Terminate]} exception after
-      the specified number of [seconds].
+      the specified timeout in [seconds].
 
       The optional [callstack] argument specifies the number of callstack
       entries to capture with the {{!Control.Terminate} [Terminate]} exception.
@@ -246,6 +263,10 @@ module Bundle : sig
       raise the exception to propagate cancelation unless {{!Control.protect}
       propagation of cancelation is forbidden}. *)
 
+  val terminate_after : ?callstack:int -> t -> seconds:float -> unit
+  (** [terminate_after ~seconds bundle] arranges to {!terminate} the bundle
+      after the specified timeout in [seconds]. *)
+
   val error : ?callstack:int -> t -> Exn_bt.t -> unit
   (** [error bundle exn_bt] first calls {!terminate} and then adds the exception
       with backtrace to the list of exceptions to be raised, unless the
diff --git a/lib/picos_structured/promise.ml b/lib/picos_structured/promise.ml
index f9317a2..c9d3a2f 100644
--- a/lib/picos_structured/promise.ml
+++ b/lib/picos_structured/promise.ml
@@ -9,11 +9,9 @@ let completed = Event.from_computation
 let is_running = Computation.is_running
 
 let try_terminate ?callstack t =
-  let terminate_bt = Control.terminate_bt ?callstack () in
-  Computation.try_cancel t terminate_bt
+  Computation.try_cancel t (Control.terminate_bt ?callstack ())
 
 let terminate ?callstack t = try_terminate ?callstack t |> ignore
 
 let terminate_after ?callstack t ~seconds =
-  let terminate_bt = Control.terminate_bt ?callstack () in
-  Computation.cancel_after t ~seconds terminate_bt
+  Computation.cancel_after t ~seconds (Control.terminate_bt ?callstack ())


Test patch:
diff --git a/test/test_structured.ml b/test/test_structured.ml
index af6320f..f72656e 100644
--- a/test/test_structured.ml
+++ b/test/test_structured.ml
@@ -236,24 +236,26 @@ let test_any_and_all_returns () =
 
 let test_race_any () =
   Test_scheduler.run ~max_domains:4 @@ fun () ->
-  let winner = ref 0 in
+  let winner = Atomic.make 0 in
   let ops =
     Test_util.shuffle
       [
         (fun () ->
-          Control.sleep ~seconds:2.9;
-          winner := 3);
+          try Control.terminate_after ~seconds:2.9 Control.block
+          with Control.Terminate ->
+            Atomic.compare_and_set winner 0 3 |> ignore);
         (fun () ->
           Control.sleep ~seconds:1.5;
-          winner := 2);
+          Atomic.compare_and_set winner 0 2 |> ignore);
         (fun () ->
-          Control.sleep ~seconds:0.1;
-          winner := 1);
+          try Control.terminate_after ~seconds:0.1 Control.block
+          with Control.Terminate ->
+            Atomic.compare_and_set winner 0 1 |> ignore);
       ]
   in
   Run.any ops;
   (* This is non-deterministic and may need to changed if flaky *)
-  assert (!winner = 1)
+  assert (Atomic.get winner = 1)
 
 let () =
   [