Add support for compressed output files, whether gzip (supported by Perfetto) or zstd (not yet supported)

This builds on previous work for compressed file destinations

Closes #154

Signed-off-by: Brian Nigito <[email protected]>

Co-authored-by: Lidya Demilew <[email protected]>
Co-authored-by: Hao Lian <[email protected]>
Co-authored-by: Tudor Brindus <[email protected]>

Author: 4 people

.github/workflows/build.yml

@@ -37,10 +37,31 @@ jobs:
           wget http://ftp.debian.org/debian/pool/main/u/upx-ucl/upx-ucl_3.96-2_amd64.deb
           sudo dpkg -i upx-ucl_3.96-2_amd64.deb
 
-      - name: Use OCaml ${{ matrix.ocaml-version }}
+      - name: "Install apt packages"
         run: |
           sudo apt-get update
           sudo apt-get install bubblewrap musl-tools
+
+      - name: Build zlib with musl
+        run: |
+          mkdir musl-zlib
+          curl -L https://zlib.net/zlib-1.3.1.tar.gz | tar -xz -C musl-zlib --strip-components=1
+          cd musl-zlib
+          CC=musl-gcc ./configure --libdir=/usr/lib/x86_64-linux-musl --includedir=/usr/include/x86_64-linux-musl
+          make -j$(nproc)
+          sudo make install
+      
+      - name: Build zstd with musl
+        run: |
+          mkdir musl-zstd
+          curl -L https://github.com/facebook/zstd/releases/download/v1.5.5/zstd-1.5.5.tar.gz | \
+            tar -xz -C musl-zstd --strip-components=1
+          cd musl-zstd
+          CC=musl-gcc make -j$(nproc)
+          sudo make INCLUDEDIR=/usr/include/x86_64-linux-musl LIBDIR=/usr/lib/x86_64-linux-musl install
+
+      - name: Use OCaml ${{ matrix.ocaml-version }}
+        run: |
           sudo wget -O /usr/local/bin/opam https://github.com/ocaml/opam/releases/download/2.1.2/opam-2.1.2-x86_64-linux
           sudo chmod a+x /usr/local/bin/opam
 

README.md

