ocaml-containers/src/unix/CCUnix.ml
2017-03-07 09:37:31 +01:00

276 lines
7.5 KiB
OCaml

(* This file is free software, part of containers. See file "license" for more details. *)
(** {1 High-level Functions on top of Unix} *)
type 'a or_error = ('a, string) Result.result
type 'a gen = unit -> 'a option
(** {2 Calling Commands} *)
let int_of_process_status = function
| Unix.WEXITED i
| Unix.WSIGNALED i
| Unix.WSTOPPED i -> i
let str_exists s p =
let rec f s p i =
if i = String.length s then false
else p s.[i] || f s p (i+1)
in
f s p 0
let rec iter_gen f g = match g() with
| None -> ()
| Some x -> f x; iter_gen f g
let finally_ f x ~h =
try
let y = f x in
ignore (h());
y
with e ->
ignore (h ());
raise e
(* print a string, but escaped if required *)
let escape_str s =
if
str_exists s
(function ' ' | '"' | '\'' | '\n' | '\t'-> true | _ -> false)
then (
let buf = Buffer.create (String.length s) in
Buffer.add_char buf '\'';
String.iter
(function
| '\'' -> Buffer.add_string buf "'\\''"
| c -> Buffer.add_char buf c
) s;
Buffer.add_char buf '\'';
Buffer.contents buf
) else s
let read_all ?(size=1024) ic =
let buf = ref (Bytes.create size) in
let len = ref 0 in
try
while true do
(* resize *)
if !len = Bytes.length !buf then (
buf := Bytes.extend !buf 0 !len;
);
assert (Bytes.length !buf > !len);
let n = input ic !buf !len (Bytes.length !buf - !len) in
len := !len + n;
if n = 0 then raise Exit; (* exhausted *)
done;
assert false (* never reached*)
with Exit ->
Bytes.sub_string !buf 0 !len
type call_result =
< stdout:string;
stderr:string;
status:Unix.process_status;
errcode:int; (** Extracted from status *)
>
let kbprintf' buf fmt k = Printf.kbprintf k buf fmt
let call_full_inner ?(bufsize=2048) ?(stdin=`Str "") ?(env=Unix.environment()) ~f cmd =
(* render the command *)
let buf = Buffer.create 256 in
kbprintf' buf cmd
(fun buf ->
let cmd = Buffer.contents buf in
let oc, ic, errc = Unix.open_process_full cmd env in
(* send stdin *)
begin match stdin with
| `Str s -> output_string ic s
| `Gen g -> iter_gen (output_string ic) g
end;
close_out ic;
(* read out and err *)
let out = read_all ~size:bufsize oc in
let err = read_all ~size:bufsize errc in
let status = Unix.close_process_full (oc, ic, errc) in
f (out,err,status)
)
let call_full ?bufsize ?stdin ?env cmd =
call_full_inner ?bufsize ?stdin ?env cmd
~f:(fun (out,err,status) ->
object
method stdout = out
method stderr = err
method status = status
method errcode = int_of_process_status status
end)
let call ?bufsize ?stdin ?env cmd =
call_full_inner ?bufsize ?stdin ?env cmd
~f:(fun (out,err,status) -> out, err, int_of_process_status status)
let call_stdout ?bufsize ?stdin ?env cmd =
call_full_inner ?bufsize ?stdin ?env cmd
~f:(fun (out,_err,_status) -> out)
type line = string
type async_call_result =
< stdout:line gen;
stderr:line gen;
stdin:line -> unit; (* send a line *)
close_in:unit; (* close stdin *)
close_err:unit;
close_out:unit;
close_all:unit; (* close all 3 channels *)
wait:Unix.process_status; (* block until the process ends *)
wait_errcode:int; (* block until the process ends, then extract errcode *)
>
let async_call ?(env=Unix.environment()) cmd =
(* render the command *)
let buf = Buffer.create 256 in
kbprintf' buf cmd
(fun buf ->
let cmd = Buffer.contents buf in
let oc, ic, errc = Unix.open_process_full cmd env in
object (self)
method stdout () =
try Some (input_line oc)
with End_of_file -> None
method stderr () =
try Some (input_line errc)
with End_of_file -> None
method stdin l = output_string ic l; output_char ic '\n'
method close_in = close_out ic
method close_out = close_in oc
method close_err = close_in errc
method close_all = close_out ic; close_in oc; close_in errc; ()
method wait = Unix.close_process_full (oc, ic, errc)
method wait_errcode = int_of_process_status self#wait
end
)
let stdout x = x#stdout
let stderr x = x#stderr
let status x = x#status
let errcode x = x#errcode
let with_in ?(mode=0o644) ?(flags=[]) file ~f =
let fd = Unix.openfile file (Unix.O_RDONLY::flags) mode in
let ic = Unix.in_channel_of_descr fd in
finally_ f ic
~h:(fun () -> Unix.close fd)
let with_out ?(mode=0o644) ?(flags=[Unix.O_CREAT; Unix.O_TRUNC]) file ~f =
let fd = Unix.openfile file (Unix.O_WRONLY::flags) mode in
let oc = Unix.out_channel_of_descr fd in
finally_ f oc
~h:(fun () -> flush oc; Unix.close fd)
let with_process_in cmd ~f =
let ic = Unix.open_process_in cmd in
finally_ f ic
~h:(fun () -> ignore (Unix.close_process_in ic))
let with_process_out cmd ~f =
let oc = Unix.open_process_out cmd in
finally_ f oc
~h:(fun () -> ignore (Unix.close_process_out oc))
type process_full = <
stdin: out_channel;
stdout: in_channel;
stderr: in_channel;
close: Unix.process_status;
>
let with_process_full ?env cmd ~f =
let env = match env with None -> Unix.environment () | Some e -> e in
let oc, ic, err = Unix.open_process_full cmd env in
let close = lazy (Unix.close_process_full (oc,ic,err)) in
let p = object
method stdin = ic
method stdout = oc
method stderr = err
method close = Lazy.force close
end in
finally_ f p ~h:(fun () -> p#close)
let with_connection addr ~f =
let ic, oc = Unix.open_connection addr in
finally_ (fun () -> f ic oc) ()
~h:(fun () -> Unix.shutdown_connection ic)
exception ExitServer
(* version of {!Unix.establish_server} that doesn't fork *)
let establish_server sockaddr ~f =
let sock =
Unix.socket (Unix.domain_of_sockaddr sockaddr) Unix.SOCK_STREAM 0 in
Unix.setsockopt sock Unix.SO_REUSEADDR true;
Unix.bind sock sockaddr;
Unix.listen sock 5;
let continue = ref true in
while !continue do
try
let s, _ = Unix.accept sock in
let ic = Unix.in_channel_of_descr s in
let oc = Unix.out_channel_of_descr s in
ignore (f ic oc)
with ExitServer ->
continue := false
done
(** {6 Locking} *)
let with_file_lock ~kind filename f =
let lock_file = Unix.openfile filename [Unix.O_CREAT; Unix.O_WRONLY] 0o644 in
let lock_action = match kind with
| `Read -> Unix.F_RLOCK
| `Write -> Unix.F_LOCK
in
Unix.lockf lock_file lock_action 0;
try
let x = f () in
Unix.lockf lock_file Unix.F_ULOCK 0;
Unix.close lock_file;
x
with e ->
Unix.lockf lock_file Unix.F_ULOCK 0;
Unix.close lock_file;
raise e
(*$R
let m = 200 in
let n = 50 in
let write_atom filename s =
with_file_lock ~kind:`Write filename
(fun () ->
CCIO.with_out ~flags:[Open_append; Open_creat]
filename (fun oc -> output_string oc s; flush oc))
in
let f filename =
for j=1 to m do
write_atom filename "foo\n"
done
in
CCIO.File.with_temp ~prefix:"containers_" ~suffix:".txt"
(fun filename ->
let a = Array.init n (fun _ -> Thread.create f filename) in
Array.iter Thread.join a;
let lines = CCIO.with_in filename CCIO.read_lines_l in
assert_equal ~printer:string_of_int (n * m) (List.length lines);
assert_bool "all valid" (List.for_all ((=) "foo") lines))
*)
module Infix = struct
let (?|) fmt = call_full fmt
let (?|&) fmt = async_call fmt
end
include Infix