From 65a8e06130a8c97a58fabc5a8a8dceb7561b43c9 Mon Sep 17 00:00:00 2001
From: Vesa Karvonen <vesa.a.j.k@gmail.com>
Date: Mon, 22 Jul 2024 19:39:32 +0300
Subject: [PATCH] Add `Control.timeout` to `Picos_structured`

---
 lib/picos_structured/control.ml           | 26 +++++++++++++++++++++++
 lib/picos_structured/picos_structured.mli | 12 +++++++++++
 test/test_server_and_client.ml            |  2 +-
 test/test_structured.ml                   |  8 +++----
 4 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/lib/picos_structured/control.ml b/lib/picos_structured/control.ml
index b7786df53..f8e432a72 100644
--- a/lib/picos_structured/control.ml
+++ b/lib/picos_structured/control.ml
@@ -7,6 +7,10 @@ exception Terminate
 
 let terminate_bt = Exn_bt.get_callstack 0 Terminate
 
+exception Timeout
+
+let timeout_bt = Exn_bt.get_callstack 0 Timeout
+
 let terminate_bt ?callstack () =
   match callstack with
   | None -> terminate_bt
@@ -63,3 +67,25 @@ let block () =
   | Some exn_bt -> Exn_bt.raise exn_bt
 
 let protect thunk = Fiber.forbid (Fiber.current ()) thunk
+
+let timeout ?(exn_bt = timeout_bt) ~seconds thunk =
+  let into = Computation.create ~mode:`LIFO () in
+  Computation.cancel_after into ~seconds exn_bt;
+  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 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.finish into;
+      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 2348ae784..03ad7bbe3 100644
--- a/lib/picos_structured/picos_structured.mli
+++ b/lib/picos_structured/picos_structured.mli
@@ -147,6 +147,18 @@ module Control : sig
 
       @raise Invalid_argument in case propagation of cancelation has been
         {{!protect} forbidden}. *)
+
+  exception Timeout
+  (** An exception that can be used to signal a timeout.
+
+      ℹ️ Unlike {!Terminate} a {!Timeout} is considered an error. *)
+
+  val timeout : ?exn_bt:Exn_bt.t -> seconds:float -> (unit -> 'a) -> 'a
+  (** [timeout ~seconds thunk] calls [thunk ()] on the current fiber with the
+      specified timeout in [seconds].
+
+      The optional [exn_bt] argument defaults to a {!Timeout} with an empty
+      initial backtrace (i.e. [Exn_bt.get_callstack 0 Timeout]). *)
 end
 
 module Promise : sig
diff --git a/test/test_server_and_client.ml b/test/test_server_and_client.ml
index aa8327761..66b765459 100644
--- a/test/test_server_and_client.ml
+++ b/test/test_server_and_client.ml
@@ -100,7 +100,7 @@ let main () =
       Unix.socket ~cloexec:true PF_INET SOCK_STREAM 0
     in
     set_nonblock socket;
-    Unix.connect socket !server_addr;
+    Control.timeout ~seconds:1.0 (fun () -> Unix.connect socket !server_addr);
     Printf.printf "  Client %s connected\n%!" id;
     let bytes = Bytes.create n in
     let n = Unix.write socket bytes 0 (Bytes.length bytes) in
diff --git a/test/test_structured.ml b/test/test_structured.ml
index af6320fce..d0864db00 100644
--- a/test/test_structured.ml
+++ b/test/test_structured.ml
@@ -241,14 +241,14 @@ let test_race_any () =
     Test_util.shuffle
       [
         (fun () ->
-          Control.sleep ~seconds:2.9;
-          winner := 3);
+          try Control.timeout ~seconds:2.9 Control.block
+          with Control.Timeout -> winner := 3);
         (fun () ->
           Control.sleep ~seconds:1.5;
           winner := 2);
         (fun () ->
-          Control.sleep ~seconds:0.1;
-          winner := 1);
+          try Control.timeout ~seconds:0.1 Control.block
+          with Control.Timeout -> winner := 1);
       ]
   in
   Run.any ops;