open Format

module Options =
  struct
    let outfile = ref "/dev/stderr"
    let page_size = ref ~-1
    let timeout = ref 30
    let max_mem = ref (512*1024)
    let interval = ref 250
    let quiet = ref false

    let spec = Arg.align [
      "-o", Arg.Set_string outfile, " output file (default stderr)";
      "-p", Arg.Set_int page_size, " page size in bytes (default autodetect)";
      "-t", Arg.Set_int timeout, " timeout in s (default 30s)";
      "-m", Arg.Set_int max_mem, " maximum memory in kiB (default 512 MiB)";
      "-i", Arg.Set_int interval, " refresh interval in ms (default 250ms)";
      "-q", Arg.Set(quiet), " don't print messages while running";
    ]
    let usage_msg = sprintf "usage: %s [options] -- <program> [arguments ...]" Sys.argv.(0)
    let find_dash a =
      let rec loop i len =
	if i < len then
	  if a.(i) = "--" then i
	  else loop (i+1) len
	else
	  raise Not_found
      in
      loop 0 (Array.length a)

    let parse_cmdline () =
      try
	let idash = find_dash Sys.argv in
	let argv = Array.sub Sys.argv 0 idash in
	let cmd = Array.sub Sys.argv (idash+1) (Array.length Sys.argv - idash - 1) in
	let () =
	  Arg.parse_argv argv spec ignore usage_msg
	in
	if Array.length cmd = 0 then (Arg.usage spec usage_msg; exit 1);
	cmd
      with
	Arg.Bad(s) -> Printf.eprintf "%s\n" s; exit 1
      | Arg.Help(s) -> Printf.printf "%s\n" s; exit 0
      | Not_found -> Arg.usage spec usage_msg; exit 1

  end


(* Utility functions *)
open Unix

let get_page_size () =
  let cin = open_process_in "getconf PAGESIZE" in
  let s = input_line cin in
    match close_process_in cin with
      WEXITED 0 -> (try int_of_string s with _ -> ~-1)
    | _ -> ~-1
;;


let usleep ms = ignore (Unix.select [] [] [] (float_of_int ms /. 1000.))

let read_proc_statm pid =
  let cin = open_in (sprintf "/proc/%i/statm" pid) in
  let m = ref ~-1 in
  Scanf.fscanf cin "%i %i %s" (fun _ i _ -> m := i);
  close_in cin;
  !m

let max_mem = ref 0

let rec monitor pid start finish mem fmt =
  let p, s = Unix.waitpid [ Unix.WNOHANG ] pid in
  if p == 0 then
    let current_mem_pages = read_proc_statm pid in
    let current_mem = current_mem_pages * !Options.page_size / 1024 in
    if current_mem >= !max_mem then max_mem := current_mem;
    let current_time = Unix.gettimeofday () in
    if (current_time > finish)
    then
      let () = eprintf "Timeout reached, killing child process\n%!" in
	 Unix.kill pid Sys.sigkill
    else
      if !max_mem >= mem
      then let () = eprintf "Memory limit reached, killing child process\n%!" in
	   Unix.kill pid Sys.sigkill
      else
	begin
	  if not !Options.quiet then
	    fprintf fmt "[% 11.3f] Memory: current=%ikb max=%ikb\n%!"
	      (current_time -. start) current_mem !max_mem;
	  usleep !Options.interval;
	  monitor pid start finish mem fmt
	end
  else
    if not !Options.quiet then
      let () =
	match s with
	  WEXITED c -> fprintf fmt "Process exited with code %i\n" c
	| WSTOPPED s -> fprintf fmt "Process stopped by signal %i\n" s
	| WSIGNALED s -> fprintf  fmt "Process killed by signal %i\n" s
      in
      fprintf fmt "Peak memory use: %ikb\n%!" !max_mem
;;


let run args start finish mem fmt =
  let pid = Unix.fork () in
  if pid == 0 then begin
    eprintf "%s\n" args.(0);
    Unix.execvp args.(0) args
  end
  else monitor pid start finish mem fmt
;;

let () =
  let args = Options.parse_cmdline () in
  if !Options.page_size <= 0 then Options.page_size := get_page_size ();
  if !Options.page_size <= 0 then
    eprintf "Warning: could not determine page size. Using 4096 (override with -p)\n%!";
  try
    let start = gettimeofday () in
    let finish = start +. float_of_int !Options.timeout in
    let cout = open_out !Options.outfile in
    let fmt = formatter_of_out_channel cout in
    run args start finish !Options.max_mem fmt;
    close_out cout;
    exit 0
  with
    Sys_error s -> eprintf "Error: %s\n%!" s; exit 2
  | e -> eprintf "%s Unknown error\n%!" (Printexc.to_string e); exit 3
;;