Repository: ocaml-multicore/picos
Pull request #91 (2024-04-10 08:33:55)

Support `Unix.wait` and related functions in `Picos_stdio`
This PR adds support for scheduler friendly `Unix.waitpid`, `Unix.wait`, and `Unix.system` on non-Windows platforms.

On Windows I see no way to provide those in a non-blocking fashion (using only the OCaml standard libraries) except by spawning threads and for that a thread pool would make a lot of sense.  So, I'm leaving that for future work.


Patch:
diff --git a/CHANGES.md b/CHANGES.md
index 5f1ac50..aaf3d27 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -1,5 +1,7 @@
 ## Next version
 
+- Scheduler friendly `waitpid`, `wait`, and `system` in `Picos_stdio.Unix` for
+  platforms other than Windows (@polytypic)
 - Added `Picos_select.configure` to allow, and sometimes require, configuring
   `Picos_select` for co-operation with libraries that also deal with signals
   (@polytypic)
diff --git a/lib/picos_select/dune b/lib/picos_select/dune
index d4acfff..8cfea8d 100644
--- a/lib/picos_select/dune
+++ b/lib/picos_select/dune
@@ -9,6 +9,7 @@
   (re_export unix)
   picos.domain
   picos.thread
+  picos.htbl
   picos_thread_atomic
   backoff
   threads.posix
diff --git a/lib/picos_select/picos_select.ml b/lib/picos_select/picos_select.ml
index 37c6f55..00a96de 100644
--- a/lib/picos_select/picos_select.ml
+++ b/lib/picos_select/picos_select.ml
@@ -1,8 +1,35 @@
 open Picos
 
-type config = { mutable intr_sig : int; mutable intr_sigs : int list }
+let handle_sigchld_bit = 0b01
+let select_thread_running_on_main_domain_bit = 0b10
 
-let config = { intr_sig = 0; intr_sigs = [] }
+type config = {
+  mutable bits : int;
+  mutable intr_sig : int;
+  mutable intr_sigs : int list;
+}
+
+let config = { bits = 0; intr_sig = 0; intr_sigs = [] }
+
+(* *)
+
+type return =
+  | Return : {
+      value : 'a;
+      computation : 'a Computation.t;
+      mutable alive : bool;
+    }
+      -> return
+
+(** We use random numbers as keys for the awaiters. *)
+module RandomInt = struct
+  type t = int
+
+  let equal = Int.equal
+  let hash = Fun.id
+end
+
+let chld_awaiters = Picos_htbl.create ~hashed_type:(module RandomInt) ()
 
 (* *)
 
@@ -217,8 +244,15 @@ and select_thread_continue s rd wr ex (rd_fds, wr_fds, ex_fds) =
   select_thread s tos rd wr ex
 
 let select_thread s =
+  if Picos_domain.is_main_domain () then
+    config.bits <- select_thread_running_on_main_domain_bit lor config.bits;
   if not Sys.win32 then begin
-    Thread.sigmask SIG_BLOCK config.intr_sigs |> ignore
+    Thread.sigmask SIG_BLOCK config.intr_sigs |> ignore;
+    Thread.sigmask
+      (if config.bits land handle_sigchld_bit <> 0 then SIG_UNBLOCK
+       else SIG_BLOCK)
+      [ Sys.sigchld ]
+    |> ignore
   end;
   begin
     try
@@ -233,27 +267,38 @@ let select_thread s =
   if s.pipe_inn != Unix.stdin then Unix.close s.pipe_inn;
   if s.pipe_out != Unix.stdin then Unix.close s.pipe_out
 
-let[@poll error] [@inline never] try_configure ~intr_sig ~intr_sigs =
+let[@poll error] [@inline never] try_configure ~intr_sig ~intr_sigs
+    ~handle_sigchld =
   config.intr_sigs == []
   && begin
+       config.bits <- Bool.to_int handle_sigchld;
        config.intr_sig <- intr_sig;
        config.intr_sigs <- intr_sigs;
        true
      end
 
