diff --git a/src/error/kind.ml b/src/error/kind.ml index 8f6cf4e20..30237f5e8 100644 --- a/src/error/kind.ml +++ b/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 diff --git a/src/error/kind.mli b/src/error/kind.mli index 4806a15f8..2ef43404f 100644 --- a/src/error/kind.mli +++ b/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 *) diff --git a/src/io/imandrakit_io.ml b/src/io/imandrakit_io.ml index 41c241f61..3cbe5e306 100644 --- a/src/io/imandrakit_io.ml +++ b/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; diff --git a/src/io/popen.ml b/src/io/popen.ml index 9fb51f57e..e1f0f6220 100644 --- a/src/io/popen.ml +++ b/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 "" self.pid -let show self = spf "" 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 diff --git a/src/io/popen.mli b/src/io/popen.mli index 08a207b01..6d1a55039 100644 --- a/src/io/popen.mli +++ b/src/io/popen.mli @@ -1,29 +1,47 @@ -(** Run sub-processes. +(** Manage sub-processes. *) - This gives more control than the equivalent {!Unix} APIs. *) +type t +(** A sub-process *) -type state +exception Killed +(** Exception indicating that a process did not run to completion. *) -type t = private { - stdin: out_channel; - stdout: in_channel; - stderr: in_channel; - pid: int; - _st: state; -} -[@@deriving show] -(** A sub-process *) +val run : + ?is_group_leader:bool -> ?env:string array -> string -> string list -> t +(** Runs subprocess with the given command and arguments. *) + +val await : t -> (int, exn) result +(** Awaits the exit of a process. *) -val run : ?env:string array -> string -> string list -> t -(** Run subprocess with given command *) +val kill : ?max_wait_s:float -> t -> unit +(** Kills a process. *) -val run_shell : ?env:string array -> string -> t -(** Run subprocess with given command *) +val kill_all : unit -> unit +(** Kills all known processes. *) -val res_code : t -> int Moonpool.Fut.t -val wait : t -> int -val kill : t -> unit val signal : t -> int -> unit +(** Sends a signal to the process. *) + +val on_exit : t -> (t -> (int, exn) result -> unit) -> unit +(** Registers a callback to be run (in a new thread) when the process exits. *) + +val pid : t -> int +(** The process identifier of the process. *) + +val stdin : t -> out_channel +(** Standard Input of the process. *) + +val stdout : t -> in_channel +(** Standard Output of the process. *) + +val stderr : t -> in_channel +(** Standard Error Output of the process. *) + +val start_time : t -> Ptime.t +(** The time the process was started. *) + +val stop_time : t -> Ptime.t option +(** The time the result of the process arrived. *) -val stopped : t -> bool -(** We know that we have stopped the process *) +val execution_time : t -> Ptime.span option +(** Wall-clock execution time. *) diff --git a/src/io/setrlimit/dune b/src/io/setrlimit/dune index 369065d6f..4535c98ac 100644 --- a/src/io/setrlimit/dune +++ b/src/io/setrlimit/dune @@ -3,6 +3,8 @@ (public_name imandrakit-io.setrlimit) (synopsis "Basic resource control via setrlimit") (c_library_flags :standard) + (preprocess + (pps ppx_deriving.std ppx_deriving_yojson)) (foreign_stubs (language c) ;(include_dirs .) diff --git a/src/io/setrlimit/imandrakit_io_setrlimit.ml b/src/io/setrlimit/imandrakit_io_setrlimit.ml index 80663821a..358f11eb6 100644 --- a/src/io/setrlimit/imandrakit_io_setrlimit.ml +++ b/src/io/setrlimit/imandrakit_io_setrlimit.ml @@ -15,21 +15,55 @@ type resource = module Raw = struct let resource_to_int = function - | RLIMIT_CORE -> 0 - | RLIMIT_CPU -> 1 - | RLIMIT_DATA -> 2 - | RLIMIT_FSIZE -> 3 - | RLIMIT_NOFILE -> 4 - | RLIMIT_STACK -> 5 - | RLIMIT_AS -> 6 - - external set : int -> int -> bool = "caml_imandrakit_setrlimit" + | RLIMIT_CORE -> 0n + | RLIMIT_CPU -> 1n + | RLIMIT_DATA -> 2n + | RLIMIT_FSIZE -> 3n + | RLIMIT_NOFILE -> 4n + | RLIMIT_STACK -> 5n + | RLIMIT_AS -> 6n + + external get : nativeint -> (nativeint * nativeint, nativeint) Result.t + = "caml_imandrakit_getrlimit" + + external set : nativeint -> nativeint -> nativeint -> bool + = "caml_imandrakit_setrlimit" end (** [set resource limit] returns [true] if setting the limit succeeded *) -let[@inline] set (r : resource) (v : int) : bool = - Raw.set (Raw.resource_to_int r) v +let[@inline] set (r : resource) (cur : nativeint) (max : nativeint) : bool = + Raw.set (Raw.resource_to_int r) cur max + +type limits = { + cur: nativeint option; (* Soft limit *) + max: nativeint option; (* Hard limit *) +} +[@@deriving show { with_path = false }] (** Like {!set}, but propagates failures @raise Failure if it fails *) -let set_exn r v : unit = if not (set r v) then failwith "setrlimit failed" +let set_exn (r : resource) (l : limits) : unit = + let cur : nativeint = Option.value l.cur ~default:Nativeint.minus_one in + let max : nativeint = Option.value l.max ~default:Nativeint.minus_one in + if not (set r cur max) then failwith "setrlimit failed" + +let set_hard_exn (r : resource) (max : nativeint option) : unit = + set_exn r { cur = None; max } + +let get (r : resource) : (limits, nativeint) Result.t = + match Raw.get (Raw.resource_to_int r) with + | Ok (cur, max) -> + Ok + { + cur = + (if cur = Nativeint.minus_one then + None + else + Some cur); + max = + (if max = Nativeint.minus_one then + None + else + Some max); + } + | Error e -> Error e diff --git a/src/io/setrlimit/libimandrakit_setrlimit_stubs.c b/src/io/setrlimit/libimandrakit_setrlimit_stubs.c index e3a3fd2d0..eea7fbdfc 100644 --- a/src/io/setrlimit/libimandrakit_setrlimit_stubs.c +++ b/src/io/setrlimit/libimandrakit_setrlimit_stubs.c @@ -1,4 +1,5 @@ +#include #include #if !defined(_WIN32) && !defined(_WIN64) @@ -15,27 +16,57 @@ int resources[7] = {RLIMIT_CORE, RLIMIT_CPU, RLIMIT_DATA, RLIMIT_FSIZE, RLIMIT_NOFILE, RLIMIT_STACK, RLIMIT_AS}; #endif -CAMLprim value caml_imandrakit_setrlimit(value _res, value _value) { - CAMLparam2(_res, _value); +CAMLprim value caml_imandrakit_setrlimit(value _resource, value _cur, value _max) { + CAMLparam3(_resource, _cur, _max); #if !defined(_WIN32) && !defined(_WIN64) - int resource = resources[Int_val(_res)]; - unsigned long limit = Int_val(_value); + int resource = resources[Nativeint_val(_resource)]; + long unsigned int cur = Nativeint_val(_cur); + long unsigned int max = Nativeint_val(_max); + bool is_ok = true; struct rlimit old_limits; if (getrlimit(resource, &old_limits) != 0) CAMLreturn(false); + long unsigned int new_max = max < old_limits.rlim_max ? max : old_limits.rlim_max; + long unsigned int new_cur = cur < old_limits.rlim_cur ? cur : old_limits.rlim_cur; + + if (new_cur > new_max) + new_cur = new_max; + const struct rlimit new_limits = { - .rlim_cur = limit < old_limits.rlim_max ? limit : old_limits.rlim_max, - .rlim_max = old_limits.rlim_max, + .rlim_cur = new_cur, + .rlim_max = new_max, }; - int res = setrlimit(resource, &new_limits); - bool isok = (res == 0); + if (old_limits.rlim_cur != new_limits.rlim_cur || + old_limits.rlim_max != new_limits.rlim_max) { + is_ok = setrlimit(resource, &new_limits) == 0; + } #else bool isok = true; #endif + CAMLreturn(Val_bool(is_ok)); +} + +CAMLprim value caml_imandrakit_getrlimit(value _resource) { + CAMLparam1(_resource); + CAMLlocal2(r, sr); + + int resource = resources[Nativeint_val(_resource)]; + + struct rlimit limits; + if (getrlimit(resource, &limits) != 0) { + r = caml_alloc(1, 1); + Store_field(r, 0, errno); + } else { + r = caml_alloc(1, 0); + sr = caml_alloc_tuple(2); + Store_field(sr, 0, caml_copy_nativeint(limits.rlim_cur)); + Store_field(sr, 1, caml_copy_nativeint(limits.rlim_max)); + Store_field(r, 0, sr); + } - CAMLreturn(Val_bool(isok)); + CAMLreturn(r); } diff --git a/src/log/logger.ml b/src/log/logger.ml index 8230b01a4..76015104b 100644 --- a/src/log/logger.ml +++ b/src/log/logger.ml @@ -122,7 +122,9 @@ module Output = struct output_string oc s; output_char oc '\n'; if autoflush then Stdlib.flush oc - with _ -> Printf.eprintf "logger: failed to log to chan\n%!") + with exc -> + Printf.eprintf "logger: failed to log to chan:\n%s\n%!" + (Printexc.to_string exc)) () let stdout () = to_chan stdout diff --git a/test/io/setrlimit/t1.ml b/test/io/setrlimit/t1.ml index 7793c298a..d2f3fc64f 100644 --- a/test/io/setrlimit/t1.ml +++ b/test/io/setrlimit/t1.ml @@ -1,5 +1,5 @@ module R = Imandrakit_io_setrlimit let () = - R.set_exn R.RLIMIT_CPU 500; + R.set_hard_exn R.RLIMIT_CPU (Some 500n); () diff --git a/test/io/setrlimit/t2.ml b/test/io/setrlimit/t2.ml index cf9817de6..e878602ac 100644 --- a/test/io/setrlimit/t2.ml +++ b/test/io/setrlimit/t2.ml @@ -2,7 +2,7 @@ module R = Imandrakit_io_setrlimit let () = (* 2 MB max *) - R.set_exn R.RLIMIT_AS 2_000_000; + R.set_hard_exn R.RLIMIT_AS (Some 2_000_000n); (* now allocate a 8MB array *) try let _arr = Sys.opaque_identity (Array.make 1_000_000 true) in