Source

spotlib / lib / xunix.ml

Full commit
  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
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 is_reg path = kind path = Some S_REG

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


module Find = struct

  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_reg : bool
    method kind : (Unix.file_kind, exn) Result.t
    method path : string
    method stat : (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 kind : (_,_) Result.t = match self#stat with
        | `Error _exn -> `Error _exn
        | `Ok stat -> `Ok stat.st_kind
      method is_dir = self#kind = `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 ~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 pth#is_dir 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 = 
        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)