utils/alarm.ml

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