Source

spotlib / lib / xunix.ml

Full commit
open Base
open Unix

(* run [f] on files in [path] *)
let folddir ~f ~init path =
  let dh = opendir path in
  Exn.protect' (fun () ->
    let rec loop st =
      try
        let st' = f st (readdir dh) in
        loop st'
      with
      | End_of_file -> st
    in
    loop init)
    ~finally:(fun _ -> closedir dh)


module Inodes = Set.Make(struct
  type t = int * int
  let compare : t -> t -> int = compare
end)

module Find = struct

  exception Prune

  class type path = object
    method base : string
    method depth : int
    method dev_inode : (int * int, exn) Result.t
    method dir : string
    method is_dir : bool
    method is_ldir : bool
    method is_reg : bool
    method kind : (Unix.file_kind, exn) Result.t
    method lkind : (Unix.file_kind, exn) Result.t
    method path : string
    method stat : (Unix.stats, exn) Result.t
    method lstat : (Unix.stats, exn) Result.t
  end

  class path_ ~dir ~base ~depth = 
    let path = match Filename.concat dir base with
      | "./." -> "."
      | s -> s
    in
    object (self)
      method dir = dir
      method base = base
      method path = path
      method depth : int = depth
      method stat : (_,_) Result.t = try `Ok (stat path) with e -> `Error e
      method lstat : (_,_) Result.t = try `Ok (lstat path) with e -> `Error e
      method kind : (_,_) Result.t = match self#stat with
        | `Error _exn -> `Error _exn
        | `Ok stat -> `Ok stat.st_kind
      method lkind : (_,_) Result.t = match self#lstat with
        | `Error _exn -> `Error _exn
        | `Ok stat -> `Ok stat.st_kind
      method is_dir = self#kind = `Ok S_DIR
      method is_ldir = self#lkind = `Ok S_DIR
      method is_reg = self#kind = `Ok S_REG
      method dev_inode : (_,_) Result.t = match self#stat with
      | `Ok stat -> `Ok (stat.st_dev, stat.st_ino)
      | `Error e -> `Error e
    end

  let prune () = raise Prune
  
  let find ?(follow_symlink=false) ~f fnames =
  
    (* visited cache *)
    let visited = ref Inodes.empty in
    let if_not_visited_then path ~f = match path#dev_inode with
      | `Error _ -> ()
      | `Ok inode ->
          if Inodes.mem inode !visited then ()
          else begin
            visited := Inodes.add inode !visited;
            f path
          end
    in
  
    let rec find_dir pth =
      try 
        f pth;
        let subdirs =
          folddir pth#path ~init:[] ~f:(fun dirs -> function
            | "." | ".." -> dirs
            | name -> 
        	let pth = new path_ ~depth:(pth#depth + 1) ~dir:pth#path ~base:name in
        	if try if follow_symlink then pth#is_dir else pth#is_ldir with _ -> false then pth::dirs
        	else begin find_non_dir pth; dirs end)
        in
        List.iter (if_not_visited_then ~f:find_dir) subdirs
      with
      | Prune -> ()
  
    and find_non_dir path = try f path with Prune -> ()
      (* Even if path is a dangling symlink, f path is called *)
      (* CR jfuruse: if the initial argument contains non existent files,
         they reach here. *)
    in
  
    List.iter (fun fname ->
      let path = 
        new path_ ~depth: 0 ~dir:(Filename.dirname fname) ~base:(Filename.basename fname)
      in
      if path#is_dir then find_dir path
      else find_non_dir path) fnames
  
end

let try_set_close_on_exec fd =
  try set_close_on_exec fd; true with Invalid_argument _ -> false

    
