utils/alarm.ml

   1
   2 let read_procmem pid =
   3   let cin = open_in (Printf.sprintf "/proc/%i/status" pid) in
   4   let matchline s =
   5     try
   6       Scanf.sscanf s " VmHWM: %i kB" (fun i -> Some i)
   7     with
   8       | _ -> None
   9   in
  10   let rec loop () =
  11     match matchline (input_line cin) with
  12         Some i -> i
  13       | None -> loop ()
  14   in
  15   let s = try loop() with _ -> -1 in
  16   close_in cin;
  17   s
  18 ;;
  19
  20 let rec monitor pid timeout mem =
  21   let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
  22   if p == 0 then
  23     if (Unix.gettimeofday() > timeout || (read_procmem pid) >= mem)
  24     then Unix.kill pid Sys.sigkill
  25     else
  26       let () = Unix.sleep 1 in
  27       monitor pid timeout mem
  28 ;;
  29
  30
  31 let run args timeout mem =
  32   let pid = Unix.fork () in
  33   if pid == 0 then
  34     Unix.execvp args.(0) args
  35   else monitor pid timeout mem
  36 ;;
  37
  38 let () =
  39   if Array.length Sys.argv < 4 then exit 1
  40   else
  41     try
  42       let timeout = Unix.gettimeofday () +. float_of_string Sys.argv.(1) in
  43       let mem = int_of_string Sys.argv.(2) in
  44       let command = Array.sub Sys.argv 3 ((Array.length Sys.argv) - 3) in
  45       run command timeout mem;
  46       exit 0
  47     with
  48         _ -> exit 2
  49 ;;