@@ -71,7 +71,7 @@ You can point magic-trace at a function such that when your application calls it
 
 1. [Here](https://raw.githubusercontent.com/janestreet/magic-trace/master/demo/demo.c)'s a sample C program to try out. It's a slightly modified version of the example in `man 3 dlopen`. Download that, build it with `gcc demo.c -ldl -o demo`, then leave it running `./demo`. We're going to use that program to learn how `dlopen` works.
 
-2. Run `magic-trace attach -pid $(pidof demo)`. When you see the message that it's successfully attached, wait a couple seconds and <kbd>Ctrl</kbd>+<kbd>C</kbd> `magic-trace`. It will output a file called `trace.fxt` in your working directory.
+2. Run `magic-trace attach -pid $(pidof demo)`. When you see the message that it's successfully attached, wait a couple seconds and <kbd>Ctrl</kbd>+<kbd>C</kbd> `magic-trace`. It will output a file called `trace.fxt.gz` in your working directory.
 
 <p align="center">
   <img src="docs/assets/stage-1.gif">

magic-trace.opam

@@ -12,6 +12,7 @@ build: [
 depends: [
   "ocaml"        {>= "4.14"}
   "async"
+  "camlzip"
   "cohttp"
   "cohttp_static_handler"
   "core"
@@ -23,6 +24,7 @@ depends: [
   "dune"         {>= "2.0.0"}
   "owee"         {>= "0.6"}
   "re"           {>= "1.8.0"}
+  "zstandard"
 ]
 synopsis: "Collects and displays high-resolution traces of what a process is doing"
 description: "https://github.com/janestreet/magic-trace"

src/tracing_tool_output.ml

@@ -129,7 +129,7 @@ type t =
 
 let param =
   let%map_open.Command output_path =
-    let default = "trace.fxt" in
+    let default = "trace.fxt.gz" in
     flag
       "output"
       (optional_with_default default string)
@@ -178,7 +178,18 @@ let write_and_maybe_serve
        serving the new trace, which is unlikely to be what the user expected. *)
     let indirect_store_path = [%string "/proc/self/fd/%{fd#Core_unix.File_descr}"] in
     let writer =
-      Tracing_zero.Writer.create_for_file ?num_temp_strs ~filename:indirect_store_path ()
+      let file_format : Tracing_zero.Writer.File_format.t =
+        if Filename.check_suffix filename ".gz"
+        then Gzip
+        else if Filename.check_suffix filename ".zst"
+        then Zstandard
+        else Uncompressed
+      in
+      Tracing_zero.Writer.create_for_file
+        ?num_temp_strs
+        ~file_format
+        ~filename:indirect_store_path
+        ()
     in
     let%bind.Deferred.Or_error res = f ~events_writer:None ~writer:(Some writer) () in
     let%map () =

vendor/tracing/src/trace.mli

@@ -51,7 +51,8 @@ open! Core
 type t
 
 (** Open a file to write trace events to in the Fuchsia Trace Format, suggested extension
-    is [.fxt].
+    is [.fxt] for an uncompressed file and [.fxt.gz] for a gzip compressed one. While [.zst]
+    will produce a Zstandard compressed file, the perfetto viewer does not yet support it.
 
     If [base_time] is provided, a time initialization record will be written which
     records what absolute time corresponds to [Time_ns.Span.zero]. *)

vendor/tracing/zero/destinations.ml

@@ -1,6 +1,5 @@
 open! Core
 
-
 let direct_file_destination ?(buffer_size = 4096 * 16) ~filename () =
   let buf = Iobuf.create ~len:buffer_size in
   let file = Core_unix.openfile ~mode:[ O_CREAT; O_TRUNC; O_RDWR ] filename in
@@ -32,7 +31,121 @@ let direct_file_destination ?(buffer_size = 4096 * 16) ~filename () =
   (module Dest : Writer_intf.Destination)
 ;;
 
-let file_destination ~filename () = direct_file_destination ~filename ()
+(* While Zstandard has the best compression, perfetto does not yet understand the format. *)
+let zstd_file_destination ?(buffer_size = 64 * 1024) ~filename () =
+  let buf = Iobuf.create ~len:buffer_size in
+  let compression_level = 5 in
+  (* Ensure the compression buffer is large enough for the worst case of an input of
+     [buffer_size]. *)
+  let compressed_buf =
+    let len =
+      buffer_size
+      |> Int64.of_int
+      |> Zstandard.compression_output_size_bound
+      |> Int64.to_int_exn
+    in
+    Iobuf.create ~len
+  in
+  let file = Core_unix.openfile ~mode:[ O_CREAT; O_TRUNC; O_CLOEXEC; O_RDWR ] filename in
+  let written = ref 0 in
+  let compression_context = Zstandard.Compression_context.create () in
+  let flush () =
+    Iobuf.rewind buf;
+    Iobuf.advance buf !written;
+    Iobuf.flip_lo buf;
+    let input =
+      Zstandard.Input.from_bigstring
+        ~pos:(Iobuf.Expert.lo buf)
+        ~len:(Iobuf.length buf)
+        (Iobuf.Expert.buf buf)
+    in
+    let output =
+      Zstandard.Output.in_buffer
+        ~pos:(Iobuf.Expert.lo compressed_buf)
+        ~len:(Iobuf.length compressed_buf)
+        (Iobuf.Expert.buf compressed_buf)
+    in
+    let compressed_length =
+      Zstandard.With_explicit_context.compress
+        compression_context
+        ~compression_level
+        ~input
+        ~output
+    in
+    Iobuf.advance compressed_buf compressed_length;
+    Iobuf.flip_lo compressed_buf;
+    Iobuf_unix.write compressed_buf file;
+    written := 0;
+    Iobuf.reset buf;
+    Iobuf.reset compressed_buf
+  in
+  let module Dest = struct
+    let next_buf ~ensure_capacity =
+      flush ();
+      if ensure_capacity > Iobuf.length buf
+      then failwith "Not enough buffer space in [zstd_file_destination]";
+      buf
+    ;;
+
+    let wrote_bytes count = written := !written + count
+
+    let close () =
+      flush ();
+      Zstandard.Compression_context.free compression_context;
+      Core_unix.close file
+    ;;
+  end
+  in
+  (module Dest : Writer_intf.Destination)
+;;
+
+let gzip_file_destination ?(buffer_size = 64 * 1024) ~filename () =
+  let buf = Iobuf.create ~len:buffer_size in
+  let bytes = Bytes.create buffer_size in
+  let file = Core_unix.openfile ~mode:[ O_CREAT; O_TRUNC; O_CLOEXEC; O_RDWR ] filename in
+  let out_channel =
+    let oc = Core_unix.out_channel_of_descr file in
+    (* Consider making the compression level an environment variable for
+       experimentation. *)
+    Gzip.open_out_chan ~level:6 oc
+  in
+  let written = ref 0 in
+  let flush () =
+    Iobuf.rewind buf;
+    Iobuf.advance buf !written;
+    Iobuf.flip_lo buf;
+    Iobuf.Peek.To_bytes.blit
+      ~src:(Iobuf.read_only buf) ~src_pos:0 ~dst:bytes ~dst_pos:0 ~len:!written;
+    Gzip.output out_channel bytes 0 !written;
+    written := 0;
+    Iobuf.reset buf;
+  in
+  let module Dest = struct
+    let next_buf ~ensure_capacity =
+      flush ();
+      if ensure_capacity > Iobuf.length buf
+      then failwith "Not enough buffer space in [gzip_file_destination]";
+      buf
+    ;;
+
+    let wrote_bytes count = written := !written + count
+
+    let close () =
+      flush ();
+      (* [close_out] also closes the underlying file descr. *)
+      Gzip.close_out out_channel
+    ;;
+  end
+  in
+  (module Dest : Writer_intf.Destination)
+;;
+
+let file_destination ?(file_format = Writer_intf.File_format.Uncompressed) ~filename () =
+  match file_format with
+  | Uncompressed -> direct_file_destination ~filename ()
+  | Gzip -> gzip_file_destination ~filename ()
+  | Zstandard -> zstd_file_destination ~filename ()
+;;
 
 let iobuf_destination buf =
   (* We give out an [Iobuf] with a shared underlying [Bigstring] but different pointers

vendor/tracing/zero/destinations.mli

@@ -7,8 +7,25 @@ val direct_file_destination
   -> unit
   -> (module Writer_intf.Destination)
 
-(** Write to a file in some way with the best available performance. *)
-val file_destination : filename:string -> unit -> (module Writer_intf.Destination)
+(** Write to a zstd compressed file using synchronous writes, not suitable for low latency
+    applications. *)
+val zstd_file_destination
+  :  ?buffer_size:int
+  -> filename:string
+  -> unit
+  -> (module Writer_intf.Destination)
+
+(** Write to a gzip compressed file using synchronous writes, not suitable for low latency
+    applications. *)
+val gzip_file_destination
+  :  ?buffer_size:int
+  -> filename:string
+  -> unit
+  -> (module Writer_intf.Destination)
+
+(** Write to a file in some way with the best available performance. [format] defaults to
+    [Uncompressed]. *)
+val file_destination : ?file_format:Writer_intf.File_format.t -> filename:string -> unit -> (module Writer_intf.Destination)
 
 (** Write to a provided [Iobuf.t], throws an exception if the buffer runs out of space.
     Mostly intended for use in tests. After the [Destination] is closed, sets the window

vendor/tracing/zero/dune

@@ -1,4 +1,4 @@
 (library (name tracing_zero) (public_name tracing.tracing_zero)
  (preprocess (pps ppx_jane))
- (libraries core core_kernel.iobuf core_unix.iobuf_unix
-  core_unix.time_stamp_counter))
+ (libraries camlzip core core_kernel.iobuf core_unix.iobuf_unix
+  core_unix.time_stamp_counter zstandard))

vendor/tracing/zero/writer.ml

@@ -620,8 +620,14 @@ module Expert = struct
   module Write_arg_unchecked = Write_arg_unchecked
 end
 
-let create_for_file ?num_temp_strs ~filename () =
-  let destination = Destinations.file_destination ~filename () in
+
+module File_format = Writer_intf.File_format
+
+(** Allocates a writer which writes to [filename] with [num_temp_strs] temporary string
+    slots (see [set_temp_string_slot]), with increases in [num_temp_strs] reducing the
+    number of strings which can be allocated with [intern_string]. *)
+let create_for_file ?num_temp_strs ?file_format ~filename () =
+  let destination = Destinations.file_destination ?file_format ~filename () in
   Expert.create ?num_temp_strs ~destination ()
 ;;
 

vendor/tracing/zero/writer.mli

@@ -11,10 +11,12 @@ open! Core
 
 type t
 
+module File_format = Writer_intf.File_format
+
 (** Allocates a writer which writes to [filename] with [num_temp_strs] temporary string
     slots (see [set_temp_string_slot]), with increases in [num_temp_strs] reducing the
     number of strings which can be allocated with [intern_string]. *)
-val create_for_file : ?num_temp_strs:int -> filename:string -> unit -> t
+val create_for_file : ?num_temp_strs:int -> ?file_format:Writer_intf.File_format.t -> filename:string -> unit -> t
 
 val close : t -> unit