Source file batUnix.ml

(*
 * BatUnix - additional and modified functions for Unix and Unix-compatible systems.
 * Copyright (C) 1996 Xavier Leroy
 * Copyright (C) 2009 David Teller, LIFO, Universite d'Orleans
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version,
 * with the special exception on linking described in file LICENSE.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *)

##V>=5##module Pervasives = Stdlib

include Unix

##V<4.8##external link : string -> string -> unit = "unix_link"
##V>=4.8####V<5.0##external link : ?follow:bool -> string -> string -> unit = "unix_link"
##V>=5.0##external link : ?follow:bool -> string -> string -> unit = "caml_unix_link"

##V<4.2##let write_substring = write
##V<4.2##let single_write_substring = single_write
##V<4.2##let send_substring = send
##V<4.2##let sendto_substring = sendto

##V<4.3##let sleepf (timeout: float): unit =
##V<4.3##  let elapsed = ref 0.0 in
##V<4.3##  while !elapsed < timeout do
##V<4.3##    let start = gettimeofday () in
##V<4.3##    begin
##V<4.3##      try ignore(select [] [] [] (timeout -. !elapsed))
##V<4.3##      with Unix_error(EINTR, _, _) -> ()
##V<4.3##    end;
##V<4.3##    let stop = gettimeofday () in
##V<4.3##    let dt = stop -. start in
##V<4.3##    elapsed := !elapsed +. dt
##V<4.3##  done;
##V<4.3##  ()

(* chronometer is useful to test sleepf *)
(*$inject
let chronometer f =
  let start = gettimeofday () in
  let res = f () in
  let stop = gettimeofday () in
  let dt = stop -. start in
  (dt, res) ;;
*)

