dups.ml

   1 open Printf
   2
   3 module Array = ArrayLabels
   4 module List  = ListLabels
   5 module StrSet = Set.Make(String)
   6 module Unix  = UnixLabels
   7
   8 module Metrics : sig
   9   type t
  10
  11   val init
  12     : unit -> t
  13   val report
  14     : t
  15     -> wall_time_all:float
  16     -> wall_time_group_by_size:float
  17     -> wall_time_group_by_head:float
  18     -> wall_time_group_by_digest:float
  19     -> proc_time_all:float
  20     -> proc_time_group_by_size:float
  21     -> proc_time_group_by_head:float
  22     -> proc_time_group_by_digest:float
  23     -> unit
  24
  25   val file_considered
  26     : t -> size:int -> unit
  27   val file_ignored
  28     : t -> size:int -> unit
  29   val file_empty
  30     : t -> unit
  31   val file_sampled
  32     : t -> unit
  33   val chunk_read
  34     : t -> size:int -> unit
  35   val file_unique_size
  36     : t -> size:int -> unit
  37   val file_unique_sample
  38     : t -> size:int -> unit
  39   val file_hashed
  40     : t -> size:int -> unit
  41   val digest
  42     : t -> unit
  43   val redundant_data
  44     : t -> size:int -> unit
  45 end = struct
  46   type t =
  47     { considered_files    : int ref
  48     ; considered_bytes    : int ref
  49     ; empty               : int ref
  50     ; ignored_files       : int ref
  51     ; ignored_bytes       : int ref
  52     ; unique_size_files   : int ref
  53     ; unique_size_bytes   : int ref
  54     ; unique_sample_files : int ref
  55     ; unique_sample_bytes : int ref
  56     ; sampled_files       : int ref
  57     ; sampled_bytes       : int ref
  58     ; hashed_files        : int ref
  59     ; hashed_bytes        : int ref
  60     ; digests             : int ref
  61     ; redundant_data      : int ref
  62     }
  63
  64   let init () =
  65     { considered_files    = ref 0
  66     ; considered_bytes    = ref 0
  67     ; empty               = ref 0
  68     ; ignored_files       = ref 0
  69     ; ignored_bytes       = ref 0
  70     ; unique_size_files   = ref 0
  71     ; unique_size_bytes   = ref 0
  72     ; sampled_files       = ref 0
  73     ; sampled_bytes       = ref 0
  74     ; hashed_files        = ref 0
  75     ; hashed_bytes        = ref 0
  76     ; unique_sample_files = ref 0
  77     ; unique_sample_bytes = ref 0
  78     ; digests             = ref 0
  79     ; redundant_data      = ref 0
  80     }
  81
  82   let add sum addend =
  83     sum := !sum + addend
  84
  85   let file_considered t ~size =
  86     incr t.considered_files;
  87     add t.considered_bytes size
  88
  89   let file_ignored {ignored_files; ignored_bytes; _} ~size =
  90     incr ignored_files;
  91     add ignored_bytes size
  92
  93   let file_empty t =
  94     incr t.empty
  95
  96   let chunk_read t ~size =
  97     add t.sampled_bytes size
  98
  99   let file_sampled t =
 100     incr t.sampled_files
 101
 102   let file_unique_size t ~size =
 103     incr t.unique_size_files;
 104     add t.unique_size_bytes size
 105
 106   let file_unique_sample t ~size =
 107     incr t.unique_sample_files;
 108     add t.unique_sample_bytes size
 109
 110   let file_hashed t ~size =
 111     incr t.hashed_files;
 112     add t.hashed_bytes size
 113
 114   let digest t =
 115     incr t.digests
 116
 117   let redundant_data t ~size =
 118     add t.redundant_data size
 119
 120   let report
 121     t
 122     ~wall_time_all
 123     ~wall_time_group_by_size
 124     ~wall_time_group_by_head
 125     ~wall_time_group_by_digest
 126     ~proc_time_all
 127     ~proc_time_group_by_size
 128     ~proc_time_group_by_head
 129     ~proc_time_group_by_digest
 130   =
 131     let b_to_mb b = (float_of_int b) /. 1024. /. 1024. in
 132     let b_to_gb b = (b_to_mb b) /. 1024. in
 133     eprintf "Total time                   : %.2f wall sec  %.2f proc sec\n%!"
 134       wall_time_all
 135       proc_time_all;
 136     eprintf "Considered                   : %8d files  %6.2f Gb\n%!"
 137       !(t.considered_files)
 138       (b_to_gb !(t.considered_bytes));
 139     eprintf "Sampled                      : %8d files  %6.2f Gb\n%!"
 140       !(t.sampled_files)
 141       (b_to_gb !(t.sampled_bytes));
 142     eprintf "Hashed                       : %8d files  %6.2f Gb  %6.2f wall sec  %6.2f proc sec\n%!"
 143       !(t.hashed_files)
 144       (b_to_gb !(t.hashed_bytes))
 145       wall_time_group_by_digest
 146       proc_time_group_by_digest;
 147     eprintf "Digests                      : %8d\n%!"
 148       !(t.digests);
 149     eprintf "Duplicates (Hashed - Digests): %8d files  %6.2f Gb\n%!"
 150       (!(t.hashed_files) - !(t.digests))
 151       (b_to_gb !(t.redundant_data));
 152     eprintf "Skipped due to 0      size   : %8d files\n%!" !(t.empty);
 153     eprintf "Skipped due to unique size   : %8d files  %6.2f Gb  %6.2f wall sec  %6.2f proc sec\n%!"
 154       !(t.unique_size_files)
 155       (b_to_gb !(t.unique_size_bytes))
 156       wall_time_group_by_size
 157       proc_time_group_by_size;
 158     eprintf "Skipped due to unique sample : %8d files  %6.2f Gb  %6.2f wall sec  %6.2f proc sec\n%!"
 159       !(t.unique_sample_files)
 160       (b_to_gb !(t.unique_sample_bytes))
 161       wall_time_group_by_head
 162       proc_time_group_by_head;
 163     eprintf "Ignored due to regex match   : %8d files  %6.2f Gb\n%!"
 164       !(t.ignored_files)
 165       (b_to_gb !(t.ignored_bytes))
 166 end
 167
 168 module M = Metrics
 169
 170 module Stream : sig
 171   type 'a t
 172
 173   val create : (unit -> 'a option) -> 'a t
 174
 175   val of_queue : 'a Queue.t -> 'a t
 176
 177   val iter : 'a t -> f:('a -> unit) -> unit
 178
 179   val bag_map : 'a t -> njobs:int -> f:('a -> 'b) -> ('a * 'b) t
 180   (** Parallel map with arbitrarily-reordered elements. *)
 181
 182   val map : 'a t -> f:('a -> 'b) -> 'b t
 183
 184   val filter : 'a t -> f:('a -> bool) -> 'a t
 185
 186   val concat : ('a t) list -> 'a t
 187
 188   val group_by : 'a t -> f:('a -> 'b) -> ('b * int * 'a list) t
 189 end = struct
 190   module S = Stream
 191
 192   type 'a t =
 193     {mutable streams : ('a S.t) list}
 194
 195   type ('input, 'output) msg_from_vassal =
 196     | Ready of int
 197     | Result of (int * ('input * 'output))
 198     | Exiting of int
 199
 200   type 'input msg_from_lord =
 201     | Job of 'input option
 202
 203   let create f =
 204     {streams = [S.from (fun _ -> f ())]}
 205
 206   let of_queue q =
 207     create (fun () ->
 208       match Queue.take q with
 209       | exception Queue.Empty ->
 210           None
 211       | x ->
 212           Some x
 213     )
 214
 215   let rec next t =
 216     match t.streams with
 217     | [] ->
 218         None
 219     | s :: streams ->
 220         (match S.next s with
 221         | exception Stream.Failure ->
 222             t.streams <- streams;
 223             next t
 224         | x ->
 225             Some x
 226         )
 227
 228   let map t ~f =
 229     create (fun () ->
 230       match next t with
 231       | None   -> None
 232       | Some x -> Some (f x)
 233     )
 234
 235   let filter t ~f =
 236     let rec filter () =
 237       match next t with
 238       | None ->
 239           None
 240       | Some x when f x ->
 241           Some x
 242       | Some _ ->
 243           filter ()
 244     in
 245     create filter
 246
 247   let iter t ~f =
 248     List.iter t.streams ~f:(S.iter f)
 249
 250   let concat ts =
 251     {streams = List.concat (List.map ts ~f:(fun {streams} -> streams))}
 252
 253   let group_by t ~f =
 254     let groups_tbl = Hashtbl.create 1_000_000 in
 255     let group_update x =
 256       let group = f x in
 257       let members =
 258         match Hashtbl.find_opt groups_tbl group with
 259         | None ->
 260             (1, [x])
 261         | Some (n, xs) ->
 262             (succ n, x :: xs)
 263       in
 264       Hashtbl.replace groups_tbl group members
 265     in
 266     iter t ~f:group_update;
 267     let groups = Queue.create () in
 268     Hashtbl.iter
 269       (fun name (length, members) -> Queue.add (name, length, members) groups)
 270       groups_tbl;
 271     of_queue groups
 272
 273   module Ipc : sig
 274     val send : out_channel -> 'a -> unit
 275     val recv : in_channel -> 'a
 276   end = struct
 277     let send oc msg =
 278       Marshal.to_channel oc msg [];
 279       flush oc
 280
 281     let recv ic =
 282       Marshal.from_channel ic
 283   end
 284
 285   let lord t ~njobs ~vassals ~ic ~ocs =
 286     eprintf "[debug] [lord] started\n%!";
 287     let active_vassals = ref njobs in
 288     let results = Queue.create () in
 289     let rec dispatch () =
 290       match Ipc.recv ic with
 291       | ((Exiting i) : ('input, 'output) msg_from_vassal) ->
 292           close_out ocs.(i);
 293           decr active_vassals;
 294           if !active_vassals = 0 then
 295             ()
 296           else
 297             dispatch ()
 298       | ((Ready i) : ('input, 'output) msg_from_vassal) ->
 299           Ipc.send ocs.(i) (Job (next t));
 300           dispatch ()
 301       | ((Result (i, result)) : ('input, 'output) msg_from_vassal) ->
 302           Queue.add result results;
 303           Ipc.send ocs.(i) (Job (next t));
 304           dispatch ()
 305     in
 306     let rec wait = function
 307       | [] -> ()
 308       | vassals ->
 309           let pid, _process_status = Unix.wait () in
 310           (* TODO: handle process_status *)
 311           wait (List.filter vassals ~f:(fun p -> p <> pid))
 312     in
 313     dispatch ();
 314     close_in ic;
 315     wait vassals;
 316     of_queue results
 317
 318   let vassal i ~f ~vassal_pipe_r ~lord_pipe_w =
 319     eprintf "[debug] [vassal %d] started\n%!" i;
 320     let ic = Unix.in_channel_of_descr vassal_pipe_r in
 321     let oc = Unix.out_channel_of_descr lord_pipe_w in
 322     let rec work msg =
 323       Ipc.send oc msg;
 324       match Ipc.recv ic with
 325       | (Job (Some x) : 'input msg_from_lord) ->
 326           work (Result (i, (x, f x)))
 327       | (Job None : 'input msg_from_lord) ->
 328           Ipc.send oc (Exiting i)
 329     in
 330     work (Ready i);
 331     close_in ic;
 332     close_out oc;
 333     exit 0
 334
 335   let bag_map t ~njobs ~f =
 336     let lord_pipe_r, lord_pipe_w = Unix.pipe () in
 337     let vassal_pipes   = Array.init njobs ~f:(fun _ -> Unix.pipe ()) in
 338     let vassal_pipes_r = Array.map vassal_pipes ~f:(fun (r, _) -> r) in
 339     let vassal_pipes_w = Array.map vassal_pipes ~f:(fun (_, w) -> w) in
 340     let vassals = ref [] in
 341     for i=0 to (njobs - 1) do
 342       begin match Unix.fork () with
 343       | 0 ->
 344           Unix.close lord_pipe_r;
 345           vassal i ~f ~lord_pipe_w ~vassal_pipe_r:vassal_pipes_r.(i)
 346       | pid ->
 347           vassals := pid :: !vassals
 348       end
 349     done;
 350     Unix.close lord_pipe_w;
 351     lord
 352       t
 353       ~njobs
 354       ~vassals:!vassals
 355       ~ic:(Unix.in_channel_of_descr lord_pipe_r)
 356       ~ocs:(Array.map vassal_pipes_w ~f:Unix.out_channel_of_descr)
 357 end
 358
 359 module In_channel : sig
 360   val lines : in_channel -> string Stream.t
 361 end = struct
 362   let lines ic =
 363     Stream.create (fun () ->
 364       match input_line ic with
 365       | exception End_of_file ->
 366           None
 367       | line ->
 368           Some line
 369     )
 370 end
 371
 372 module File : sig
 373   type t =
 374     { path : string
 375     ; size : int
 376     }
 377
 378   val find : string -> t Stream.t
 379   (** Find all files in the directory tree, starting from the given root path *)
 380
 381   val lookup : string Stream.t -> t Stream.t
 382   (** Lookup file info for given paths *)
 383
 384   val filter_out_unique_sizes : t Stream.t -> metrics:M.t -> t Stream.t
 385   val filter_out_unique_heads : t Stream.t -> len:int -> metrics:M.t -> t Stream.t
 386 end = struct
 387   type t =
 388     { path : string
 389     ; size : int
 390     }
 391
 392   let lookup paths =
 393     Stream.map paths ~f:(fun path ->
 394       let {Unix.st_size = size; _} = Unix.lstat path in
 395       {path; size}
 396     )
 397
 398   let find root =
 399     let dirs  = Queue.create () in
 400     let files = Queue.create () in
 401     let explore parent =
 402       Array.iter (Sys.readdir parent) ~f:(fun child ->
 403         let path = Filename.concat parent child in
 404         let {Unix.st_kind = file_kind; st_size; _} = Unix.lstat path in
 405         match file_kind with
 406         | Unix.S_REG ->
 407             let file = {path; size = st_size} in
 408             Queue.add file files
 409         | Unix.S_DIR ->
 410             Queue.add path dirs
 411         | Unix.S_CHR
 412         | Unix.S_BLK
 413         | Unix.S_LNK
 414         | Unix.S_FIFO
 415         | Unix.S_SOCK ->
 416             ()
 417       )
 418     in
 419     explore root;
 420     let rec next () =
 421       match Queue.is_empty files, Queue.is_empty dirs with
 422       | false, _     -> Some (Queue.take files)
 423       | true , true  -> None
 424       | true , false ->
 425           explore (Queue.take dirs);
 426           next ()
 427     in
 428     Stream.create next
 429
 430   let filter_out_singletons files ~group ~handle_singleton =
 431     let q = Queue.create () in
 432     Stream.iter (Stream.group_by files ~f:group) ~f:(fun group ->
 433       let (_, n, members) = group in
 434       if n > 1 then
 435         List.iter members ~f:(fun m -> Queue.add m q)
 436       else
 437         handle_singleton group
 438     );
 439     Stream.of_queue q
 440
 441   let filter_out_unique_sizes files ~metrics =
 442     filter_out_singletons
 443       files
 444       ~group:(fun {size; _} -> size)
 445       ~handle_singleton:(fun (size, _, _) -> M.file_unique_size metrics ~size)
 446
 447   let head path ~len ~metrics =
 448     let buf = Bytes.make len ' ' in
 449     let ic = open_in_bin path in
 450     let rec read pos len =
 451       assert (len >= 0);
 452       if len = 0 then
 453         ()
 454       else begin
 455         let chunk_size = input ic buf pos len in
 456         M.chunk_read metrics ~size:chunk_size;
 457         if chunk_size = 0 then (* EOF *)
 458           ()
 459         else
 460           read (pos + chunk_size) (len - chunk_size)
 461       end
 462     in
 463     read 0 len;
 464     close_in ic;
 465     Bytes.to_string buf
 466
 467   let filter_out_unique_heads files ~len ~metrics =
 468     filter_out_singletons
 469       files
 470       ~group:(fun {path; _} ->
 471         M.file_sampled metrics;
 472         head path ~len ~metrics
 473       )
 474       ~handle_singleton:(fun (_, _, files) ->
 475         let {size; _} = List.hd files in  (* Guaranteed non-empty *)
 476         M.file_unique_sample metrics ~size
 477       )
 478 end
 479
 480 type input =
 481   | Stdin
 482   | Directories of string list
 483
 484 type output =
 485   | Stdout
 486   | Directory of string
 487
 488 type opt =
 489   { input  : input
 490   ; output : output
 491   ; ignore : string -> bool
 492   ; sample : int
 493   ; njobs  : int
 494   }
 495
 496 let make_input_stream input ignore ~metrics =
 497   let input =
 498     match input with
 499     | Stdin ->
 500         File.lookup (In_channel.lines stdin)
 501     | Directories paths ->
 502         let paths = StrSet.elements (StrSet.of_list paths) in
 503         Stream.concat (List.map paths ~f:File.find)
 504   in
 505   Stream.filter input ~f:(fun {File.path; size} ->
 506     M.file_considered metrics ~size;
 507     let empty = size = 0 in
 508     let ignored = ignore path in
 509     if empty then M.file_empty metrics;
 510     if ignored then M.file_ignored metrics ~size;
 511     (not empty) && (not ignored)
 512   )
 513
 514 let make_output_fun = function
 515   | Stdout ->
 516       fun digest n_files files ->
 517         printf "%s %d\n%!" (Digest.to_hex digest) n_files;
 518         List.iter files ~f:(fun {File.path; _} ->
 519           printf "    %S\n%!" path
 520         )
 521   | Directory dir ->
 522       fun digest _ files ->
 523         let digest = Digest.to_hex digest in
 524         let dir = Filename.concat dir (String.sub digest 0 2) in
 525         Unix.mkdir dir ~perm:0o700;
 526         let oc = open_out (Filename.concat dir digest) in
 527         List.iter files ~f:(fun {File.path; _} ->
 528           output_string oc (sprintf "%S\n%!" path)
 529         );
 530         close_out oc
 531
 532 let time_wall () =
 533   Unix.gettimeofday ()
 534
 535 let time_proc () =
 536   Sys.time ()
 537
 538 let main {input; output; ignore; sample = sample_len; njobs} =
 539   let wt0_all = time_wall () in
 540   let pt0_all = time_proc () in
 541   let metrics = M.init () in
 542   let output = make_output_fun  output in
 543   let input  = make_input_stream input ignore ~metrics in
 544   (* TODO: Make a nice(r) abstraction to re-assemble pieces in the pipeline:
 545    *
 546    * from input           to files_by_size
 547    * from files_by_size   to files_by_sample
 548    * from files_by_sample to files_by_digest
 549    * from files_by_digest to output
 550    *
 551    * input |> files_by_size |> files_by_sample |> files_by_digest |> output
 552    *)
 553
 554   let files = input in
 555
 556   let wt0_group_by_size = time_wall () in
 557   let pt0_group_by_size = time_proc () in
 558   eprintf "[debug] filtering-out files with unique size\n%!";
 559   let files = File.filter_out_unique_sizes files ~metrics in
 560   let pt1_group_by_size = time_proc () in
 561   let wt1_group_by_size = time_wall () in
 562
 563   let wt0_group_by_sample = wt1_group_by_size in
 564   let pt0_group_by_sample = pt1_group_by_size in
 565   eprintf "[debug] filtering-out files with unique heads\n%!";
 566   let files = File.filter_out_unique_heads files ~len:sample_len ~metrics in
 567   let pt1_group_by_sample = time_proc () in
 568   let wt1_group_by_sample = time_wall () in
 569
 570   let wt0_group_by_digest = wt1_group_by_sample in
 571   let pt0_group_by_digest = pt1_group_by_sample in
 572   eprintf "[debug] hashing\n%!";
 573   let groups =
 574     if njobs > 1 then
 575       let digests =
 576         Stream.bag_map files ~njobs ~f:(fun {File.path; _} -> Digest.file path)
 577       in
 578       Stream.map (Stream.group_by digests ~f:(fun (_, d) -> d)) ~f:(
 579         fun (digest, n, file_digest_pairs) ->
 580           let files =
 581             List.map file_digest_pairs ~f:(fun (file, _) ->
 582               M.file_hashed metrics ~size:file.File.size;
 583               file
 584             )
 585           in
 586           (digest, n, files)
 587       )
 588     else
 589       Stream.group_by files ~f:(fun {File.path; size} ->
 590         M.file_hashed metrics ~size;
 591         Digest.file path
 592       )
 593   in
 594   let pt1_group_by_digest = time_proc () in
 595   let wt1_group_by_digest = time_wall () in
 596
 597   eprintf "[debug] reporting\n%!";
 598   Stream.iter groups ~f:(fun (d, n, files) ->
 599     M.digest metrics;
 600     if n > 1 then
 601       M.redundant_data metrics ~size:(n * (List.hd files).File.size);
 602       output d n files
 603   );
 604
 605   let pt1_all = time_proc () in
 606   let wt1_all = time_wall () in
 607
 608   M.report metrics
 609     ~wall_time_all:            (wt1_all             -. wt0_all)
 610     ~wall_time_group_by_size:  (wt1_group_by_size   -. wt0_group_by_size)
 611     ~wall_time_group_by_head:  (wt1_group_by_sample -. wt0_group_by_sample)
 612     ~wall_time_group_by_digest:(wt1_group_by_digest -. wt0_group_by_digest)
 613     ~proc_time_all:            (pt1_all             -. pt0_all)
 614     ~proc_time_group_by_size:  (pt1_group_by_size   -. pt0_group_by_size)
 615     ~proc_time_group_by_head:  (pt1_group_by_sample -. pt0_group_by_sample)
 616     ~proc_time_group_by_digest:(pt1_group_by_digest -. pt0_group_by_digest)
 617
 618 let get_opt () : opt =
 619   let assert_ test x msg =
 620     if not (test x) then begin
 621       eprintf "%s\n%!" msg;
 622       exit 1
 623     end
 624   in
 625   let assert_file_exists path =
 626     assert_ Sys.file_exists path (sprintf "File does not exist: %S" path)
 627   in
 628   let assert_file_is_dir path =
 629     assert_ Sys.is_directory path (sprintf "File is not a directory: %S" path)
 630   in
 631   let input  = ref Stdin in
 632   let output = ref Stdout in
 633   let ignore = ref (fun _ -> false) in
 634   let sample = ref 512 in
 635   let njobs  = ref 8 in
 636   let spec =
 637     [ ( "-out"
 638       , Arg.String (fun path ->
 639           assert_file_exists path;
 640           assert_file_is_dir path;
 641           output := Directory path
 642         )
 643       , " Output to this directory instead of stdout."
 644       )
 645     ; ( "-ignore"
 646       , Arg.String (fun regexp ->
 647           let regexp = Str.regexp regexp in
 648           ignore := fun string -> Str.string_match regexp string 0)
 649       , " Ignore file paths which match this regexp pattern (see Str module)."
 650       )
 651     ; ( "-sample"
 652       , Arg.Set_int sample
 653       , (sprintf " Byte size of file samples to use. Default: %d" !sample)
 654       )
 655     ; ( "-j"
 656       , Arg.Set_int njobs
 657       , (sprintf " Number of parallel jobs. Default: %d" !njobs)
 658       )
 659     ]
 660   in
 661   Arg.parse
 662     (Arg.align spec)
 663     (fun path ->
 664       assert_file_exists path;
 665       assert_file_is_dir path;
 666       match !input with
 667       | Stdin ->
 668           input := Directories [path]
 669       | Directories paths ->
 670           input := Directories (path :: paths)
 671     )
 672     "";
 673   assert_
 674     (fun x -> x > 0)
 675     !sample
 676     (sprintf "Sample size cannot be negative: %d" !sample);
 677   { input  = !input
 678   ; output = !output
 679   ; ignore = !ignore
 680   ; sample = !sample
 681   ; njobs  = !njobs
 682   }
 683
 684 let () =
 685   main (get_opt ())