Fix rlimits and Popen

src/error/kind.ml

@@ -12,7 +12,6 @@ let of_twine st v =
   make ~name ()
 
 let generic_internal_error : t = make ~name:"GenericInternalError" ()
-let timeout : t = make ~name:"Timeout" ()
 let todo : t = make ~name:"Todo" ()
 
 module As_key = struct

src/error/kind.mli

@@ -15,9 +15,6 @@ val name : t -> string
 val generic_internal_error : t
 (** Any internal error that's not more specific *)
 
-val timeout : t
-(** Timeout *)
-
 val todo : t
 (** Not implemented yet *)
 

src/io/imandrakit_io.ml

@@ -69,10 +69,10 @@ let get_pid = Unix.getpid
 
 let with_signal ?(signal = Sys.sigint) ~on_sig f =
   Sys.catch_break false;
-  let handler = Sys.signal signal (Sys.Signal_handle on_sig) in
+  let previous = Sys.signal signal (Sys.Signal_handle on_sig) in
   let old_mask = Thread.sigmask Unix.SIG_UNBLOCK [ signal ] in
   Fun.protect f ~finally:(fun () ->
-      Sys.set_signal signal handler;
+      Sys.set_signal signal previous;
       ignore (Thread.sigmask Unix.SIG_BLOCK old_mask : _ list);
       Sys.catch_break true)
 
@@ -84,10 +84,8 @@ let block_signals () =
     ignore
       (Unix.sigprocmask Unix.SIG_BLOCK
          [
-           Sys.sigterm;
            Sys.sigpipe;
            Sys.sigint;
-           Sys.sigchld;
            Sys.sigalrm;
            Sys.sigusr1;
            Sys.sigusr2;

src/io/popen.ml

@@ -1,58 +1,135 @@
-open Moonpool
 module Log = (val Imandrakit_log.Logger.mk_log_str "x.popen")
 
-type state = {
-  stopped: bool Atomic.t;
-  res_code: int Fut.t;
-  promise_code: int Fut.promise;
-}
-
 type t = {
+  pid: int;
   stdin: out_channel;
   stdout: in_channel;
   stderr: in_channel;
-  pid: int;
-  _st: state;
+  mutable exit_code: (int, exn) result option;
+  exit_code_mutex: Mutex.t;
+  exit_code_condition: Condition.t;
+  start_time: Ptime.t;
+  mutable stop_time: Ptime.t option;
+  mutable on_exit: (t -> (int, exn) result -> unit) list;
+  is_group_leader: bool;
 }
-(** A sub-process *)
-
-let pp out self = Fmt.fprintf out "<pid=%d>" self.pid
-let show self = spf "<pid=%d>" self.pid
-let[@inline] stopped self = Atomic.get self._st.stopped
-
-let kill_and_close_ (self : t) =
-  let already_stopped = Atomic.exchange self._st.stopped true in
-  if not already_stopped then (
-    Log.debug (fun k -> k "(popen.kill-and-close :pid %d)" self.pid);
-    (try Unix.kill self.pid 15 with _ -> ());
-    close_out_noerr self.stdin;
-    close_in_noerr self.stdout;
-    close_in_noerr self.stderr;
-    (* just to be sure, wait a second and kill dash nine *)
-    ignore
-      (Thread.create
-         (fun () ->
-           Thread.delay 1.;
-           try Unix.kill self.pid 9 with _ -> ())
-         ()
-        : Thread.t);
-
-    (* kill zombies *)
-    let code =
+
+exception Killed
+
+(* Processes that we started. Global because all threads need access to it. *)
+let g_running_processes : t list ref = ref []
+let g_running_processes_mtx : Mutex.t = Mutex.create ()
+let g_more_to_reap : bool Atomic.t = Atomic.make false
+
+(* Initialization flag. *)
+let g_initialized : bool ref = ref false
+let g_initialized_mtx : Mutex.t = Mutex.create ()
+
+(* Reaper thread *)
+let g_reaper : Thread.t option ref = ref None
+let g_reaper_mtx : Mutex.t = Mutex.create ()
+let g_reaper_condition : Condition.t = Condition.create ()
+
+let fulfill (p : t) (r : (int, exn) result) (pid : int) : unit =
+  Mutex.lock p.exit_code_mutex;
+  p.exit_code <- Some r;
+  p.stop_time <- Some (Ptime_clock.now ());
+  Condition.broadcast p.exit_code_condition;
+  Mutex.unlock p.exit_code_mutex;
+  List.iter (fun f -> ignore (Thread.create (fun _ -> f p r))) p.on_exit
+
+let reap_one (p : t) : bool =
+  try
+    let wpid, wstatus = Unix.waitpid [ WNOHANG ] p.pid in
+    if wpid <> p.pid then
+      true
+    else (
+      match wstatus with
+      | WEXITED c ->
+        Log.debug (fun k -> k "(@[resolve :ok %d :c %d@])" p.pid c);
+        fulfill p (Ok c) p.pid;
+        false
+      | WSIGNALED c ->
+        Log.debug (fun k -> k "(@[resolve :error %d :c %d@])" p.pid c);
+        fulfill p (Error Killed) p.pid;
+        false
+      | WSTOPPED _ ->
+        (* Unreachable without WUNTRACED. *)
+        true
+    )
+  with Unix.Unix_error (Unix.ECHILD, _, _) -> true
+
+let rec reap () =
+  if Mutex.try_lock g_running_processes_mtx then (
+    while Atomic.exchange g_more_to_reap false do
       try
-        if not (String.equal Sys.os_type "Win32") then
-          ignore (Unix.sigprocmask Unix.SIG_BLOCK [ Sys.sigchld ]);
-        fst @@ Unix.waitpid [] self.pid
-      with _ -> max_int
-    in
-    Fut.fulfill_idempotent self._st.promise_code @@ Ok code
+        let r = List.filter reap_one !g_running_processes in
+        g_running_processes := r;
+        if List.length r <= 2 then
+          Log.debug (fun k ->
+              k "(@[remaining %a@])" (Fmt.Dump.list Fmt.int)
+                (List.map (fun x -> x.pid) r))
+        else
+          Log.debug (fun k -> k "(@[remaining :n %d@])" (List.length r))
+      with exc ->
+        Log.warn (fun k ->
+            k "(@[reap :exception '%s'@])" (Printexc.to_string exc))
+    done;
+    Mutex.unlock g_running_processes_mtx
   )
 
-let run_ ?(env = Unix.environment ()) cmd args : t =
-  (* block sigpipe *)
+let reaper () =
+  (* This thread, once kicked off, will run forever. Once there are no more
+     processes to reap, it will remain blocked and therefore won't consume any
+     time or additional memory. *)
+  while true do
+    Mutex.lock g_reaper_mtx;
+    Condition.wait g_reaper_condition g_reaper_mtx;
+    Atomic.set g_more_to_reap true;
+    reap ();
+    Mutex.unlock g_reaper_mtx
+  done
+
+let init () =
+  Mutex.protect g_initialized_mtx (fun _ ->
+      if not !g_initialized then (
+        let old_handler = Sys.signal Sys.sigchld Sys.Signal_ignore in
+
+        ignore
+          (Sys.set_signal Sys.sigchld
+             (Sys.Signal_handle
+                (fun _ ->
+                  Log.debug (fun k -> k "(sigchld)");
+
+                  (* Note: when we get here, we could be running in the same
+                     thread that just locked an exit code mutex in [await] and
+                     if [reap] were to lock that same mutex, OCaml would hate
+                     us. *)
+                  Atomic.set g_more_to_reap true;
+                  if Mutex.try_lock g_reaper_mtx then (
+                    Condition.signal g_reaper_condition;
+                    Mutex.unlock g_reaper_mtx
+                  );
+
+                  match old_handler with
+                  | Sys.Signal_handle h -> h Sys.sigchld
+                  | _ -> ())));
+
+        g_reaper := Some (Thread.create reaper ());
+
+        ignore (Unix.sigprocmask Unix.SIG_UNBLOCK [ Sys.sigchld ]);
+
+        g_initialized := true
+      ))
+
+let spawn (is_group_leader : bool) (env : string array) (cmd : string)
+    (args : string array) : t =
+  init ();
+
   if not (String.equal Sys.os_type "Win32") then
-    ignore (Unix.sigprocmask Unix.SIG_BLOCK [ Sys.sigpipe; Sys.sigchld ]);
-  (* make pipes, to give the appropriate ends to the subprocess *)
+    ignore (Unix.sigprocmask Unix.SIG_BLOCK [ Sys.sigpipe ]);
+
+  (* Make pipes, to give the appropriate ends to the subprocess *)
   let stdout, p_stdout = Unix.pipe () in
   let stderr, p_stderr = Unix.pipe () in
   let p_stdin, stdin = Unix.pipe () in
@@ -63,47 +140,140 @@ let run_ ?(env = Unix.environment ()) cmd args : t =
   let stdout = Unix.in_channel_of_descr stdout in
   let stderr = Unix.in_channel_of_descr stderr in
   let stdin = Unix.out_channel_of_descr stdin in
-  let pid = Unix.create_process_env cmd args env p_stdin p_stdout p_stderr in
-  let res_code, promise_code = Fut.make () in
-  Log.debug (fun k ->
-      k "Opened subprocess pid=%d cmd=%S args=[…%d]" pid cmd (Array.length args));
-  (* close the subprocess ends in here *)
+  let pid =
+    Unix.create_process_env cmd
+      (Array.append [| cmd |] args)
+      env p_stdin p_stdout p_stderr
+  in
+  (* Close the subprocess ends in here *)
   Unix.close p_stdout;
   Unix.close p_stdin;
   Unix.close p_stderr;
-  let p =
+  let r =
     {
+      pid;
       stdin;
       stdout;
       stderr;
-      pid;
-      _st = { stopped = Atomic.make false; res_code; promise_code };
+      exit_code = None;
+      exit_code_mutex = Mutex.create ();
+      exit_code_condition = Condition.create ();
+      on_exit = [];
+      start_time = Ptime_clock.now ();
+      stop_time = None;
+      is_group_leader;
     }
   in
-  p
-
-let run ?env cmd args : t = run_ ?env cmd (Array.of_list (cmd :: args))
-let res_code self = self._st.res_code
-let run_shell ?env cmd : t = run_ ?env "/bin/sh" [| "/bin/sh"; "-c"; cmd |]
-
-let kill self =
-  Log.debug (fun k -> k "(popen.kill %a)" pp self);
-  kill_and_close_ self
-
-let signal self s = Unix.kill self.pid s
-
-let wait (self : t) : int =
-  Log.debug (fun k -> k "(popen.wait %a)" pp self);
-  let res =
-    try
-      if not (String.equal Sys.os_type "Win32") then
-        ignore (Unix.sigprocmask Unix.SIG_BLOCK [ Sys.sigchld ]);
-      snd @@ Unix.waitpid [] self.pid
-    with _ -> Unix.WEXITED 0
+  Log.debug (fun k ->
+      k "(spawn :pid %d :cmd '%s' :args '%s')" r.pid cmd
+        (String.concat " " (Array.to_list args)));
+  Mutex.protect g_running_processes_mtx (fun x ->
+      g_running_processes := r :: !g_running_processes);
+  r
+
+let run ?(is_group_leader = false) ?(env = Unix.environment ()) (cmd : string)
+    (args : string list) : t =
+  spawn is_group_leader env cmd (Array.of_list args)
+
+let pid_alive (pid : int) =
+  try
+    Unix.kill pid 0;
+    true
+  with Unix.Unix_error (Unix.ESRCH, _, _) -> false
+
+let pid_is_gone ~(pid : int) ~(max_wait_s : float) =
+  let deadline = Unix.gettimeofday () +. max_wait_s in
+  let rec loop () =
+    if not (pid_alive pid) then
+      true
+    else if Unix.gettimeofday () > deadline then
+      false
+    else (
+      Unix.sleepf 0.1;
+      loop ()
+    )
   in
-  kill_and_close_ self;
-  let res =
-    match res with
-    | Unix.WEXITED i | Unix.WSTOPPED i | Unix.WSIGNALED i -> i
+  loop ()
+
+let await (self : t) : (int, exn) result =
+  Log.debug (fun k -> k "(await %d)" self.pid);
+  Mutex.lock self.exit_code_mutex;
+  let r =
+    while Option.is_none self.exit_code do
+      Condition.wait self.exit_code_condition self.exit_code_mutex
+    done;
+    Option.value self.exit_code
+      ~default:
+        (Error (Failure "Exit code of process unexpectedly not present."))
   in
-  res
+  Mutex.unlock self.exit_code_mutex;
+  r
+
+let kill ?(max_wait_s = 0.5) self =
+  Log.debug (fun k -> k "(kill %d)" self.pid);
+  let max_wait_s = max 0.0 max_wait_s in
+  try
+    let pgid =
+      if self.is_group_leader then
+        -self.pid
+      else
+        self.pid
+    in
+
+    (try Unix.kill pgid Sys.sigterm with
+    | Unix.Unix_error (Unix.ESRCH, _, _) ->
+      (* Perhaps it hasn't become a group leader yet. *)
+      (try Unix.kill self.pid Sys.sigterm with
+      | Unix.Unix_error (Unix.ESRCH, _, _) ->
+        (* Perhaps it just became a group leader. *)
+        (try Unix.kill pgid Sys.sigterm with _ -> ())
+      | exc ->
+        Log.debug (fun k ->
+            k "(@[kill :exception1@ '%s'@])" (Printexc.to_string exc)))
+    | exc ->
+      Log.debug (fun k ->
+          k "(@[kill :exception2@ '%s'@])" (Printexc.to_string exc)));
+
+    if not (pid_is_gone ~pid:pgid ~max_wait_s:(max_wait_s *. 0.75)) then (
+      Log.debug (fun k -> k "(hard-kill %d)" pgid);
+      (try Unix.kill (-self.pid) Sys.sigkill with _ -> ());
+      (try Unix.kill self.pid Sys.sigkill with _ -> ());
+      if
+        (not (pid_is_gone ~pid:pgid ~max_wait_s:(max_wait_s *. 0.25)))
+        && max_wait_s <> 0.0
+      then
+        Log.warn (fun k ->
+            k
+              "Could not verify that PID %d was killed successfully; child \
+               processes may be leaked."
+              pgid)
+    );
+
+    reap ()
+  with
+  | Unix.Unix_error (Unix.ESRCH, _, _) -> (* Ok, nothing to kill *) ()
+  | exc ->
+    Log.warn (fun k ->
+        k
+          "Child processes may be leaked due to exception raised while \
+           attempting to kill process %d: %s"
+          self.pid (Printexc.to_string exc))
+
+let kill_all () =
+  Mutex.protect g_running_processes_mtx (fun x ->
+      List.iter kill !g_running_processes)
+
+let signal (self : t) (s : int) = Unix.kill self.pid s
+
+let on_exit (self : t) (f : t -> (int, exn) result -> unit) : unit =
+  self.on_exit <- f :: self.on_exit
+
+let pid (self : t) : int = self.pid
+let stdin (self : t) : out_channel = self.stdin
+let stdout (self : t) : in_channel = self.stdout
+let stderr (self : t) : in_channel = self.stderr
+let start_time (self : t) : Ptime.t = self.start_time
+let stop_time (self : t) : Ptime.t option = self.stop_time
+
+let execution_time (self : t) : Ptime.span option =
+  Option.map (fun x -> Ptime.diff x self.start_time) self.stop_time