+let is_sigchld signum = signum = Sys.sigchld
 let is_intr_sig signum = signum = config.intr_sig
 
-let rec configure ?(intr_sig = Sys.sigusr2) () =
+let rec configure ?(intr_sig = Sys.sigusr2) ?(handle_sigchld = true) () =
   if not (Picos_thread.is_main_thread ()) then
     invalid_arg
       "Picos_select must be configured from the main thread on the main domain";
   assert (Sys.sigabrt = -1 && Sys.sigxfsz < Sys.sigabrt);
   if intr_sig < Sys.sigxfsz || 0 <= intr_sig || intr_sig = Sys.sigchld then
     invalid_arg "Invalid interrupt signal number";
-  if not (try_configure ~intr_sig ~intr_sigs:[ intr_sig ]) then
+  if not (try_configure ~intr_sig ~intr_sigs:[ intr_sig ] ~handle_sigchld) then
     invalid_arg "Picos_select.configure already configured";
 
   if not Sys.win32 then begin
+    if config.bits land handle_sigchld_bit <> 0 then begin
+      let previously_blocked = Thread.sigmask SIG_BLOCK [ Sys.sigchld ] in
+      assert (not (List.exists is_sigchld previously_blocked));
+      let old_behavior =
+        Sys.signal Sys.sigchld (Sys.Signal_handle handle_signal)
+      in
+      assert (old_behavior == Signal_default)
+    end;
     begin
       let previously_blocked = Thread.sigmask SIG_BLOCK config.intr_sigs in
       assert (not (List.exists is_intr_sig previously_blocked));
@@ -265,7 +310,13 @@ let rec configure ?(intr_sig = Sys.sigusr2) () =
   end
 
 and handle_signal signal =
-  if signal = config.intr_sig then
+  if signal = Sys.sigchld then begin
+    Picos_htbl.remove_all chld_awaiters
+    |> Seq.iter @@ fun (_, Return r) ->
+       r.alive <- false;
+       Computation.return r.computation r.value
+  end
+  else if signal = config.intr_sig then
     let (Req r) = Picos_thread.TLS.get intr_key in
     Computation.return r.computation Signaled
 
@@ -452,3 +503,32 @@ module Intr = struct
           decr Backoff.default
         end
 end
+
+(* *)
+
+let rec insert return =
+  let id = Random.bits () in
+  if Picos_htbl.try_add chld_awaiters id return then id else insert return
+
+let[@alert "-handler"] return_on_sigchld computation value =
+  if config.bits land select_thread_running_on_main_domain_bit = 0 then
+    (* Ensure there is at least one thread handling [Sys.sigchld] signals. *)
+    get () |> ignore;
+  let return = Return { value; computation; alive = true } in
+  let id = insert return in
+  let remover =
+    Trigger.from_action id return @@ fun _trigger id (Return this_r as this) ->
+    if this_r.alive then begin
+      this_r.alive <- false;
+      (* It should be extremely rare, but possible, that the return was already
+         removed and another added just at this point and so we must account for
+         the possibility and make sure that whatever we remove is completed. *)
+      match Picos_htbl.remove_exn chld_awaiters id with
+      | Return that_r as that ->
+          if this != that then
+            Computation.return that_r.computation that_r.value
+      | exception Not_found -> ()
+    end
+  in
+  if not (Computation.try_attach computation remover) then
+    Trigger.signal remover
diff --git a/lib/picos_select/picos_select.mli b/lib/picos_select/picos_select.mli
index 8537ccc..5d872f0 100644
--- a/lib/picos_select/picos_select.mli
+++ b/lib/picos_select/picos_select.mli
@@ -70,15 +70,31 @@ module Intr : sig
       {!req} must be cleared exactly once! *)
 end
 
