(** {2 Control} *)

(** Execute f and then g, no matter what.  Returns the result of f.
  If f raises an exception catch it, execute g, and then re-raise the
  exception.  Returns the value returned by f.  This is a good place
  to do things like close open file descriptors.  A secondary intent
  of this code is to get a nice traceback, but note that it is
  possible for the execution of g to obliterate the existing
  traceback.  More generally, if Unix.chdir is executed and argv[0] is
  a relative pathname, no traceback is possible because the system
  can't locate the executable. *)
let and_then f g =
  let res =
    try f ()
    with
      exn ->
	g ();
	raise exn
  in
  g ();
  res

let stopwatch thunk =
  let start = Unix.gettimeofday () in
  let result = thunk () in
  let stop = Unix.gettimeofday ()  in
  stop -. start, result;; 

(** Return the first index between min and max for which (f i) is true,
  or raise Not_found. *)
let index f min max =
  let rec loop i =
    if i > max
    then raise Not_found
    else if f i
    then i
    else loop (i + 1) in
  loop min

(* Perform BEFORE and AFTER actions surrounding THUNK, subject to
   these rules:

   1. If BEFORE succeeds, always do AFTER
   2. If BEFORE fails, never do THUNK or AFTER

   Example:

   let with_proc_mounted f =
     if Sys.file_exists "/proc/mounts"
     then f
     else wrap (fun () -> Sys.command "mount -t proc none /proc") f
	       (fun () -> Sys.command "umount /proc")
*)

exception Before_failed of exn
exception After_failed of exn

let wrap before thunk after =
  try
    (try before () with exn -> raise (Before_failed exn));
    let result = thunk () in
    (try after () with exn -> raise (After_failed exn));
    result
  with
    (* If BEFORE fails, skip THUNK and AFTER. *)
    Before_failed exn -> raise exn
      (* If AFTER fails, its a bad thing, but there's nothing for it. *)
  | After_failed exn -> raise exn
  | exn ->
      (* If THUNK fails, run after and then re-raise exception *)
      (try after () with _ -> ());
      raise exn

let pid_re = Pcre.regexp "[1-9][0-9]*"

(** [waitfn count ready] - This function performs .01 second wait.
    The COUNT parameters allow a user supplied version to give up
    after a certain number of tries, while the READY flag allows the
    function to log a message about how long it had to wait. *)
let waitfn count ready =
  let wait = 0.01 in
  if not ready
  then ignore (Unix.select [] [] [] wait)

exception Timeout

(** [do_locked ?timeout ?waitfn lockfile f] *)
let do_locked ?timeout ?(waitfn=waitfn) lockfile f =
  let rec loop count =
    (match timeout with Some maxcount -> if count > maxcount then raise Timeout else () | _ -> ());
    let pid = string_of_int (Unix.getpid ()) in
    let len = String.length pid in
    try
      (* Try to create the file.  The O_EXCL flag means that if
	 this function succeeds we were able to create the file
	 and no other process was. *)
      let fd = Unix.openfile lockfile [Unix.O_CREAT; Unix.O_EXCL; Unix.O_WRONLY] 0o644 in
      if Unix.write fd pid 0 len != len then failwith "Couldn't write lock file";
      Unix.close fd;
      (* This call won't do a wait, but it could write out
	 a log message about how long it waited. *)
      waitfn count true;
      let result = f () in
      Unix.unlink lockfile;
      result
    with
      Unix.Unix_error (Unix.EEXIST, _, _) ->
	(* We couldn't create the lock file, see if the process that
	   owns it still exists. *)
	let pid = File.to_string lockfile in
	if Pcre.pmatch ~rex:pid_re pid && Sys.file_exists ("/proc/" ^ pid) then begin
	  waitfn count false;
	  loop (count + 1)
	end
	else begin
	  (* The process holding the lock died, break lock and retry *)
	  Unix.unlink lockfile;
	  loop count
	end
(*  | Unix.Unix_error (error, arg1, arg2) as exn ->
	prerr_endline (Unix.error_message error);
	raise exn*)
  in
  loop 0