Source

spotlib / lib / xunix.ml

The default branch has multiple heads

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)
;;

type path = 
    { dir : string;
      base : string;
      path : string; (* dir / name *)
      stat : [ `Ok of stats | `Error of exn ];
      depth : int;
    }

let path ~depth ~dir base =
  let path = match Filename.concat dir base with
    | "./." -> "."
    | s -> s
  in
  { dir = dir;
    base = base;
    path = path;
    depth = depth; 
    stat = try `Ok (stat path) with e -> `Error e;
  }
;;

let kind path =
  match path.stat with
  | `Error _exn -> None
  | `Ok stat -> Some stat.st_kind
;;

let is_dir path = kind path = Some S_DIR

let dev_inode path =
  match path.stat with
  | `Ok stat -> Some (stat.st_dev, stat.st_ino)
  | `Error _ -> None

exception Prune

let prune () = raise Prune

let find ~f fnames =

  (* visited cache *)
  let visited = ref Inodes.empty in
  let if_not_visited_then path ~f = match dev_inode path with
    | None -> ()
    | Some 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 = path ~depth:(pth.depth + 1) ~dir:pth.path name in
      	if try is_dir pth 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 *)
  in

  List.iter (fun fname ->
    let path = 
      path ~depth: 0 ~dir:(Filename.dirname fname) (Filename.basename fname)
    in
    if is_dir path then find_dir path
    else find_non_dir path) fnames
;;

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)