Merge PR #9881: Protect some I/O routines from SIGALRM

Ack-by: SkySkimmer
Reviewed-by: ejgallego
Ack-by: maximedenes

2 files changed, 29 insertions, 0 deletions

diff --git a/lib/control.ml b/lib/control.ml
index ffb3584f1e..9054507e46 100644
--- a/lib/control.ml
+++ b/lib/control.ml
@@ -89,3 +89,21 @@ let timeout_fun_ref = ref timeout_fun
 let set_timeout f = timeout_fun_ref := f
 
 let timeout n f e = !timeout_fun_ref.timeout n f e
+
+let protect_sigalrm f x =
+  let timed_out = ref false in
+  let timeout_handler _ = timed_out := true in
+  try
+    let old_handler = Sys.signal Sys.sigalrm (Sys.Signal_handle timeout_handler) in
+    try
+      let res = f x in
+      Sys.set_signal Sys.sigalrm old_handler;
+      match !timed_out, old_handler with
+      | true, Sys.Signal_handle f -> f Sys.sigalrm; res
+      | _, _ -> res
+    with e ->
+      let e = Backtrace.add_backtrace e in
+      Sys.set_signal Sys.sigalrm old_handler;
+      Exninfo.iraise e
+  with Invalid_argument _ -> (* This happens on Windows, as handling SIGALRM does not seem supported *)
+    f x
diff --git a/lib/control.mli b/lib/control.mli
index 59e2a15158..640d41a4f7 100644
--- a/lib/control.mli
+++ b/lib/control.mli
@@ -29,3 +29,14 @@ val timeout : int -> ('a -> 'b) -> 'a -> exn -> 'b
    API and it is scheduled to go away. *)
 type timeout = { timeout : 'a 'b. int -> ('a -> 'b) -> 'a -> exn -> 'b }
 val set_timeout : timeout -> unit
+
+(** [protect_sigalrm f x] computes [f x], but if SIGALRM is received during that
+    computation, the signal handler is executed only once the computation is
+    terminated. Otherwise said, it makes the execution of [f] atomic w.r.t.
+    handling of SIGALRM.
+
+    This is useful for example to prevent the implementation of `Timeout` to
+    interrupt I/O routines, generating ill-formed output.
+
+*)
+val protect_sigalrm : ('a -> 'b) -> 'a -> 'b