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       try
  33         let idash = find_dash Sys.argv in
  34         let argv = Array.sub Sys.argv 0 idash in
  35         let cmd = Array.sub Sys.argv (idash+1) (Array.length Sys.argv - idash - 1) in
  36         let () =
  37           Arg.parse_argv argv spec ignore usage_msg
  38         in
  39         if Array.length cmd = 0 then (Arg.usage spec usage_msg; exit 1);
  40         cmd
  41       with
  42         Arg.Bad(s) -> Printf.eprintf "%s\n" s; exit 1
  43       | Arg.Help(s) -> Printf.printf "%s\n" s; exit 0
  44       | Not_found -> Arg.usage spec usage_msg; exit 1
  45
  46   end
  47
  48
  49 (* Utility functions *)
  50 open Unix
  51
  52 let get_page_size () =
  53   let cin = open_process_in "getconf PAGESIZE" in
  54   let s = input_line cin in
  55     match close_process_in cin with
  56       WEXITED 0 -> (try int_of_string s with _ -> ~-1)
  57     | _ -> ~-1
  58 ;;
  59
  60
  61 let usleep ms = ignore (Unix.select [] [] [] (float_of_int ms /. 1000.))
  62
  63 let read_proc_statm pid =
  64   let cin = open_in (sprintf "/proc/%i/statm" pid) in
  65   let m = ref ~-1 in
  66   Scanf.fscanf cin "%i %i %s" (fun _ i _ -> m := i);
  67   close_in cin;
  68   !m
  69
  70 let max_mem = ref 0
  71
  72 let rec monitor pid start finish mem fmt =
  73   let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
  74   if p == 0 then
  75     let current_mem_pages = read_proc_statm pid in
  76     let current_mem = current_mem_pages * !Options.page_size / 1024 in
  77     if current_mem >= !max_mem then max_mem := current_mem;
  78     let current_time = Unix.gettimeofday () in
  79     if (current_time > finish)
  80     then
  81       let () = eprintf "Timeout reached, killing child process\n%!" in
  82          Unix.kill pid Sys.sigkill
  83     else
  84       if !max_mem >= mem
  85       then let () = eprintf "Memory limit reached, killing child process\n%!" in
  86            Unix.kill pid Sys.sigkill
  87       else
  88         begin
  89           if not !Options.quiet then
  90             fprintf fmt "[% 11.3f] Memory: current=%ikb max=%ikb\n%!"
  91               (current_time -. start) current_mem !max_mem;
  92           usleep !Options.interval;
  93           monitor pid start finish mem fmt
  94         end
  95   else
  96     if not !Options.quiet then
  97       let () =
  98         match s with
  99           WEXITED c -> fprintf fmt "Process exited with code %i\n" c
 100         | WSTOPPED s -> fprintf fmt "Process stopped by signal %i\n" s
 101         | WSIGNALED s -> fprintf  fmt "Process killed by signal %i\n" s
 102       in
 103       fprintf fmt "Peak memory use: %ikb\n%!" !max_mem
 104 ;;
 105
 106
 107 let run args start finish mem fmt =
 108   let pid = Unix.fork () in
 109   if pid == 0 then begin
 110     eprintf "%s\n" args.(0);
 111     Unix.execvp args.(0) args
 112   end
 113   else monitor pid start finish mem fmt
 114 ;;
 115
 116 let () =
 117   let args = Options.parse_cmdline () in
 118   if !Options.page_size <= 0 then Options.page_size := get_page_size ();
 119   if !Options.page_size <= 0 then
 120     eprintf "Warning: could not determine page size. Using 4096 (override with -p)\n%!";
 121   try
 122     let start = gettimeofday () in
 123     let finish = start +. float_of_int !Options.timeout in
 124     let cout = open_out !Options.outfile in
 125     let fmt = formatter_of_out_channel cout in
 126     run args start finish !Options.max_mem fmt;
 127     close_out cout;
 128     exit 0
 129   with
 130     Sys_error s -> eprintf "Error: %s\n%!" s; exit 2
 131   | e -> eprintf "%s Unknown error\n%!" (Printexc.to_string e); exit 3
 132 ;;