faster gzip_io, using the new extlib IO.as_buffered interface

Makefile

@@ -29,6 +29,7 @@ test:
 
 benchs:
 		RUN_BENCHS=1 dune exec --profile bench $(DUNEFLAGS) bench/bench_sharded_hash_trie.exe
+		dune exec --release $(DUNEFLAGS) bench/bench_gzip_io.exe
 
 doc:
 		dune build $(DUNEFLAGS) @doc

bench/bench_gzip_io.ml

@@ -0,0 +1,122 @@
+(* Gzip_io.input: IO.read_line throughput, one byte at a time ("by_char") vs
+   directly on the channel ("current").
+
+   dune exec --release bench/bench_gzip_io.exe *)
+
+open Devkit
+
+(* hide any [as_buffered] so that [IO.read_line] goes through [read] one byte at a time;
+   also returns the number of bytes read so far *)
+let unbuffered_counting inp =
+  let n = ref 0 in
+  ( IO.create_in
+      ~read:(fun () ->
+        let c = IO.read inp in
+        incr n;
+        c)
+      ~input:(fun buf pos len ->
+        let k = IO.input inp buf pos len in
+        n := !n + k;
+        k)
+      ~close:(fun () -> IO.close_in inp),
+    fun () -> !n )
+
+(* [current] reads through the buffer directly, so there is no byte count to get *)
+let impls =
+  [ ( "by_char",
+      fun z ->
+        let inp, count = unbuffered_counting (Gzip_io.input (IO.input_string z)) in
+        inp, Some count );
+    "current", fun z -> Gzip_io.input (IO.input_string z), None ]
+
+let lines_of ~count len_of =
+  let b = Buffer.create (1 lsl 20) in
+  for i = 0 to count - 1 do
+    Buffer.add_string b (String.init (len_of i) (fun j -> Char.chr (Char.code 'a' + ((i + j) mod 26))));
+    Buffer.add_char b '\n'
+  done;
+  Buffer.contents b
+
+let kib = 1024
+
+(* each dataset is generated on demand to keep memory low; sized so that [current] runs for ~100ms+ *)
+let datasets =
+  [ "small lines (~0-100B)", (fun () -> lines_of ~count:3_000_000 (fun i -> i * 7 mod 100));
+    "medium lines (~1-4KiB)", (fun () -> lines_of ~count:60_000 (fun i -> kib + (i * 131 mod (3 * kib))));
+    "big lines (~100-300KiB)", (fun () -> lines_of ~count:1200 (fun i -> (100 * kib) + (i * 37 * kib mod (200 * kib))));
+    "mixed", (fun () -> lines_of ~count:600_000 (fun i -> if i mod 1000 = 0 then 200 * kib else i mod 80)) ]
+
+let time f =
+  let t0 = Unix.gettimeofday () in
+  let r = f () in
+  r, Unix.gettimeofday () -. t0
+
+let count_lines inp =
+  let n = ref 0 in
+  (try
+     while true do
+       ignore (IO.read_line inp : string);
+       incr n
+     done
+   with IO.No_more_input -> ());
+  IO.close_in inp;
+  !n
+
+(* just decompress, no parsing: one [IO.read] per byte *)
+let through_chars inp =
+  let n = ref 0 in
+  (try
+     while true do
+       ignore (IO.read inp : char);
+       incr n
+     done
+   with IO.No_more_input -> ());
+  IO.close_in inp;
+  !n
+
+(* just decompress, no parsing: [IO.input] into a 64 KiB buffer *)
+let through_chunks inp =
+  let buf = Bytes.create (64 * kib) in
+  let n = ref 0 in
+  (try
+     while true do
+       n := !n + IO.input inp buf 0 (Bytes.length buf)
+     done
+   with IO.No_more_input -> ());
+  IO.close_in inp;
+  !n
+
+(* what is done with the input, per implementation; results (line or byte counts) must agree
+   across implementations *)
+let modes =
+  [ "lines", [ "by_char", count_lines; "current", count_lines ];
+    "through", [ "by_char", through_chars; "current", through_chunks ] ]
+
+let () =
+  ignore (Unix.alarm 900 : int);
+  let runs = 3 in
+  List.iter (fun (dname, gen) ->
+    let raw = gen () in
+    let z = Gzip_io.string raw in
+    Printf.printf "%s: %s raw, %s gzip\n%!" dname
+      (Action.bytes_string (String.length raw)) (Action.bytes_string (String.length z));
+    List.iter (fun (mname, consumers) ->
+      let expected = ref None in
+      List.iter (fun (iname, consume) ->
+        let mk = List.assoc iname impls in
+        let best = ref infinity in
+        for _ = 1 to runs do
+          let (result, count), dt = time (fun () ->
+            let inp, count = mk z in
+            consume inp, count)
+          in
+          (match !expected with
+           | None -> expected := Some result
+           | Some e -> if e <> result then failwith (Printf.sprintf "%s/%s/%s: result mismatch" dname mname iname));
+          Stdlib.Option.iter (fun count -> if count () <> String.length raw then failwith (Printf.sprintf "%s/%s/%s: byte count mismatch" dname mname iname)) count;
+          best := min !best dt
+        done;
+        Printf.printf "  %-8s %-8s %8.1f MB/s  (%.3fs)\n%!" mname iname (float (String.length raw) /. 1e6 /. !best) !best)
+        consumers)
+      modes)
+    datasets

