From ae5b9f9b96a702ae8a8f61a31e5aa012283ce185 Mon Sep 17 00:00:00 2001 From: =?utf8?q?Kim=20Nguy=E1=BB=85n?= Date: Fri, 17 Feb 2012 18:02:05 +0100 Subject: [PATCH] New version of the alarm utility. --- utils/alarm.ml | 139 +++++++++++++++++++++++++++++++++++-------------- 1 file changed, 99 insertions(+), 40 deletions(-) diff --git a/utils/alarm.ml b/utils/alarm.ml index d1b7b30..5a9c83f 100644 --- a/utils/alarm.ml +++ b/utils/alarm.ml @@ -1,56 +1,115 @@ -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 args = ref [] + 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 argument s = args := s :: !args + let usage_msg = sprintf "usage: %s [options] [arguments ...]" Sys.argv.(0) + + let parse_cmdline () = + let _ = Arg.parse spec argument usage_msg in + match !args with + [] -> Arg.usage spec usage_msg; exit 1 + | l -> Array.of_list (List.rev l) + + 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 "[%f] Memory: current=%ikb max=%ikb\n%!" + (current_time -. start) current_mem !max_mem; + usleep !Options.interval; + monitor pid start finish mem fmt + end + else + 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 ;; -- 2.17.1