Add Control.timeout to Picos_structured

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

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

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

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;