module Process_status = struct

  type t = [ `Exited of int | `Signaled of int | `Stopped of int ]  

  let failwith ?name = 
    let name = match name with None -> "" | Some n -> n ^ ": " in
    function
      | `Exited n   -> Exn.failwithf "%sprocess exited with id %d" name n
      | `Signaled n -> Exn.failwithf "%sprocess killed by signal %d" name n
      | `Stopped n  -> Exn.failwithf "%sprocess stopped by signal %d" name n

  let convert = function
    | WEXITED n   -> `Exited n
    | WSIGNALED n -> `Signaled n
    | WSTOPPED n  -> `Stopped n
end

let open_proc_full cmdargs input output error toclose =
  let cmd = match cmdargs with
    | x :: _ -> x
    | _ -> invalid_arg "Xunix.gen_open_proc_full"
  in
  let cmdargs = Array.of_list cmdargs in
  let cloexec = List.for_all try_set_close_on_exec toclose in
  match fork() with
    0 ->
      dup2 input stdin; close input;
      dup2 output stdout; close output;
      dup2 error stderr; close error;
      if not cloexec then List.iter close toclose;
      begin try execvp cmd cmdargs with _ -> exit 127
      end (* never return *)
  | id -> id


let open_process_full cmdargs =
  let (in_read, in_write) = pipe() in
  let (out_read, out_write) = pipe() in
  let (err_read, err_write) = pipe() in
  let pid = open_proc_full cmdargs
    out_read in_write err_write [in_read; out_write; err_read]
  in
  close out_read;
  close in_write;
  close err_write;
  pid, (in_read, out_write, err_read)

    
let open_shell_process_full cmd = open_process_full [ "/bin/sh"; "-c"; cmd ]


let buf_flush_limit = 100000

    
let command_aux readers stat =
  let read_buflen = 4096 in
  let read_buf = String.create read_buflen in

  let try_read_lines fd buf : (string list * bool (* eof *)) =
    let read_bytes = 
      try Some (read fd read_buf 0 read_buflen) with
      | Unix_error ((EAGAIN | EWOULDBLOCK), _, _) -> None
    in
    match read_bytes with
    | None -> [], false
    | Some 0 -> (* eof *)
        let s = Buffer.contents buf in
        (if s = "" then [] else [s]), true
    | Some len ->
        let buffer_old_len = Buffer.length buf in
        Buffer.add_substring buf read_buf 0 len;

        let pos_in_buffer pos = buffer_old_len + pos in
        
        let rec get_lines st from_in_buffer pos =  
          match
            if pos >= len then None
            else Xstring.index_from_to read_buf pos (len-1) '\n'
          with
          | None ->
              let rem =
                Buffer.sub buf
                  from_in_buffer
                  (Buffer.length buf - from_in_buffer)
              in
              Buffer.clear buf;
              if String.length rem > buf_flush_limit then rem :: st
              else begin
                Buffer.add_string buf rem; st
              end
          | Some pos ->
              let next_from_in_buffer = pos_in_buffer pos + 1 in
              let line =
                Buffer.sub buf
                  from_in_buffer
                  (next_from_in_buffer - from_in_buffer)
              in
              get_lines (line :: st) next_from_in_buffer (pos + 1)
        in
        List.rev (get_lines [] 0 0), false
  in

  let rec loop readers stat =
    if readers = [] then stat (* no more reader and no need to loop *)
    else begin
      let fds = List.map (fun (fd, _, _) -> fd) readers in 
      let readables, _, _ = select fds [] [](*?*) (-1.0)(*?*) in
      let readers', stat = 
        List.fold_right (fun (fd, buf, fs as reader) (st, stat) ->
          if not (List.mem fd readables) then
            (reader :: st, stat)
          else begin
            let rec loop stat =
              let lines, is_eof = try_read_lines fd buf in
              if lines <> [] then begin
                let stat = 
                  List.fold_left (fun stat line ->
                    List.fold_left (fun stat f -> f stat (`Read line)) stat fs) stat lines
                in
                if not is_eof then loop stat else is_eof, stat
              end else is_eof, stat 
            in
            match loop stat with
            | true (*eof*), stat ->
	        (* reached eof. remove the reader *)
	        let stat = List.fold_left (fun stat f -> f stat `EOF) stat fs in
                close fd; 
	        st, stat
            | false, stat -> reader :: st, stat
          end) readers ([], stat)
      in
      loop readers' stat
    end
  in
  loop readers stat


let rec waitpid_non_intr pid =
  try 
    waitpid [] pid 
  with Unix_error (EINTR, _, _) -> waitpid_non_intr pid

    
let command_wrapper (pid, (out, in_, err)) f ~init:stat =
  try
    close in_;
    set_nonblock out;
    set_nonblock err;
    
    let buf_out = Buffer.create buf_flush_limit in
    let buf_err = Buffer.create buf_flush_limit in

    let stat = command_aux
      [out, buf_out, [fun stat s -> f stat (`Out, s)];
       err, buf_err, [fun stat s -> f stat (`Err, s)]] stat 
    in
    Process_status.convert & snd & waitpid_non_intr pid, 
    stat
  with
  | e ->
      (* kill really ? *)
      kill pid 9;
      ignore (waitpid_non_intr pid);
      raise e


let print_all = function
  | `Err, `Read s -> prerr_endline & Xstring.chop_newline s
  | `Out, `Read s -> print_endline & Xstring.chop_newline s
  | _ -> ()

let command ?(f=print_all) cmd = 
  fst & command_wrapper (open_process_full cmd) (fun () -> f) ~init:()

let shell_command ?(f=print_all) cmd = 
  fst & command_wrapper (open_shell_process_full cmd) (fun () -> f) ~init:() 

let command' cmd = command_wrapper (open_process_full cmd)
let shell_command' cmd = command_wrapper (open_shell_process_full cmd)

let shell_command_stdout cmd = 
  match 
    command_wrapper (open_shell_process_full cmd) ~init:[] (fun rev -> function 
      | `Err, `Read s -> prerr_endline & Xstring.chop_newline s; rev
      | `Out, `Read s ->  s :: rev
      | _ -> rev)
  with
  | stat, rev -> stat, List.rev rev


let cmp p1 p2 = 
  match command ["cmp"; p1; p2] ~f:ignore with
  | `Exited 0 -> `Same
  | `Exited 1 -> `Different
  | `Exited 2 -> `Error
  | _ -> `Error (* something extremely wrong happened *)

let gen_timed get minus f v = 
  let t1 = get () in
  let res = f v  in
  let t2 = get () in
  res, minus t2 t1

let timed f v = gen_timed Unix.gettimeofday (-.) f v

module Process_times = struct
  type t = process_times
  let (-) pt1 pt2 = {
    tms_utime = pt1.tms_utime  -. pt2.tms_utime;
    tms_stime = pt1.tms_stime  -. pt2.tms_stime;
    tms_cutime = pt1.tms_utime -. pt2.tms_cutime;
    tms_cstime = pt1.tms_utime -. pt2.tms_cstime;
  }
  let timed f v = gen_timed Unix.times (-) f v
end

let file path = 
  match 
    command' ["/usr/bin/file"; path] ~init:[] & fun revls -> function
      | `Out, `Read s -> s::revls
      | _ -> revls
  with
  | `Exited 0, [] -> `Error `Empty_result
  | `Exited 0, lines -> `Ok (Xlist.last lines)
  | st, _ -> `Error st

let mkdir ?(perm=0o700) s =
  match File.Test._d' s with
  | `Error ENOENT -> 
      begin try
	mkdir s perm; (* CR jfuruse: use umask? *)
	`Ok
      with
      | Unix_error (e,_,_) -> `Error e
      end
  | `TRUE st -> `Already_exists st (* CR jfuruse: perm check ? *)
  | `FALSE st -> `Not_a_directory st
  | `Error e -> `Error e


let mkdtemp template =
  match Xstring.is_postfix' "XXXXXX" template with
  | None -> 
      Exn.invalid_argf "Unix.mkdtemp must take an argument whose postfix is \"XXXXXX\""
  | Some prefix ->
      let rec find () =
        let d = !% "%s%06d" prefix & Random.int 1000000 in
        if Sys.file_exists d then find ()
        else d
      in
      let d = find () in
      Unix.mkdir d 0o700;
      d

let with_dtemp template f =
  let d = mkdtemp template in
  Exn.protect' (fun () -> f d) ~finally:(fun () ->
    if ksprintf Sys.command "/bin/rm -rf %s" d <> 0 then
      Exn.failwithf "Unix.with_dtemp: cleaning tempdir %s failed" d)
  
let with_chdir ?(at_failure=(fun exn -> raise exn)) dir f =
  let cwd = Unix.getcwd () in
  match Exn.catch ~f:Unix.chdir dir with
  | `Error exn -> at_failure exn
  | `Ok () ->
      Exn.protect' f ~finally:(fun () -> Unix.chdir cwd)