utils/alarm.ml

   1 open Format
   2
   3 module Options =
   4   struct
   5     let outfile = ref "/dev/stderr"
   6     let page_size = ref ~-1
   7     let timeout = ref 30
   8     let max_mem = ref (512*1024)
   9     let interval = ref 250
  10     let quiet = ref false
  11
  12     let spec = Arg.align [
  13       "-o", Arg.Set_string outfile, " output file (default stderr)";
  14       "-p", Arg.Set_int page_size, " page size in bytes (default autodetect)";
  15       "-t", Arg.Set_int timeout, " timeout in s (default 30s)";
  16       "-m", Arg.Set_int max_mem, " maximum memory in kiB (default 512 MiB)";
  17       "-i", Arg.Set_int interval, " refresh interval in ms (default 250ms)";
  18       "-q", Arg.Set(quiet), " don't print messages while running";
  19     ]
  20     let usage_msg = sprintf "usage: %s [options] -- <program> [arguments ...]" Sys.argv.(0)
  21     let find_dash a =
  22       let rec loop i len =
  23         if i < len then
  24           if a.(i) = "--" then i
  25           else loop (i+1) len
  26         else
  27           raise Not_found
  28       in
  29       loop 0 (Array.length a)
  30
  31     let parse_cmdline () =
  32       let idash = find_dash Sys.argv in
  33       let argv = Array.sub Sys.argv 0 idash in
  34       let cmd = Array.sub Sys.argv (idash+1) (Array.length Sys.argv - idash - 1) in
  35       let () =
  36         try
  37           Arg.parse_argv argv spec ignore usage_msg
  38         with
  39           Arg.Bad(s) -> Printf.eprintf "%s\n" s; exit 1
  40         | Arg.Help(s) -> Printf.printf "%s\n" s; exit 0
  41         | Not_found -> Arg.usage spec usage_msg; exit 1
  42       in
  43       if Array.length cmd = 0 then (Arg.usage spec usage_msg; exit 1);
  44       cmd
  45   end
  46
  47
  48 (* Utility functions *)
  49 open Unix
  50
  51 let get_page_size () =
  52   let cin = open_process_in "getconf PAGESIZE" in
  53   let s = input_line cin in
  54     match close_process_in cin with
  55       WEXITED 0 -> (try int_of_string s with _ -> ~-1)
  56     | _ -> ~-1
  57 ;;
  58
  59
  60 let usleep ms = ignore (Unix.select [] [] [] (float_of_int ms /. 1000.))
  61
  62 let read_proc_statm pid =
  63   let cin = open_in (sprintf "/proc/%i/statm" pid) in
  64   let m = ref ~-1 in
  65   Scanf.fscanf cin "%i %i %s" (fun _ i _ -> m := i);
  66   close_in cin;
  67   !m
  68
  69 let max_mem = ref 0
  70
  71 let rec monitor pid start finish mem fmt =
  72   let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
  73   if p == 0 then
  74     let current_mem_pages = read_proc_statm pid in
  75     let current_mem = current_mem_pages * !Options.page_size / 1024 in
  76     if current_mem >= !max_mem then max_mem := current_mem;
  77     let current_time = Unix.gettimeofday () in
  78     if (current_time > finish)
  79     then
  80       let () = eprintf "Timeout reached, killing child process\n%!" in
  81          Unix.kill pid Sys.sigkill
  82     else
  83       if !max_mem >= mem
  84       then let () = eprintf "Memory limit reached, killing child process\n%!" in
  85            Unix.kill pid Sys.sigkill
  86       else
  87         begin
  88           if not !Options.quiet then
  89             fprintf fmt "[% 11.3f] Memory: current=%ikb max=%ikb\n%!"
  90               (current_time -. start) current_mem !max_mem;
  91           usleep !Options.interval;
  92           monitor pid start finish mem fmt
  93         end
  94   else
  95     let () =
  96       match s with
  97         WEXITED c -> fprintf fmt "Process exited with code %i\n" c
  98       | WSTOPPED s -> fprintf fmt "Process stopped by signal %i\n" s
  99       | WSIGNALED s -> fprintf  fmt "Process killed by signal %i\n" s
 100     in
 101     fprintf fmt "Peak memory use: %ikb\n%!" !max_mem
 102 ;;
 103
 104
 105 let run args start finish mem fmt =
 106   let pid = Unix.fork () in
 107   if pid == 0 then begin
 108     eprintf "%s\n" args.(0);
 109     Unix.execvp args.(0) args
 110   end
 111   else monitor pid start finish mem fmt
 112 ;;
 113
 114 let () =
 115   let args = Options.parse_cmdline () in
 116   if !Options.page_size <= 0 then Options.page_size := get_page_size ();
 117   if !Options.page_size <= 0 then
 118     eprintf "Warning: could not determine page size. Using 4096 (override with -p)\n%!";
 119   try
 120     let start = gettimeofday () in
 121     let finish = start +. float_of_int !Options.timeout in
 122     let cout = open_out !Options.outfile in
 123     let fmt = formatter_of_out_channel cout in
 124     run args start finish !Options.max_mem fmt;
 125     close_out cout;
 126     exit 0
 127   with
 128     Sys_error s -> eprintf "Error: %s\n%!" s; exit 2
 129   | e -> eprintf "%s Unknown error\n%!" (Printexc.to_string e); exit 3
 130 ;;