Fix bug in commandline parsing.
[SXSI/xpathcomp.git] / utils / alarm.ml
index d1b7b30..0eeb823 100644 (file)
-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
 ;;