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 " VmRSS: %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         Printf.eprintf "Memory: %i\n%!" s;
  16   close_in cin;
  17   s
  18 ;;
  19 let max_mem = ref 0
  20 let rec monitor pid timeout mem =
  21   let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
  22   if p == 0 then
  23     let current_mem = read_procmem pid in
  24     if current_mem >= !max_mem then max_mem := current_mem;
  25     if (Unix.gettimeofday() > timeout)
  26     then let () = Printf.eprintf "Timeout reached, killing child process\n%!" in
  27          Unix.kill pid Sys.sigkill
  28     else if !max_mem >= mem
  29     then let () = Printf.eprintf "Memory limit reached, killing child process\n%!" in
  30          Unix.kill pid Sys.sigkill
  31     else
  32       let () = Unix.sleep 1 in
  33       monitor pid timeout mem
  34 ;;
  35
  36
  37 let run args timeout mem =
  38   let pid = Unix.fork () in
  39   if pid == 0 then
  40     Unix.execvp args.(0) args
  41   else monitor pid timeout mem
  42 ;;
  43
  44 let () =
  45   if Array.length Sys.argv < 4 then exit 1
  46   else
  47     try
  48       let timeout = Unix.gettimeofday () +. float_of_string Sys.argv.(1) in
  49       let mem = int_of_string Sys.argv.(2) in
  50       let command = Array.sub Sys.argv 3 ((Array.length Sys.argv) - 3) in
  51       run command timeout mem;
  52       Printf.printf "Child process used %i kB of memory\n%!" !max_mem;
  53       exit 0
  54     with
  55         _ -> exit 2
  56 ;;