(* do not underestimate the imprecission of sleepf
   and so don't be too harsh when testing it *)
(*$T sleepf
let dt, _ = chronometer (fun () -> sleepf 0.002) in \
dt >= 0.002
*)

let run_and_read cmd =
  (* This code is before the open of BatInnerIO
     to avoid using batteries' wrapped IOs *)
  let string_of_file fn =
    let buff_size = 1024 in
    let buff = Buffer.create buff_size in
    let ic = open_in fn in
    let line_buff = Bytes.create buff_size in
    begin
      let was_read = ref (input ic line_buff 0 buff_size) in
      while !was_read <> 0 do
        BatBytesCompat.buffer_add_subbytes buff line_buff 0 !was_read;
        was_read := input ic line_buff 0 buff_size;
      done;
      close_in ic;
    end;
    Buffer.contents buff
  in
  let tmp_fn = Filename.temp_file "" "" in
  let cmd_to_run = cmd ^ " > " ^ tmp_fn in
  let status = Unix.system cmd_to_run in
  let output = string_of_file tmp_fn in
  Unix.unlink tmp_fn;
  (status, output)

(*$T run_and_read
  run_and_read "echo"            = (WEXITED 0, "\n")
  run_and_read "echo toto"       = (WEXITED 0, "toto\n")
  run_and_read "seq 1 3"         = (WEXITED 0, "1\n2\n3\n")
  run_and_read "printf 'abc'"    = (WEXITED 0, "abc")
*)

open BatInnerIO

(**
   {6 Thread-safety internals}
*)
let lock = ref BatConcurrent.nolock

(**
   {6 Tracking additional information on inputs/outputs}

   {b Note} Having [input]/[output] as objects would have made this
   easier. Here, we need to maintain an external weak hashtable to
   track low-level information on our [input]s/[output]s.
*)

module Wrapped_in  = BatInnerWeaktbl.Make(Input) (*input  -> in_channel *)
module Wrapped_out = BatInnerWeaktbl.Make(Output)(*output -> out_channel*)
let wrapped_in     = Wrapped_in.create 16
let wrapped_out    = Wrapped_out.create 16

let input_add k v =
  BatConcurrent.sync !lock (Wrapped_in.add wrapped_in k) v

let input_get k =
  BatConcurrent.sync !lock (Wrapped_in.find wrapped_in) k

let output_add k v =
  BatConcurrent.sync !lock (Wrapped_out.add wrapped_out k) v

let output_get k =
  BatConcurrent.sync !lock (Wrapped_out.find wrapped_out) k

let wrap_in ?autoclose ?cleanup cin =
  let input = BatInnerIO.input_channel ?autoclose ?cleanup cin in
  BatConcurrent.sync !lock (Wrapped_in.add wrapped_in input) cin;
  input

let wrap_out ?cleanup cout =
  let output = cast_output (BatInnerIO.output_channel ?cleanup cout) in
  BatConcurrent.sync !lock (Wrapped_out.add wrapped_out output) cout;
  output

let _ =
  input_add stdin Pervasives.stdin;
  output_add stdout Pervasives.stdout;
  output_add stderr Pervasives.stderr

(**
   {6 File descriptors}
*)

let input_of_descr ?autoclose ?cleanup fd =
  wrap_in ?autoclose ?cleanup (in_channel_of_descr fd)

let descr_of_input cin =
  try  descr_of_in_channel (input_get cin)
  with Not_found -> invalid_arg "Unix.descr_of_input"

let output_of_descr ?cleanup fd =
  wrap_out ?cleanup (out_channel_of_descr fd)

let descr_of_output cout =
  try  descr_of_out_channel (output_get (cast_output cout))
  with Not_found -> invalid_arg "Unix.descr_of_output"

let in_channel_of_descr fd = input_of_descr ~autoclose:false ~cleanup:true fd
let descr_of_in_channel = descr_of_input
let out_channel_of_descr fd = output_of_descr  ~cleanup:true fd
let descr_of_out_channel = descr_of_output


(**
   {6 Processes}
*)

let open_process_in ?autoclose ?(cleanup=true) s =
  wrap_in ?autoclose ~cleanup (open_process_in s)

let open_process_out ?(cleanup=true) s =
  wrap_out ~cleanup (open_process_out s)

let open_process ?autoclose ?(cleanup=true) s =
  let (cin, cout) = open_process s in
  (wrap_in ?autoclose cin, wrap_out ~cleanup cout)

(*$T open_process
  let s = "hello world" in let r,w = open_process "cat" in \
  Printf.fprintf w "%s\n" s; IO.close_out w; \
  IO.read_line r = s

  try \
    let r,w = open_process "cat" in \
      Printf.fprintf w "hello world\n"; \
      IO.close_out w; \
      while true do ignore (input_char r) done; false \
  with e -> e=IO.No_more_input || e=End_of_file
*)


let open_process_full ?autoclose ?(cleanup=true) s args =
  let (a,b,c) = open_process_full s args in
  (wrap_in ?autoclose ~cleanup a, wrap_out ~cleanup b, wrap_in ?autoclose ~cleanup c)

(**@TODO in a future version, [close_process_in] should also work
   on processes opened with [open_process] or [open_process_full].
   Same thing for [close_process_out].*)

let close_process_in cin =
  try close_process_in (input_get cin)
  with Not_found -> raise (Unix_error(EBADF, "close_process_in", ""))

let close_process_out cout =
  try close_process_out (output_get cout)
  with Not_found -> raise (Unix_error(EBADF, "close_process_out", ""))

let close_process (cin, cout) =
  try
    let pin = input_get cin
    and pout= output_get cout in
    close_process (pin, pout)
  with Not_found -> raise (Unix_error(EBADF, "close_process", ""))

let close_process_full (cin, cout, cin2) =
  try close_process_full (input_get cin, output_get cout, input_get cin2)
  with Not_found -> raise (Unix_error(EBADF, "close_process_full", ""))

(**
   {6 Network}
*)

let shutdown_connection cin =
  try shutdown_connection (input_get cin)
  with Not_found -> invalid_arg "Unix.shutdown_connection"

let open_connection ?autoclose addr =
  let (cin, cout) = open_connection addr in
  let (cin',cout')= (wrap_in ?autoclose ~cleanup:true cin, wrap_out ~cleanup:true cout) in
  let close ()    = shutdown_connection cin' in
  (inherit_in  cin' ~close, inherit_out cout' ~close)

let establish_server ?autoclose ?cleanup f addr =
  let f' cin cout = f (wrap_in ?autoclose ?cleanup cin) (wrap_out cout) in
  establish_server f' addr

(**
   {6 Tools}
*)

let is_directory fn = (lstat fn).st_kind = S_DIR

let rec restart_on_EINTR f x =
  try f x
  with Unix_error(EINTR, _, _) -> restart_on_EINTR f x

(**
   {6 Locking}
*)

let with_locked_file ~kind filename f =
  let perms =
    [O_CREAT ; match kind with `Read -> O_RDONLY | `Write -> O_RDWR] in
  let lock_file = openfile filename perms 0o644 in
  let lock_action = match kind with
    | `Read -> F_RLOCK
    | `Write -> F_LOCK
  in
  lockf lock_file lock_action 0;
  BatInnerPervasives.finally
    (fun () ->
      (* Although the user might expect EINTR to interrupt locking, we must
       * not allow such interrupt here since there is no way to restart the
       * unlock: *)
      ignore (restart_on_EINTR (lseek lock_file 0) SEEK_SET);
      restart_on_EINTR (lockf lock_file F_ULOCK) 0;
      restart_on_EINTR close lock_file)
    f lock_file