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