+(** {1 Processes} *)
+
+val return_on_sigchld : 'a Computation.t -> 'a -> unit
+(** [return_on_sigchld computation value] arranges for [computation] to be
+    {{!Picos.Computation.return} returned} with given [value] on next
+    {!Sys.sigchld}.  Completion of the [computation] before a {!Sys.sigchld} is
+    received effectively cancels the arrangement.
+
+    ⚠️ The mechanism uses the {!Sys.sigchld} signal which should not be used for
+    other purposes. *)
+
 (** {1 Configuration} *)
 
-val configure : ?intr_sig:int -> unit -> unit
-(** [configure ~intr_sig ()] can, and sometimes must, be called by an
-    application to configure the use of signals by this module.
+val configure : ?intr_sig:int -> ?handle_sigchld:bool -> unit -> unit
+(** [configure ~intr_sig ~handle_sigchld ()] can, and sometimes must, be called
+    by an application to configure the use of signals by this module.
 
     The optional [intr_sig] argument can be used to specify the signal used by
     the {{!Intr} interrupt} mechanism.  The default is to use {!Sys.sigusr2}.
 
+    The optional [handle_sigchld] argument can be used to specify whether this
+    module should setup handling of {!Sys.sigchld}.  The default is [true].
+    When explicitly specified as [~handle_sigchld:false], the application should
+    arrange to call {!handle_signal} whenever a {!Sys.sigchld} signal occurs.
+
     ⚠️ This module must always be configured before use.  Unless this module has
     been explicitly configured, calling a method of this module from the main
     thread on the main domain will automatically configure this module with
diff --git a/lib/picos_stdio/picos_stdio.ml b/lib/picos_stdio/picos_stdio.ml
index 4f2d0f1..0b73f61 100644
--- a/lib/picos_stdio/picos_stdio.ml
+++ b/lib/picos_stdio/picos_stdio.ml
@@ -276,6 +276,94 @@ module Unix = struct
 
   (* *)
 
+  module Wait_flag = struct
+    let nohang_bit = 0b10
+    let untraced_bit = 0b01
+
+    (* Note that this is optimized to the identity function. *)
+    let to_int = function WNOHANG -> 0 | WUNTRACED -> 1
+    let to_bit flag = nohang_bit - to_int flag
+    let () = assert (to_bit WNOHANG = nohang_bit)
+    let () = assert (to_bit WUNTRACED = untraced_bit)
+
+    let rec to_bits flags bits =
+      match flags with
+      | [] -> bits
+      | flag :: flags -> to_bits flags (bits lor to_bit flag)
+
+    let to_bits flags = to_bits flags 0
+    let to_flags = [| []; [ WUNTRACED ]; [ WNOHANG ]; [ WNOHANG; WUNTRACED ] |]
+    let to_flags bits = Array.unsafe_get to_flags bits
+  end
+
+  let rec waitpid_unix ~bits ~pid =
+    if bits land Wait_flag.nohang_bit <> 0 then
+      Unix.waitpid (Wait_flag.to_flags bits) pid
+    else
+      let computation = Computation.create () in
+      Picos_select.return_on_sigchld computation ();
+      match
+        Unix.waitpid (Wait_flag.to_flags (bits lor Wait_flag.nohang_bit)) pid
+      with
+      | exception Unix_error (EINTR, _, _) -> waitpid_unix ~bits ~pid
+      | (pid_or_0, _) as result ->
+          if pid_or_0 = 0 then begin
+            Computation.await computation;
+            waitpid_unix ~bits ~pid
+          end
+          else begin
+            Computation.finish computation;
+            result
+          end
+      | exception exn ->
+          Computation.finish computation;
+          raise exn
+
+  let waitpid_win32 ~bits ~pid =
+    if bits land Wait_flag.nohang_bit <> 0 then
+      Unix.waitpid (Wait_flag.to_flags bits) pid
+    else
+      (* One way to provide a scheduler friendly [waitpid] on Windows would be
+         to use a thread pool to run blocking operations on.  PR for a thread
+         pool implemented in Picos would be welcome! *)
+      invalid_arg
+        "Picos_stdio.Unix.waitpid is currently not supported on Windows \
+         without WNOHANG"
+
+  let waitpid flags pid =
+    let bits = Wait_flag.to_bits flags in
+    if Sys.win32 then
+      if pid <> -1 then waitpid_win32 ~bits ~pid
+      else begin
+        (* This should raise? *)
+        Unix.waitpid flags pid
+      end
+    else waitpid_unix ~bits ~pid
+
+  let wait () =
+    if not Sys.win32 then waitpid_unix ~bits:0 ~pid:(-1)
+    else begin
+      (* This should raise [Invalid_argument] *)
+      Unix.wait ()
+    end
+
+  (* *)
+
+  let sh = "/bin/sh"
+
+  let system cmd =
+    if Sys.win32 then
+      (* One way to provide a scheduler friendly [system] on Windows would be to
+         use a thread pool to run blocking operations on.  PR for a thread pool
+         implemented in Picos would be welcome! *)
+      invalid_arg
+        "Picos_stdio.Unix.system is currently not supported on Windows"
+    else
+      create_process sh [| sh; "-c"; cmd |] stdin stdout stderr
+      |> waitpid [] |> snd
+
+  (* *)
+
   (*let select _ = failwith "TODO"*)
 
   (* *)
diff --git a/lib/picos_stdio/picos_stdio.mli b/lib/picos_stdio/picos_stdio.mli
index be816fa..99f32ec 100644
--- a/lib/picos_stdio/picos_stdio.mli
+++ b/lib/picos_stdio/picos_stdio.mli
@@ -27,7 +27,13 @@ module Unix : sig
       - {!sleep}, and
       - {!sleepf}
 
-      also block in a scheduler friendly manner.
+      also block in a scheduler friendly manner.  Additionally
+
+      - {!wait},
+      - {!waitpid}, and
+      - {!system}
+
+      also block in a scheduler friendly manner except on Windows.
 
       ⚠️ This module uses {!Picos_select} and you may need to
       {{!Picos_select.configure} configure} it at start of your application.
@@ -144,11 +150,9 @@ module Unix : sig
   val execvp : string -> string array -> 'a
   val execvpe : string -> string array -> string array -> 'a
   val fork : unit -> int
-
-  (*val wait : unit -> int * process_status*)
-  (*val waitpid : wait_flag list -> int -> int * process_status*)
-  (*val system : string -> process_status*)
-
+  val wait : unit -> int * process_status
+  val waitpid : wait_flag list -> int -> int * process_status
+  val system : string -> process_status
   val _exit : int -> 'a
   val getpid : unit -> int
   val getppid : unit -> int


Test patch:
diff --git a/test/dune b/test/dune
index ebdb85d..bd5bc2f 100644
--- a/test/dune
+++ b/test/dune
@@ -52,6 +52,13 @@
 
 ;;
 
+(test
+ (name test_stdio)
+ (modules test_stdio)
+ (libraries test_scheduler picos.stdio alcotest))
+
+;;
+
 (test
  (name test_select)
  (modules test_select)
diff --git a/test/test_stdio.ml b/test/test_stdio.ml
new file mode 100644
index 0000000..e20ccad
--- /dev/null
+++ b/test/test_stdio.ml
@@ -0,0 +1,18 @@
+open Picos_stdio
+
+let test_system_unix () =
+  Test_scheduler.run @@ fun () ->
+  assert (Unix.system "exit 101" = Unix.WEXITED 101);
+  assert (Unix.system "echo Hello world!" = Unix.WEXITED 0);
+  assert (Unix.system "this-is-not-supposed-to-exist" = Unix.WEXITED 127);
+  match Unix.wait () with
+  | _ -> assert false
+  | exception Unix.Unix_error (ECHILD, _, _) -> ()
+
+let () =
+  [
+    ( "Unix",
+      if Sys.win32 then []
+      else [ Alcotest.test_case "system" `Quick test_system_unix ] );
+  ]
+  |> Alcotest.run "Picos_stdio"