-let read_procmem pid =
- let cin = open_in (Printf.sprintf "/proc/%i/status" pid) in
- let matchline s =
- try
- Scanf.sscanf s " VmRSS: %i kB" (fun i -> Some i)
- with
- | _ -> None
- in
- let rec loop () =
- match matchline (input_line cin) with
- Some i -> i
- | None -> loop ()
- in
- let s = try loop() with _ -> -1 in
- Printf.eprintf "Memory: %i\n%!" s;
- close_in cin;
- s
+open Format
+
+module Options =
+ struct
+ let outfile = ref "/dev/stderr"
+ let page_size = ref ~-1
+ let timeout = ref 30
+ let max_mem = ref (512*1024)
+ let interval = ref 250
+ let quiet = ref false
+
+ let spec = Arg.align [
+ "-o", Arg.Set_string outfile, " output file (default stderr)";
+ "-p", Arg.Set_int page_size, " page size in bytes (default autodetect)";
+ "-t", Arg.Set_int timeout, " timeout in s (default 30s)";
+ "-m", Arg.Set_int max_mem, " maximum memory in kiB (default 512 MiB)";
+ "-i", Arg.Set_int interval, " refresh interval in ms (default 250ms)";
+ "-q", Arg.Set(quiet), " don't print messages while running";
+ ]
+ let usage_msg = sprintf "usage: %s [options] -- <program> [arguments ...]" Sys.argv.(0)
+ let find_dash a =
+ let rec loop i len =
+ if i < len then
+ if a.(i) = "--" then i
+ else loop (i+1) len
+ else
+ raise Not_found
+ in
+ loop 0 (Array.length a)
+
+ let parse_cmdline () =
+ try
+ let idash = find_dash Sys.argv in
+ let argv = Array.sub Sys.argv 0 idash in
+ let cmd = Array.sub Sys.argv (idash+1) (Array.length Sys.argv - idash - 1) in
+ let () =
+ Arg.parse_argv argv spec ignore usage_msg
+ in
+ if Array.length cmd = 0 then (Arg.usage spec usage_msg; exit 1);
+ cmd
+ with
+ Arg.Bad(s) -> Printf.eprintf "%s\n" s; exit 1
+ | Arg.Help(s) -> Printf.printf "%s\n" s; exit 0
+ | Not_found -> Arg.usage spec usage_msg; exit 1
+
+ end
+
+
+(* Utility functions *)
+open Unix
+
+let get_page_size () =
+ let cin = open_process_in "getconf PAGESIZE" in
+ let s = input_line cin in
+ match close_process_in cin with
+ WEXITED 0 -> (try int_of_string s with _ -> ~-1)
+ | _ -> ~-1
;;
+
+
+let usleep ms = ignore (Unix.select [] [] [] (float_of_int ms /. 1000.))
+
+let read_proc_statm pid =
+ let cin = open_in (sprintf "/proc/%i/statm" pid) in
+ let m = ref ~-1 in
+ Scanf.fscanf cin "%i %i %s" (fun _ i _ -> m := i);
+ close_in cin;
+ !m
+
let max_mem = ref 0
-let rec monitor pid timeout mem =
+
+let rec monitor pid start finish mem fmt =
let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
if p == 0 then
- let current_mem = read_procmem pid in
+ let current_mem_pages = read_proc_statm pid in
+ let current_mem = current_mem_pages * !Options.page_size / 1024 in
if current_mem >= !max_mem then max_mem := current_mem;
- if (Unix.gettimeofday() > timeout)
- then let () = Printf.eprintf "Timeout reached, killing child process\n%!" in
- Unix.kill pid Sys.sigkill
- else if !max_mem >= mem
- then let () = Printf.eprintf "Memory limit reached, killing child process\n%!" in
+ let current_time = Unix.gettimeofday () in
+ if (current_time > finish)
+ then
+ let () = eprintf "Timeout reached, killing child process\n%!" in
Unix.kill pid Sys.sigkill
else
- let () = Unix.sleep 1 in
- monitor pid timeout mem
+ if !max_mem >= mem
+ then let () = eprintf "Memory limit reached, killing child process\n%!" in
+ Unix.kill pid Sys.sigkill
+ else
+ begin
+ if not !Options.quiet then
+ fprintf fmt "[% 11.3f] Memory: current=%ikb max=%ikb\n%!"
+ (current_time -. start) current_mem !max_mem;
+ usleep !Options.interval;
+ monitor pid start finish mem fmt
+ end
+ else
+ if not !Options.quiet then
+ let () =
+ match s with
+ WEXITED c -> fprintf fmt "Process exited with code %i\n" c
+ | WSTOPPED s -> fprintf fmt "Process stopped by signal %i\n" s
+ | WSIGNALED s -> fprintf fmt "Process killed by signal %i\n" s
+ in
+ fprintf fmt "Peak memory use: %ikb\n%!" !max_mem
;;
-let run args timeout mem =
+let run args start finish mem fmt =
let pid = Unix.fork () in
- if pid == 0 then
+ if pid == 0 then begin
+ eprintf "%s\n" args.(0);
Unix.execvp args.(0) args
- else monitor pid timeout mem
+ end
+ else monitor pid start finish mem fmt
;;
let () =
- if Array.length Sys.argv < 4 then exit 1
- else
- try
- let timeout = Unix.gettimeofday () +. float_of_string Sys.argv.(1) in
- let mem = int_of_string Sys.argv.(2) in
- let command = Array.sub Sys.argv 3 ((Array.length Sys.argv) - 3) in
- run command timeout mem;
- Printf.printf "Child process used %i kB of memory\n%!" !max_mem;
- exit 0
- with
- _ -> exit 2
+ let args = Options.parse_cmdline () in
+ if !Options.page_size <= 0 then Options.page_size := get_page_size ();
+ if !Options.page_size <= 0 then
+ eprintf "Warning: could not determine page size. Using 4096 (override with -p)\n%!";
+ try
+ let start = gettimeofday () in
+ let finish = start +. float_of_int !Options.timeout in
+ let cout = open_out !Options.outfile in
+ let fmt = formatter_of_out_channel cout in
+ run args start finish !Options.max_mem fmt;
+ close_out cout;
+ exit 0
+ with
+ Sys_error s -> eprintf "Error: %s\n%!" s; exit 2
+ | e -> eprintf "%s Unknown error\n%!" (Printexc.to_string e); exit 3
;;