bench/dune

@@ -1,6 +1,14 @@
 ; optional: needs saturn and benchmark. Build with make benchs.
+
 (executable
  (name bench_sharded_hash_trie)
+ (modules bench_sharded_hash_trie)
  (optional)
- (enabled_if (= %{profile} bench))
+ (enabled_if
+  (= %{profile} bench))
  (libraries devkit saturn benchmark unix))
+
+(executable
+ (name bench_gzip_io)
+ (modules bench_gzip_io)
+ (libraries devkit extlib unix))

devkit.opam

@@ -13,7 +13,7 @@ build: [
 depends: [
   "ocaml" {>= "5.0"}
   "dune" {>= "2.3"}
-  ("extlib" {>= "1.7.1"} | "extlib-compat" {>= "1.7.1"})
+  ("extlib" {>= "1.9"} | "extlib-compat" {>= "1.7.1"})
   "ounit2"
   "qcheck-core" {with-test}
   "benchmark" {with-dev-setup}

gzip_io.ml

@@ -1,11 +1,38 @@
 (** gzip IO *)
 
+let buf_size = 65536
+
 let input io =
-  let iz = Gzip_stream.open_in io in
-  IO.create_in
-    ~read:(fun () -> Gzip_stream.input_char iz)
-    ~input:(Gzip_stream.input iz)
-    ~close:(fun () -> Gzip_stream.close_in iz)
+  let stream = Gzip_stream.open_in io in
+  (* Gzip_stream.input returns 0 at end of stream, which [Buf_in.create] turns into [No_more_input] *)
+  let buf_in = IO.Buf_in.create (Bytes.create buf_size) ~fill:(fun buf -> Gzip_stream.input stream buf 0 (Bytes.length buf)) in
+  let read () =
+    IO.Buf_in.refill_if_empty buf_in;
+    assert (buf_in.len > 0);
+    let c = Bytes.unsafe_get buf_in.buf buf_in.pos in
+    IO.Buf_in.consume buf_in 1;
+    c
+  in
+  let input buf pos len =
+    if buf_in.len = 0 then (
+      (* passthrough *)
+      let n = Gzip_stream.input stream buf pos len in
+      if n = 0 && len > 0 then raise IO.No_more_input;
+      n
+    )
+    else (
+      let n = min len buf_in.len in
+      Bytes.blit buf_in.buf buf_in.pos buf pos n;
+      IO.Buf_in.consume buf_in n;
+      n
+    )
+  in
+  let buffered = Some buf_in in
+  IO.create_in_full ()
+    ~read
+    ~input
+    ~close:(fun () -> Gzip_stream.close_in stream)
+    ~as_buffered:(fun () -> buffered)
 
 let output io =
   let oz = Gzip_stream.open_out io in

gzip_io.mli

@@ -0,0 +1,16 @@
+(** gzip IO *)
+
+(** [input io] decompresses the gzip data read from [io] (a single gzip member) *)
+val input : IO.input -> IO.input
+
+(** [output io] compresses what is written to it as gzip and writes that to [io] *)
+val output : 'a IO.output -> 'a IO.output
+
+val input_ch : Stdlib.in_channel -> IO.input
+val output_ch : Stdlib.out_channel -> unit IO.output
+
+(** [string s] is [s] compressed as gzip *)
+val string : string -> string
+
+(** [to_string s] decompresses gzip data [s] *)
+val to_string : string -> string

tests/dune

@@ -1,3 +1,9 @@
 (tests
  (names sharded_hash_trie_test cache_count_test logger_domains_test)
+ (modules sharded_hash_trie_test cache_count_test logger_domains_test)
  (libraries devkit qcheck-core qcheck-core.runner unix))
+
+(test
+ (name gzip_io_test)
+ (modules gzip_io_test)
+ (libraries devkit extlib unix))

tests/gzip_io_test.ml

@@ -0,0 +1,124 @@
+(* Gzip_io.input: reading lines back from compressed data *)
+open Devkit
+
+(* hide any [as_buffered] so that [IO.read_line] goes through [read] one byte at a time *)
+let unbuffered inp =
+  IO.create_in
+    ~read:(fun () -> IO.read inp)
+    ~input:(fun buf pos len -> IO.input inp buf pos len)
+    ~close:(fun () -> IO.close_in inp)
+
+let by_char z = unbuffered (Gzip_io.input (IO.input_string z))
+let current z = Gzip_io.input (IO.input_string z)
+let reference raw = unbuffered (IO.input_string raw)
+
+let read_lines inp =
+  let rec loop acc =
+    match IO.read_line inp with
+    | l -> loop (l :: acc)
+    | exception IO.No_more_input -> List.rev acc
+  in
+  let l = loop [] in
+  IO.close_in inp;
+  l
+
+let failures = ref 0
+
+let check name raw =
+  let z = Gzip_io.string raw in
+  let expected = read_lines (reference raw) in
+  let ok impl f =
+    let got = f z in
+    if got <> expected then begin
+      incr failures;
+      Printf.printf "FAIL %s [%s]: %d lines expected, got %d\n%!" name impl (List.length expected) (List.length got)
+    end
+  in
+  ok "by_char" (fun z -> read_lines (by_char z));
+  ok "current" (fun z -> read_lines (current z));
+  Printf.printf "ok   %s (%d bytes, %d lines)\n%!" name (String.length raw) (List.length expected)
+
+let kib = 1024
+let block = 64 * kib
+
+let random_text ~seed ~alphabet n =
+  let st = Random.State.make [| seed |] in
+  String.init n (fun _ -> alphabet.[Random.State.int st (String.length alphabet)])
+
+(* lines of [len_of i] 'x' (with a few distinct chars) joined by '\n' *)
+let lines_of ~count len_of =
+  let b = Buffer.create (1 lsl 20) in
+  for i = 0 to count - 1 do
+    let n = len_of i in
+    Buffer.add_string b (String.init n (fun j -> Char.chr (Char.code 'a' + ((i + j) mod 26))));
+    Buffer.add_char b '\n'
+  done;
+  Buffer.contents b
+
+
+(* random mix of IO ops on the gzip input and on a reference over the raw data, results must match *)
+let mixed seed =
+  let raw = random_text ~seed ~alphabet:"ab\n\r xyz{}\":0123" (3 * block) in
+  let z = Gzip_io.string raw in
+  let a = current z and b = reference raw in
+  let st = Random.State.make [| seed; 42 |] in
+  let size () =
+    if Random.State.int st 10 = 0 then Random.State.int st (128 * kib + 1) else Random.State.int st 513
+  in
+  let res f = match f () with v -> Ok v | exception IO.No_more_input -> Error () in
+  let fail op = incr failures; Printf.printf "FAIL mixed seed %d: %s differs\n%!" seed op in
+  (try
+     for _ = 1 to 300 do
+       match Random.State.int st 4 with
+       | 0 -> if res (fun () -> IO.read a) <> res (fun () -> IO.read b) then (fail "read"; raise Exit)
+       | 1 -> if res (fun () -> IO.read_line a) <> res (fun () -> IO.read_line b) then (fail "read_line"; raise Exit)
+       | 2 ->
+         let n = size () in
+         if res (fun () -> IO.nread_string a n) <> res (fun () -> IO.nread_string b n) then (fail "nread_string"; raise Exit)
+       | _ ->
+         let n = size () in
+         let buf = Bytes.create n in
+         (match res (fun () -> IO.input a buf 0 n) with
+          | Error () -> if res (fun () -> IO.input b (Bytes.create 1) 0 1) <> Error () then (fail "input eof"; raise Exit)
+          | Ok 0 when n > 0 -> fail "input returned 0 instead of raising"; raise Exit
+          | Ok k ->
+            (* [input] may return short, read exactly as much from the reference *)
+            (match res (fun () -> IO.really_nread b k) with
+             | Ok exp when Bytes.sub_string buf 0 k = Bytes.to_string exp -> ()
+             | _ -> fail "input"; raise Exit))
+     done
+   with Exit -> ());
+  IO.close_in a;
+  IO.close_in b;
+  Printf.printf "ok   mixed ops seed %d\n%!" seed
+
+let closed () =
+  let inp = current (Gzip_io.string "abc\ndef\n") in
+  ignore (IO.read_line inp : string);
+  IO.close_in inp;
+  match IO.read_line inp with
+  | _ -> incr failures; print_endline "FAIL read_line after close_in returned"
+  | exception IO.Input_closed -> print_endline "ok   read_line after close_in"
+  | exception e -> incr failures; Printf.printf "FAIL read_line after close_in raised %s\n%!" (Printexc.to_string e)
+
+let () =
+  ignore (Unix.alarm 300 : int);
+  check "empty" "";
+  check "no trailing newline" "hello";
+  check "single newline" "\n";
+  check "empty lines" "\n\n\na\n\n";
+  check "crlf" "a\r\nb\r\n\r\nc";
+  check "cr runs" "a\r\r\r\nb\r\r\rc\r\r\r";
+  check "line > 64KiB" (String.make (block + 100) 'x' ^ "\nend\n");
+  check "line ends at 64KiB" (String.make (block - 1) 'x' ^ "\nend\n");
+  check "newline at 64KiB+1" (String.make block 'x' ^ "\nend\n");
+  check "cr run across 64KiB" (String.make (block - 10) 'x' ^ String.make 20 '\r' ^ "\ny\n");
+  check "small lines (~8 MiB)" (lines_of ~count:600_000 (fun i -> i mod 24));
+  check "big lines (~8 MiB)" (lines_of ~count:40 (fun i -> 100 * kib + (i * 37 * kib mod (200 * kib))));
+  check "mixed small and big lines" (lines_of ~count:20_000 (fun i -> if i mod 500 = 0 then 3 * block + i else i mod 100));
+  for seed = 0 to 2 do
+    check (Printf.sprintf "random text %d" seed) (random_text ~seed ~alphabet:"ab\n\r xyz{}\":0123" (5 * block))
+  done;
+  for seed = 0 to 19 do mixed seed done;
+  closed ();
+  if !failures > 0 then (Printf.printf "%d failures\n%!" !failures; exit 1)