Merge pull request #36 from c-cube/simon/fuchsia-via-subscriber-2025-05-02

full refactoring of TEF and fuchsia backends

.github/workflows/main.yml

@@ -51,6 +51,11 @@ jobs:
       - run: opam install hmap
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
 
+      - run: opam install picos_aux
+        if: matrix.ocaml-compiler != '4.08.x' && matrix.ocaml-compiler != '4.12.x'
+      - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
+        if: matrix.ocaml-compiler != '4.08.x' && matrix.ocaml-compiler != '4.12.x'
+
       - run: opam install mtime
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
 

bench/bench_fuchsia_write.ml

@@ -1,23 +1,34 @@
-open Trace_fuchsia_write
+open Trace_fuchsia
+open Trace_fuchsia.Writer
 module B = Benchmark
 
 let pf = Printf.printf
 
-let encode_1_span (out : Output.t) () =
+let encode_1000_span (bufs : Buf_chain.t) () =
-  Event.Duration_complete.encode out ~name:"span" ~t_ref:(Thread_ref.Ref 5)
+  for _i = 1 to 1000 do
-    ~time_ns:100_000L ~end_time_ns:5_000_000L ~args:[] ()
+    Event.Duration_complete.encode bufs ~name:"span" ~t_ref:(Thread_ref.Ref 5)
+      ~time_ns:100_000L ~end_time_ns:5_000_000L ~args:[] ()
+  done;
+  Buf_chain.ready_all_non_empty bufs;
+  Buf_chain.pop_ready bufs ~f:ignore;
+  ()
 
-let encode_3_span (out : Output.t) () =
+let encode_300_span (bufs : Buf_chain.t) () =
-  Event.Duration_complete.encode out ~name:"outer" ~t_ref:(Thread_ref.Ref 5)
+  for _i = 1 to 100 do
-    ~time_ns:100_000L ~end_time_ns:5_000_000L ~args:[] ();
+    Event.Duration_complete.encode bufs ~name:"outer" ~t_ref:(Thread_ref.Ref 5)
-  Event.Duration_complete.encode out ~name:"inner" ~t_ref:(Thread_ref.Ref 5)
+      ~time_ns:100_000L ~end_time_ns:5_000_000L ~args:[] ();
-    ~time_ns:180_000L ~end_time_ns:4_500_000L ~args:[] ();
+    Event.Duration_complete.encode bufs ~name:"inner" ~t_ref:(Thread_ref.Ref 5)
-  Event.Instant.encode out ~name:"hello" ~time_ns:1_234_567L
+      ~time_ns:180_000L ~end_time_ns:4_500_000L ~args:[] ();
-    ~t_ref:(Thread_ref.Ref 5)
+    Event.Instant.encode bufs ~name:"hello" ~time_ns:1_234_567L
-    ~args:[ "x", `Int 42 ]
+      ~t_ref:(Thread_ref.Ref 5)
-    ()
+      ~args:[ "x", A_int 42 ]
+      ()
+  done;
+  Buf_chain.ready_all_non_empty bufs;
+  Buf_chain.pop_ready bufs ~f:ignore;
+  ()
 
-let time_per_iter_ns (samples : B.t list) : float =
+let time_per_iter_ns n_per_iter (samples : B.t list) : float =
   let n_iters = ref 0L in
   let time = ref 0. in
   List.iter
@@ -25,34 +36,32 @@ let time_per_iter_ns (samples : B.t list) : float =
       n_iters := Int64.add !n_iters s.iters;
       time := !time +. s.stime +. s.utime)
     samples;
-  !time *. 1e9 /. Int64.to_float !n_iters
+  !time *. 1e9 /. (Int64.to_float !n_iters *. float n_per_iter)
 
 let () =
   let buf_pool = Buf_pool.create () in
-  let out =
+  let bufs = Buf_chain.create ~sharded:false ~buf_pool () in
-    Output.create ~buf_pool
-      ~send_buf:(fun buf -> Buf_pool.recycle buf_pool buf)
-      ()
-  in
 
-  let samples = B.throughput1 4 ~name:"encode_1_span" (encode_1_span out) () in
+  let samples =
+    B.throughput1 4 ~name:"encode_1000_span" (encode_1000_span bufs) ()
+  in
   B.print_gc samples;
 
   let [ (_, samples) ] = samples [@@warning "-8"] in
-
+  let iter_per_ns = time_per_iter_ns 1000 samples in
-  let iter_per_ns = time_per_iter_ns samples in
   pf "%.3f ns/iter\n" iter_per_ns;
 
   ()
 
 let () =
   let buf_pool = Buf_pool.create () in
-  let out =
+  let bufs = Buf_chain.create ~sharded:false ~buf_pool () in
-    Output.create ~buf_pool
+  let samples =
-      ~send_buf:(fun buf -> Buf_pool.recycle buf_pool buf)
+    B.throughput1 4 ~name:"encode_300_span" (encode_300_span bufs) ()
-      ()
   in
-
-  let samples = B.throughput1 4 ~name:"encode_3_span" (encode_3_span out) () in
   B.print_gc samples;
+
+  let [ (_, samples) ] = samples [@@warning "-8"] in
+  let iter_per_ns = time_per_iter_ns 300 samples in
+  pf "%.3f ns/iter\n" iter_per_ns;
   ()

bench/dune

@@ -20,4 +20,4 @@
 (executable
  (name bench_fuchsia_write)
  (modules bench_fuchsia_write)
- (libraries benchmark trace-fuchsia.write))
+ (libraries benchmark trace-fuchsia))

dune-project

@@ -28,6 +28,7 @@
  (depopts
   hmap
   unix
+  (picos_aux (>= 0.6))
   (mtime
    (>= 2.0)))
  (tags

src/event/dune

@@ -0,0 +1,7 @@
+
+(library
+ (name trace_event)
+ (public_name trace.event)
+ (synopsis "Turns subscriber callbacks into an event type")
+ (libraries
+  (re_export trace.core) (re_export trace.subscriber)))

src/event/event.ml

@@ -1,26 +1,28 @@
 open Trace_core
 module Sub = Trace_subscriber
 
-(** An event, specialized for TEF *)
+(** An event with TEF/fuchsia semantics *)
 type t =
   | E_tick
+  | E_init of { time_ns: int64 }
+  | E_shutdown of { time_ns: int64 }
   | E_message of {
       tid: int;
       msg: string;
-      time_us: float;
+      time_ns: int64;
       data: (string * Sub.user_data) list;
     }
   | E_define_span of {
       tid: int;
       name: string;
-      time_us: float;
+      time_ns: int64;
       id: span;
       fun_name: string option;
       data: (string * Sub.user_data) list;
     }
   | E_exit_span of {
       id: span;
-      time_us: float;
+      time_ns: int64;
     }
   | E_add_data of {
       id: span;
@@ -29,7 +31,7 @@ type t =
   | E_enter_manual_span of {
       tid: int;
       name: string;
-      time_us: float;
+      time_ns: int64;
       id: trace_id;
       flavor: Sub.flavor option;
       fun_name: string option;
@@ -38,7 +40,7 @@ type t =
   | E_exit_manual_span of {
       tid: int;
       name: string;
-      time_us: float;
+      time_ns: int64;
       flavor: Sub.flavor option;
       data: (string * Sub.user_data) list;
       id: trace_id;
@@ -46,7 +48,7 @@ type t =
   | E_counter of {
       name: string;
       tid: int;
-      time_us: float;
+      time_ns: int64;
       n: float;
     }
   | E_name_process of { name: string }
@@ -54,3 +56,8 @@ type t =
       tid: int;
       name: string;
     }
+  | E_extension_event of {
+      tid: int;
+      time_ns: int64;
+      ext: Trace_core.extension_event;
+    }

src/event/subscriber.ml

@@ -0,0 +1,53 @@
+open Trace_core
+open Event
+
+type event_consumer = { on_event: Event.t -> unit } [@@unboxed]
+(** Callback for events *)
+
+module Callbacks : Sub.Callbacks.S with type st = event_consumer = struct
+  type st = event_consumer
+
+  let on_init (self : st) ~time_ns = self.on_event (E_init { time_ns })
+  let on_shutdown (self : st) ~time_ns = self.on_event (E_shutdown { time_ns })
+
+  let on_name_process (self : st) ~time_ns:_ ~tid:_ ~name : unit =
+    self.on_event @@ E_name_process { name }
+
+  let on_name_thread (self : st) ~time_ns:_ ~tid ~name : unit =
+    self.on_event @@ E_name_thread { tid; name }
+
+  let[@inline] on_enter_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
+      ~__LINE__:_ ~time_ns ~tid ~data ~name span : unit =
+    self.on_event
+    @@ E_define_span { tid; name; time_ns; id = span; fun_name; data }
+
+  let on_exit_span (self : st) ~time_ns ~tid:_ span : unit =
+    self.on_event @@ E_exit_span { id = span; time_ns }
+
+  let on_add_data (self : st) ~data span =
+    if data <> [] then self.on_event @@ E_add_data { id = span; data }
+
+  let on_message (self : st) ~time_ns ~tid ~span:_ ~data msg : unit =
+    self.on_event @@ E_message { tid; time_ns; msg; data }
+
+  let on_counter (self : st) ~time_ns ~tid ~data:_ ~name f : unit =
+    self.on_event @@ E_counter { name; n = f; time_ns; tid }
+
+  let on_enter_manual_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
+      ~__LINE__:_ ~time_ns ~tid ~parent:_ ~data ~name ~flavor ~trace_id _span :
+      unit =
+    self.on_event
+    @@ E_enter_manual_span
+         { id = trace_id; time_ns; tid; data; name; fun_name; flavor }
+
+  let on_exit_manual_span (self : st) ~time_ns ~tid ~name ~data ~flavor
+      ~trace_id (_ : span) : unit =
+    self.on_event
+    @@ E_exit_manual_span { tid; id = trace_id; name; time_ns; data; flavor }
+
+  let on_extension_event (self : st) ~time_ns ~tid ext : unit =
+    self.on_event @@ E_extension_event { tid; time_ns; ext }
+end
+
+let subscriber (consumer : event_consumer) : Sub.t =
+  Sub.Subscriber.Sub { st = consumer; callbacks = (module Callbacks) }

src/fuchsia/bg_thread.ml

@@ -1,62 +0,0 @@
-open Common_
-
-type out =
-  [ `Stdout
-  | `Stderr
-  | `File of string
-  ]
-
-type event =
-  | E_write_buf of Buf.t
-  | E_tick
-
-type state = {
-  buf_pool: Buf_pool.t;
-  oc: out_channel;
-  events: event B_queue.t;
-}
-
-let with_out_ (out : out) f =
-  let oc, must_close =
-    match out with
-    | `Stdout -> stdout, false
-    | `Stderr -> stderr, false
-    | `File path -> open_out path, true
-  in
-
-  if must_close then (
-    let finally () = close_out_noerr oc in
-    Fun.protect ~finally (fun () -> f oc)
-  ) else
-    f oc
-
-let handle_ev (self : state) (ev : event) : unit =
-  match ev with
-  | E_tick -> flush self.oc
-  | E_write_buf buf ->
-    output self.oc buf.buf 0 buf.offset;
-    Buf_pool.recycle self.buf_pool buf
-
-let bg_loop (self : state) : unit =
-  let continue = ref true in
-
-  while !continue do
-    match B_queue.pop_all self.events with
-    | exception B_queue.Closed -> continue := false
-    | evs -> List.iter (handle_ev self) evs
-  done
-
-let bg_thread ~buf_pool ~out ~(events : event B_queue.t) () : unit =
-  let@ oc = with_out_ out in
-  let st = { oc; buf_pool; events } in
-  bg_loop st
-
-(** Thread that simply regularly "ticks", sending events to the background
-    thread so it has a chance to write to the file, and call [f()] *)
-let tick_thread events : unit =
-  try
-    while true do
-      Thread.delay 0.5;
-      B_queue.push events E_tick
-    done
-  with B_queue.Closed -> ()

src/fuchsia/buf.ml

@@ -8,12 +8,14 @@ type t = {
 let empty : t = { buf = Bytes.empty; offset = 0 }
 
 let create (n : int) : t =
+  (* multiple of 8-bytes size *)
   let buf = Bytes.create (round_to_word n) in
   { buf; offset = 0 }
 
 let[@inline] clear self = self.offset <- 0
 let[@inline] available self = Bytes.length self.buf - self.offset
 let[@inline] size self = self.offset
+let[@inline] is_empty self = self.offset = 0
 
 (* see below: we assume little endian *)
 let () = assert (not Sys.big_endian)

src/fuchsia/buf_chain.ml

@@ -0,0 +1,140 @@
+(** A set of buffers in use, and a set of ready buffers *)
+
+open Common_
+
+(** Buffers in use *)
+type buffers =
+  | B_one of { mutable buf: Buf.t }
+  | B_many of Buf.t Lock.t array
+      (** mask(thread id) -> buffer. This reduces contention *)
+
+type t = {
+  bufs: buffers;
+  has_ready: bool A.t;
+  ready: Buf.t Queue.t Lock.t;
+      (** Buffers that are full (enough) and must be written *)
+  buf_pool: Buf_pool.t;
+}
+(** A set of buffers, some of which are ready to be written *)
+
+open struct
+  let shard_log = 4
+  let shard = 1 lsl shard_log
+  let shard_mask = shard - 1
+end
+
+let create ~(sharded : bool) ~(buf_pool : Buf_pool.t) () : t =
+  let bufs =
+    if sharded then (
+      let bufs =
+        Array.init shard (fun _ -> Lock.create @@ Buf_pool.alloc buf_pool)
+      in
+      B_many bufs
+    ) else
+      B_one { buf = Buf_pool.alloc buf_pool }
+  in
+  {
+    bufs;
+    buf_pool;
+    has_ready = A.make false;
+    ready = Lock.create @@ Queue.create ();
+  }
+
+open struct
+  let put_in_ready (self : t) buf : unit =
+    if Buf.size buf > 0 then (
+      let@ q = Lock.with_ self.ready in
+      A.set self.has_ready true;
+      Queue.push buf q
+    )
+
+  let assert_available buf ~available =
+    if Buf.available buf < available then (
+      let msg =
+        Printf.sprintf
+          "fuchsia: buffer is too small (available: %d bytes, needed: %d bytes)"
+          (Buf.available buf) available
+      in
+      failwith msg
+    )
+end
+
+(** Move all non-empty buffers to [ready] *)
+let ready_all_non_empty (self : t) : unit =
+  let@ q = Lock.with_ self.ready in
+  match self.bufs with
+  | B_one r ->
+    if not (Buf.is_empty r.buf) then (
+      Queue.push r.buf q;
+      A.set self.has_ready true;
+      r.buf <- Buf.empty
+    )
+  | B_many bufs ->
+    Array.iter
+      (fun buf ->
+        Lock.update buf (fun buf ->
+            if Buf.size buf > 0 then (
+              Queue.push buf q;
+              A.set self.has_ready true;
+              Buf.empty
+            ) else
+              buf))
+      bufs
+
+let[@inline] has_ready self : bool = A.get self.has_ready
+
+(** Get access to ready buffers, then clean them up automatically *)
+let pop_ready (self : t) ~(f : Buf.t Queue.t -> 'a) : 'a =
+  let@ q = Lock.with_ self.ready in
+  let res = f q in
+
+  (* clear queue *)
+  Queue.iter (Buf_pool.recycle self.buf_pool) q;
+  Queue.clear q;
+  A.set self.has_ready false;
+  res
+
+(** Maximum size available, in words, for a single message *)
+let[@inline] max_size_word (_self : t) : int = fuchsia_buf_size lsr 3
+
+(** Obtain a buffer with at least [available_word] 64-bit words *)
+let with_buf (self : t) ~(available_word : int) (f : Buf.t -> 'a) : 'a =
+  let available = available_word lsl 3 in
+  match self.bufs with
+  | B_one r ->
+    if Buf.available r.buf < available_word then (
+      put_in_ready self r.buf;
+      r.buf <- Buf_pool.alloc self.buf_pool
+    );
+    assert_available r.buf ~available;
+    f r.buf
+  | B_many bufs ->
+    let tid = Thread.(id (self ())) in
+    let masked_tid = tid land shard_mask in
+    let buf_lock = bufs.(masked_tid) in
+    let@ buf = Lock.with_ buf_lock in
+    let buf =
+      if Buf.available buf < available then (
+        put_in_ready self buf;
+        let new_buf = Buf_pool.alloc self.buf_pool in
+        assert_available new_buf ~available;
+        Lock.set_while_locked buf_lock new_buf;
+        new_buf
+      ) else
+        buf
+    in
+    f buf
+
+(** Dispose of resources (here, recycle buffers) *)
+let dispose (self : t) : unit =
+  match self.bufs with
+  | B_one r ->
+    Buf_pool.recycle self.buf_pool r.buf;
+    r.buf <- Buf.empty
+  | B_many bufs ->
+    Array.iter
+      (fun buf_lock ->
+        let@ buf = Lock.with_ buf_lock in
+        Buf_pool.recycle self.buf_pool buf;
+        Lock.set_while_locked buf_lock Buf.empty)
+      bufs

src/fuchsia/buf_pool.ml

@@ -0,0 +1,23 @@
+open Common_
+open Trace_private_util
+
+type t = Buf.t Rpool.t
+
+let create ?(max_size = 64) () : t =
+  Rpool.create ~max_size ~clear:Buf.clear
+    ~create:(fun () -> Buf.create fuchsia_buf_size)
+    ()
+
+let alloc = Rpool.alloc
+let[@inline] recycle self buf = if buf != Buf.empty then Rpool.recycle self buf
+
+let with_ (self : t) f =
+  let x = alloc self in
+  try
+    let res = f x in
+    recycle self x;
+    res
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    recycle self x;
+    Printexc.raise_with_backtrace e bt

src/fuchsia/common_.ml

@@ -1,12 +1,22 @@
 module A = Trace_core.Internal_.Atomic_
-module FWrite = Trace_fuchsia_write
+module Sub = Trace_subscriber
-module B_queue = Trace_private_util.B_queue
-module Buf = FWrite.Buf
-module Buf_pool = FWrite.Buf_pool
-module Output = FWrite.Output
 
 let on_tracing_error =
   ref (fun s -> Printf.eprintf "trace-fuchsia error: %s\n%!" s)
 
 let ( let@ ) = ( @@ )
 let spf = Printf.sprintf
+
+let with_lock lock f =
+  Mutex.lock lock;
+  try
+    let res = f () in
+    Mutex.unlock lock;
+    res
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    Mutex.unlock lock;
+    Printexc.raise_with_backtrace e bt
+
+(** Buffer size we use. *)
+let fuchsia_buf_size = 1 lsl 16

src/fuchsia/dune

@@ -6,8 +6,8 @@
  (libraries
   trace.core
   trace.private.util
+  trace.subscriber
   thread-local-storage
-  (re_export trace-fuchsia.write)
   bigarray
   mtime
   mtime.clock.os

src/fuchsia/exporter.ml

@@ -0,0 +1,61 @@
+(** An exporter, takes buffers with fuchsia events, and writes them somewhere *)
+
+open Common_
+
+type t = {
+  write_bufs: Buf.t Queue.t -> unit;
+      (** Takes buffers and writes them somewhere. The buffers are only valid
+          during this call and must not be stored. The queue must not be
+          modified. *)
+  flush: unit -> unit;  (** Force write *)
+  close: unit -> unit;  (** Close underlying resources *)
+}
+(** An exporter, takes buffers and writes them somewhere. This should be
+    thread-safe if used in a threaded environment. *)
+
+open struct
+  let with_lock lock f =
+    Mutex.lock lock;
+    try
+      let res = f () in
+      Mutex.unlock lock;
+      res
+    with e ->
+      let bt = Printexc.get_raw_backtrace () in
+      Mutex.unlock lock;
+      Printexc.raise_with_backtrace e bt
+end
+
+(** Export to the channel
+    @param close_channel if true, closing the exporter will close the channel *)
+let of_out_channel ~close_channel oc : t =
+  let lock = Mutex.create () in
+  let closed = ref false in
+  let flush () =
+    let@ () = with_lock lock in
+    flush oc
+  in
+  let close () =
+    let@ () = with_lock lock in
+    if not !closed then (
+      closed := true;
+      if close_channel then close_out_noerr oc
+    )
+  in
+  let write_bufs bufs =
+    if not (Queue.is_empty bufs) then
+      let@ () = with_lock lock in
+      Queue.iter (fun (buf : Buf.t) -> output oc buf.buf 0 buf.offset) bufs
+  in
+  { flush; close; write_bufs }
+
+let of_buffer (buffer : Buffer.t) : t =
+  let buffer = Lock.create buffer in
+  let write_bufs bufs =
+    if not (Queue.is_empty bufs) then
+      let@ buffer = Lock.with_ buffer in
+      Queue.iter
+        (fun (buf : Buf.t) -> Buffer.add_subbytes buffer buf.buf 0 buf.offset)
+        bufs
+  in
+  { flush = ignore; close = ignore; write_bufs }

src/fuchsia/fcollector.ml

@@ -1,410 +0,0 @@
-open Trace_core
-open Common_
-module TLS = Thread_local_storage
-module Int_map = Map.Make (Int)
-
-let pid = Unix.getpid ()
-
-(** Thread-local stack of span info *)
-module Span_info_stack : sig
-  type t
-
-  val create : unit -> t
-
-  val push :
-    t ->
-    span ->
-    name:string ->
-    start_time_ns:int64 ->
-    data:(string * user_data) list ->
-    unit
-
-  val pop : t -> int64 * string * int64 * (string * user_data) list
-  val find_ : t -> span -> int option
-  val add_data : t -> int -> (string * user_data) list -> unit
-end = struct
-  module BA = Bigarray
-  module BA1 = Bigarray.Array1
-
-  type int64arr = (int64, BA.int64_elt, BA.c_layout) BA1.t
-
-  type t = {
-    mutable len: int;
-    mutable span: int64arr;
-    mutable start_time_ns: int64arr;
-    mutable name: string array;
-    mutable data: (string * user_data) list array;
-  }
-
-  let init_size_ = 1
-
-  let create () : t =
-    {
-      len = 0;
-      span = BA1.create BA.Int64 BA.C_layout init_size_;
-      start_time_ns = BA1.create BA.Int64 BA.C_layout init_size_;
-      name = Array.make init_size_ "";
-      data = Array.make init_size_ [];
-    }
-
-  let[@inline] cap self = Array.length self.name
-
-  let grow_ (self : t) : unit =
-    let new_cap = 2 * cap self in
-    let new_span = BA1.create BA.Int64 BA.C_layout new_cap in
-    BA1.blit self.span (BA1.sub new_span 0 self.len);
-    let new_startime_ns = BA1.create BA.Int64 BA.C_layout new_cap in
-    BA1.blit self.start_time_ns (BA1.sub new_startime_ns 0 self.len);
-    let new_name = Array.make new_cap "" in
-    Array.blit self.name 0 new_name 0 self.len;
-    let new_data = Array.make new_cap [] in
-    Array.blit self.data 0 new_data 0 self.len;
-    self.span <- new_span;
-    self.start_time_ns <- new_startime_ns;
-    self.name <- new_name;
-    self.data <- new_data
-
-  let push (self : t) (span : int64) ~name ~start_time_ns ~data =
-    if cap self = self.len then grow_ self;
-    BA1.set self.span self.len span;
-    BA1.set self.start_time_ns self.len start_time_ns;
-    Array.set self.name self.len name;
-    Array.set self.data self.len data;
-    self.len <- self.len + 1
-
-  let pop (self : t) =
-    assert (self.len > 0);
-    self.len <- self.len - 1;
-
-    let span = BA1.get self.span self.len in
-    let name = self.name.(self.len) in
-    let start_time_ns = BA1.get self.start_time_ns self.len in
-    let data = self.data.(self.len) in
-
-    (* avoid holding onto old values *)
-    Array.set self.name self.len "";
-    Array.set self.data self.len [];
-
-    span, name, start_time_ns, data
-
-  let[@inline] add_data self i d : unit =
-    assert (i < self.len);
-    self.data.(i) <- List.rev_append d self.data.(i)
-
-  exception Found of int
-
-  let[@inline] find_ (self : t) span : _ option =
-    try
-      for i = self.len - 1 downto 0 do
-        if Int64.equal (BA1.get self.span i) span then raise_notrace (Found i)
-      done;
-
-      None
-    with Found i -> Some i
-end
-
-type async_span_info = {
-  flavor: [ `Sync | `Async ] option;
-  name: string;
-  mutable data: (string * user_data) list;
-}
-
-let key_async_data : async_span_info Meta_map.key = Meta_map.Key.create ()
-
-open struct
-  let state_id_ = A.make 0
-
-  (* re-raise exception with its backtrace *)
-  external reraise : exn -> 'a = "%reraise"
-end
-
-type per_thread_state = {
-  tid: int;
-  state_id: int;  (** ID of the current collector state *)
-  local_span_id_gen: int A.t;  (** Used for thread-local spans *)
-  mutable thread_ref: FWrite.Thread_ref.t;
-  mutable out: Output.t option;
-  spans: Span_info_stack.t;  (** In-flight spans *)
-}
-
-type state = {
-  active: bool A.t;
-  events: Bg_thread.event B_queue.t;
-  span_id_gen: int A.t;  (** Used for async spans *)
-  bg_thread: Thread.t;
-  buf_pool: Buf_pool.t;
-  next_thread_ref: int A.t;  (** in [0x01 .. 0xff], to allocate thread refs *)
-  per_thread: per_thread_state Int_map.t A.t array;
-      (** the state keeps tabs on thread-local state, so it can flush writers at
-          the end. This is a tid-sharded array of maps. *)
-}
-
-let[@inline] mk_trace_id (self : state) : trace_id =
-  let n = A.fetch_and_add self.span_id_gen 1 in
-  let b = Bytes.create 8 in
-  Bytes.set_int64_le b 0 (Int64.of_int n);
-  Bytes.unsafe_to_string b
-
-let key_thread_local_st : per_thread_state TLS.t = TLS.create ()
-
-let[@inline never] mk_thread_local_st () =
-  let tid = Thread.id @@ Thread.self () in
-  let st =
-    {
-      tid;
-      state_id = A.get state_id_;
-      thread_ref = FWrite.Thread_ref.inline ~pid ~tid;
-      local_span_id_gen = A.make 0;
-      out = None;
-      spans = Span_info_stack.create ();
-    }
-  in
-  TLS.set key_thread_local_st st;
-  st
-
-let[@inline] get_thread_local_st () =
-  match TLS.get_opt key_thread_local_st with
-  | Some k -> k
-  | None -> mk_thread_local_st ()
-
-let out_of_st (st : state) : Output.t =
-  FWrite.Output.create () ~buf_pool:st.buf_pool ~send_buf:(fun buf ->
-      try B_queue.push st.events (E_write_buf buf) with B_queue.Closed -> ())
-
-module C
-    (St : sig
-      val st : state
-    end)
-    () =
-struct
-  open St
-
-  let state_id = 1 + A.fetch_and_add state_id_ 1
-
-  (** prepare the thread's state *)
-  let[@inline never] update_or_init_local_state (self : per_thread_state) : unit
-      =
-    (* get an output *)
-    let out = out_of_st st in
-    self.out <- Some out;
-
-    (* try to allocate a thread ref for current thread *)
-    let th_ref = A.fetch_and_add st.next_thread_ref 1 in
-    if th_ref <= 0xff then (
-      self.thread_ref <- FWrite.Thread_ref.ref th_ref;
-      FWrite.Thread_record.encode out ~as_ref:th_ref ~tid:self.tid ~pid ()
-    );
-
-    (* add to [st]'s list of threads *)
-    let shard_of_per_thread = st.per_thread.(self.tid land 0b1111) in
-    while
-      let old = A.get shard_of_per_thread in
-      not
-        (A.compare_and_set shard_of_per_thread old
-           (Int_map.add self.tid self old))
-    do
-      ()
-    done;
-
-    let on_exit _ =
-      while
-        let old = A.get shard_of_per_thread in
-        not
-          (A.compare_and_set shard_of_per_thread old
-             (Int_map.remove self.tid old))
-      do
-        ()
-      done;
-      Option.iter Output.flush self.out
-    in
-
-    (* after thread exits, flush output and remove from global list *)
-    Gc.finalise on_exit (Thread.self ());
-    ()
-
-  (** Obtain the output for the current thread *)
-  let[@inline] get_thread_output () : Output.t * per_thread_state =
-    let tls = get_thread_local_st () in
-    if tls.state_id != state_id || tls.out == None then
-      update_or_init_local_state tls;
-    let out =
-      match tls.out with
-      | None -> assert false
-      | Some o -> o
-    in
-    out, tls
-
-  let close_per_thread (tls : per_thread_state) =
-    Option.iter Output.flush tls.out
-
-  (** flush all outputs *)
-  let flush_all_outputs_ () =
-    Array.iter
-      (fun shard ->
-        let tls_l = A.get shard in
-        Int_map.iter (fun _tid tls -> close_per_thread tls) tls_l)
-      st.per_thread
-
-  let shutdown () =
-    if A.exchange st.active false then (
-      flush_all_outputs_ ();
-
-      B_queue.close st.events;
-      (* wait for writer thread to be done. The writer thread will exit
-         after processing remaining events because the queue is now closed *)
-      Thread.join st.bg_thread
-    )
-
-  let enter_span ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : span =
-    let tls = get_thread_local_st () in
-    let span = Int64.of_int (A.fetch_and_add tls.local_span_id_gen 1) in
-    let time_ns = Time.now_ns () in
-    Span_info_stack.push tls.spans span ~name ~data ~start_time_ns:time_ns;
-    span
-
-  let exit_span span : unit =
-    let out, tls = get_thread_output () in
-    let end_time_ns = Time.now_ns () in
-
-    let span', name, start_time_ns, data = Span_info_stack.pop tls.spans in
-    if span <> span' then
-      !on_tracing_error
-        (spf "span mismatch: top is %Ld, expected %Ld" span' span)
-    else
-      FWrite.Event.Duration_complete.encode out ~name ~t_ref:tls.thread_ref
-        ~time_ns:start_time_ns ~end_time_ns ~args:data ()
-
-  let with_span ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name f =
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-    let span = Int64.of_int (A.fetch_and_add tls.local_span_id_gen 1) in
-    Span_info_stack.push tls.spans span ~start_time_ns:time_ns ~data ~name;
-
-    let[@inline] exit () : unit =
-      let end_time_ns = Time.now_ns () in
-
-      let _span', _, _, data = Span_info_stack.pop tls.spans in
-      assert (span = _span');
-      FWrite.Event.Duration_complete.encode out ~name ~time_ns ~end_time_ns
-        ~t_ref:tls.thread_ref ~args:data ()
-    in
-
-    try
-      let x = f span in
-      exit ();
-      x
-    with exn ->
-      exit ();
-      reraise exn
-
-  let add_data_to_span span data =
-    let tls = get_thread_local_st () in
-    match Span_info_stack.find_ tls.spans span with
-    | None -> !on_tracing_error (spf "unknown span %Ld" span)
-    | Some idx -> Span_info_stack.add_data tls.spans idx data
-
-  let enter_manual_span ~(parent : explicit_span_ctx option) ~flavor
-      ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : explicit_span =
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-
-    (* get the id, or make a new one *)
-    let trace_id =
-      match parent with
-      | Some m -> m.trace_id
-      | None -> mk_trace_id st
-    in
-
-    FWrite.Event.Async_begin.encode out ~name ~args:data ~t_ref:tls.thread_ref
-      ~time_ns ~async_id:trace_id ();
-    {
-      span = 0L;
-      trace_id;
-      meta = Meta_map.(empty |> add key_async_data { name; flavor; data = [] });
-    }
-
-  let exit_manual_span (es : explicit_span) : unit =
-    let { name; data; flavor = _ } = Meta_map.find_exn key_async_data es.meta in
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-
-    FWrite.Event.Async_end.encode out ~name ~t_ref:tls.thread_ref ~time_ns
-      ~args:data ~async_id:es.trace_id ()
-
-  let add_data_to_manual_span (es : explicit_span) data =
-    let m = Meta_map.find_exn key_async_data es.meta in
-    m.data <- List.rev_append data m.data
-
-  let message ?span:_ ~data msg : unit =
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-    FWrite.Event.Instant.encode out ~name:msg ~time_ns ~t_ref:tls.thread_ref
-      ~args:data ()
-
-  let counter_float ~data name f =
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-    FWrite.Event.Counter.encode out ~name:"c" ~time_ns ~t_ref:tls.thread_ref
-      ~args:((name, `Float f) :: data)
-      ()
-
-  let counter_int ~data name i =
-    let out, tls = get_thread_output () in
-    let time_ns = Time.now_ns () in
-    FWrite.Event.Counter.encode out ~name:"c" ~time_ns ~t_ref:tls.thread_ref
-      ~args:((name, `Int i) :: data)
-      ()
-
-  let name_process name : unit =
-    let out, _tls = get_thread_output () in
-    FWrite.Kernel_object.(encode out ~name ~ty:ty_process ~kid:pid ~args:[] ())
-
-  let name_thread name : unit =
-    let out, tls = get_thread_output () in
-    FWrite.Kernel_object.(
-      encode out ~name ~ty:ty_thread ~kid:tls.tid
-        ~args:[ "process", `Kid pid ]
-        ())
-
-  let extension_event _ = ()
-end
-
-let create ~out () : collector =
-  let buf_pool = Buf_pool.create () in
-  let events = B_queue.create () in
-
-  let bg_thread =
-    Thread.create (Bg_thread.bg_thread ~buf_pool ~out ~events) ()
-  in
-
-  let st =
-    {
-      active = A.make true;
-      buf_pool;
-      bg_thread;
-      events;
-      span_id_gen = A.make 0;
-      next_thread_ref = A.make 1;
-      per_thread = Array.init 16 (fun _ -> A.make Int_map.empty);
-    }
-  in
-
-  let _tick_thread = Thread.create (fun () -> Bg_thread.tick_thread events) in
-
-  (* write header *)
-  let out = out_of_st st in
-  FWrite.Metadata.Magic_record.encode out;
-  FWrite.Metadata.Initialization_record.(
-    encode out ~ticks_per_secs:default_ticks_per_sec ());
-  FWrite.Metadata.Provider_info.encode out ~id:0 ~name:"ocaml-trace" ();
-  Output.flush out;
-  Output.dispose out;
-
-  let module Coll =
-    C
-      (struct
-        let st = st
-      end)
-      ()
-  in
-  (module Coll)

src/fuchsia/fcollector.mli

@@ -1,3 +0,0 @@
-open Trace_core
-
-val create : out:Bg_thread.out -> unit -> collector

src/fuchsia/lock.ml

@@ -0,0 +1,28 @@
+type 'a t = {
+  mutex: Mutex.t;
+  mutable content: 'a;
+}
+
+let create content : _ t = { mutex = Mutex.create (); content }
+
+let with_ (self : _ t) f =
+  Mutex.lock self.mutex;
+  try
+    let x = f self.content in
+    Mutex.unlock self.mutex;
+    x
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    Mutex.unlock self.mutex;
+    Printexc.raise_with_backtrace e bt
+
+let[@inline] update self f = with_ self (fun x -> self.content <- f x)
+
+let[@inline] update_map l f =
+  with_ l (fun x ->
+      let x', y = f x in
+      l.content <- x';
+      y)
+
+let[@inline] set_while_locked (self : 'a t) (x : 'a) = self.content <- x
+

src/fuchsia/lock.mli

@@ -0,0 +1,10 @@
+type 'a t
+(** A value protected by a mutex *)
+
+val create : 'a -> 'a t
+val with_ : 'a t -> ('a -> 'b) -> 'b
+val update : 'a t -> ('a -> 'a) -> unit
+val update_map : 'a t -> ('a -> 'a * 'b) -> 'b
+
+val set_while_locked : 'a t -> 'a -> unit
+(** Change the value while inside [with_] or similar. *)

src/fuchsia/subscriber.ml

@@ -0,0 +1,172 @@
+open Common_
+open Trace_core
+module Span_tbl = Trace_subscriber.Span_tbl
+
+let on_tracing_error = on_tracing_error
+
+type span_info = {
+  tid: int;
+  name: string;
+  start_ns: int64;
+  mutable data: (string * Sub.user_data) list;
+      (* NOTE: thread safety: this is supposed to only be modified by the thread
+      that's running this (synchronous, stack-abiding) span. *)
+}
+(** Information we store about a span begin event, to emit a complete event when
+    we meet the corresponding span end event *)
+
+type t = {
+  active: bool A.t;
+  pid: int;
+  spans: span_info Span_tbl.t;
+  buf_chain: Buf_chain.t;
+  exporter: Exporter.t;
+}
+(** Subscriber state *)
+
+open struct
+  (** Write the buffers that are ready *)
+  let[@inline] write_ready_ (self : t) =
+    if Buf_chain.has_ready self.buf_chain then
+      Buf_chain.pop_ready self.buf_chain ~f:self.exporter.write_bufs
+
+  let print_non_closed_spans_warning spans =
+    let module Str_set = Set.Make (String) in
+    let spans = Span_tbl.to_list spans in
+    if spans <> [] then (
+      !on_tracing_error
+      @@ Printf.sprintf "warning: %d spans were not closed" (List.length spans);
+      let names =
+        List.fold_left
+          (fun set (_, span) -> Str_set.add span.name set)
+          Str_set.empty spans
+      in
+      Str_set.iter
+        (fun name ->
+          !on_tracing_error @@ Printf.sprintf "  span %S was not closed" name)
+        names;
+      flush stderr
+    )
+end
+
+let close (self : t) : unit =
+  if A.exchange self.active false then (
+    Buf_chain.ready_all_non_empty self.buf_chain;
+    write_ready_ self;
+    self.exporter.close ();
+
+    print_non_closed_spans_warning self.spans
+  )
+
+let[@inline] active self = A.get self.active
+
+let flush (self : t) : unit =
+  Buf_chain.ready_all_non_empty self.buf_chain;
+  write_ready_ self;
+  self.exporter.flush ()
+
+let create ?(buf_pool = Buf_pool.create ()) ~pid ~exporter () : t =
+  let buf_chain = Buf_chain.create ~sharded:true ~buf_pool () in
+  { active = A.make true; buf_chain; exporter; pid; spans = Span_tbl.create () }
+
+module Callbacks = struct
+  type st = t
+
+  let on_init (self : st) ~time_ns:_ =
+    Writer.Metadata.Magic_record.encode self.buf_chain;
+    Writer.Metadata.Initialization_record.(
+      encode self.buf_chain ~ticks_per_secs:default_ticks_per_sec ());
+    Writer.Metadata.Provider_info.encode self.buf_chain ~id:0
+      ~name:"ocaml-trace" ();
+    (* make sure we write these immediately so they're not out of order *)
+    Buf_chain.ready_all_non_empty self.buf_chain;
+
+    write_ready_ self
+
+  let on_shutdown (self : st) ~time_ns:_ = close self
+
+  let on_name_process (self : st) ~time_ns:_ ~tid:_ ~name : unit =
+    Writer.Kernel_object.(
+      encode self.buf_chain ~name ~ty:ty_process ~kid:self.pid ~args:[] ());
+    write_ready_ self
+
+  let on_name_thread (self : st) ~time_ns:_ ~tid ~name : unit =
+    Writer.Kernel_object.(
+      encode self.buf_chain ~name ~ty:ty_thread ~kid:tid
+        ~args:[ "process", A_kid (Int64.of_int self.pid) ]
+        ());
+    write_ready_ self
+
+  (* add function name, if provided, to the metadata *)
+  let add_fun_name_ fun_name data : _ list =
+    match fun_name with
+    | None -> data
+    | Some f -> ("function", Sub.U_string f) :: data
+
+  let[@inline] on_enter_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
+      ~__LINE__:_ ~time_ns ~tid ~data ~name span : unit =
+    let data = add_fun_name_ fun_name data in
+    let info = { tid; name; start_ns = time_ns; data } in
+    (* save the span so we find it at exit *)
+    Span_tbl.add self.spans span info
+
+  let on_exit_span (self : st) ~time_ns ~tid:_ span : unit =
+    match Span_tbl.find_exn self.spans span with
+    | exception Not_found ->
+      !on_tracing_error (Printf.sprintf "cannot find span %Ld" span)
+    | { tid; name; start_ns; data } ->
+      Span_tbl.remove self.spans span;
+      Writer.(
+        Event.Duration_complete.encode self.buf_chain ~name
+          ~t_ref:(Thread_ref.inline ~pid:self.pid ~tid)
+          ~time_ns:start_ns ~end_time_ns:time_ns ~args:(args_of_user_data data)
+          ());
+      write_ready_ self
+
+  let on_add_data (self : st) ~data span =
+    if data <> [] then (
+      try
+        let info = Span_tbl.find_exn self.spans span in
+        info.data <- List.rev_append data info.data
+      with Not_found ->
+        !on_tracing_error (Printf.sprintf "cannot find span %Ld" span)
+    )
+
+  let on_message (self : st) ~time_ns ~tid ~span:_ ~data msg : unit =
+    Writer.(
+      Event.Instant.encode self.buf_chain
+        ~t_ref:(Thread_ref.inline ~pid:self.pid ~tid)
+        ~name:msg ~time_ns ~args:(args_of_user_data data) ());
+    write_ready_ self
+
+  let on_counter (self : st) ~time_ns ~tid ~data ~name n : unit =
+    Writer.(
+      Event.Counter.encode self.buf_chain
+        ~t_ref:(Thread_ref.inline ~pid:self.pid ~tid)
+        ~name ~time_ns
+        ~args:((name, A_float n) :: args_of_user_data data)
+        ());
+    write_ready_ self
+
+  let on_enter_manual_span (self : st) ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_
+      ~time_ns ~tid ~parent:_ ~data ~name ~flavor:_ ~trace_id _span : unit =
+    Writer.(
+      Event.Async_begin.encode self.buf_chain ~name
+        ~args:(args_of_user_data data)
+        ~t_ref:(Thread_ref.inline ~pid:self.pid ~tid)
+        ~time_ns ~async_id:trace_id ());
+    write_ready_ self
+
+  let on_exit_manual_span (self : st) ~time_ns ~tid ~name ~data ~flavor:_
+      ~trace_id (_ : span) : unit =
+    Writer.(
+      Event.Async_end.encode self.buf_chain ~name ~args:(args_of_user_data data)
+        ~t_ref:(Thread_ref.inline ~pid:self.pid ~tid)
+        ~time_ns ~async_id:trace_id ());
+    write_ready_ self
+
+  let on_extension_event _ ~time_ns:_ ~tid:_ _ev = ()
+end
+
+let subscriber (self : t) : Sub.t =
+  Sub.Subscriber.Sub { st = self; callbacks = (module Callbacks) }

src/fuchsia/subscriber.mli

@@ -0,0 +1,20 @@
+type t
+(** Main subscriber state. *)
+
+val create : ?buf_pool:Buf_pool.t -> pid:int -> exporter:Exporter.t -> unit -> t
+(** Create a subscriber state. *)
+
+val flush : t -> unit
+val close : t -> unit
+val active : t -> bool
+
+module Callbacks : Trace_subscriber.Callbacks.S with type st = t
+
+val subscriber : t -> Trace_subscriber.t
+(** Subscriber that writes json into this writer *)
+
+(**/**)
+
+val on_tracing_error : (string -> unit) ref
+
+(**/**)

src/fuchsia/trace_fuchsia.ml

@@ -1,31 +1,50 @@
 open Common_
+module Buf = Buf
+module Buf_chain = Buf_chain
+module Buf_pool = Buf_pool
+module Exporter = Exporter
+module Subscriber = Subscriber
+module Writer = Writer
 
 type output =
-  [ `Stdout
+  [ `File of string
-  | `Stderr
+  | `Exporter of Exporter.t
-  | `File of string
   ]
 
-let collector = Fcollector.create
+let get_out_ (out : [< output ]) : Exporter.t =
+  match out with
+  | `File path ->
+    let oc = open_out path in
+    Exporter.of_out_channel ~close_channel:true oc
+  | `Exporter e -> e
+
+let subscriber ~out () : Sub.t =
+  let exporter = get_out_ out in
+  let pid =
+    if !Trace_subscriber.Private_.mock then
+      2
+    else
+      Unix.getpid ()
+  in
+  let sub = Subscriber.create ~pid ~exporter () in
+  Subscriber.subscriber sub
+
+let collector ~out () = Sub.collector @@ subscriber ~out ()
 
 let setup ?(out = `Env) () =
   match out with
-  | `Stderr -> Trace_core.setup_collector @@ Fcollector.create ~out:`Stderr ()
+  | `File path -> Trace_core.setup_collector @@ collector ~out:(`File path) ()
-  | `Stdout -> Trace_core.setup_collector @@ Fcollector.create ~out:`Stdout ()
+  | `Exporter _ as out ->
-  | `File path ->
+    let sub = subscriber ~out () in
-    Trace_core.setup_collector @@ Fcollector.create ~out:(`File path) ()
+    Trace_core.setup_collector @@ Sub.collector sub
   | `Env ->
     (match Sys.getenv_opt "TRACE" with
     | Some ("1" | "true") ->
       let path = "trace.fxt" in
-      let c = Fcollector.create ~out:(`File path) () in
+      let c = collector ~out:(`File path) () in
       Trace_core.setup_collector c
-    | Some "stdout" ->
-      Trace_core.setup_collector @@ Fcollector.create ~out:`Stdout ()
-    | Some "stderr" ->
-      Trace_core.setup_collector @@ Fcollector.create ~out:`Stderr ()
     | Some path ->
-      let c = Fcollector.create ~out:(`File path) () in
+      let c = collector ~out:(`File path) () in
       Trace_core.setup_collector c
     | None -> ())
 
@@ -33,6 +52,24 @@ let with_setup ?out () f =
   setup ?out ();
   Fun.protect ~finally:Trace_core.shutdown f
 
+module Mock_ = struct
+  let now = ref 0
+
+  (* used to mock timing *)
+  let get_now_ns () : int64 =
+    let x = !now in
+    incr now;
+    Int64.(mul (of_int x) 1000L)
+
+  let get_tid_ () : int = 3
+end
+
 module Internal_ = struct
+  let mock_all_ () =
+    Sub.Private_.mock := true;
+    Sub.Private_.get_now_ns_ := Mock_.get_now_ns;
+    Sub.Private_.get_tid_ := Mock_.get_tid_;
+    ()
+
   let on_tracing_error = on_tracing_error
 end

src/fuchsia/trace_fuchsia.mli

@@ -6,22 +6,23 @@
      trace format}. This reduces the tracing overhead compared to [trace-tef],
     at the expense of simplicity. *)
 
-val collector :
+module Buf = Buf
-  out:[ `File of string | `Stderr | `Stdout ] -> unit -> Trace_core.collector
+module Buf_chain = Buf_chain
-(** Make a collector that writes into the given output. See {!setup} for more
+module Buf_pool = Buf_pool
-    details. *)
+module Exporter = Exporter
+module Subscriber = Subscriber
+module Writer = Writer
 
 type output =
-  [ `Stdout
+  [ `File of string
-  | `Stderr
+  | `Exporter of Exporter.t
-  | `File of string
   ]
-(** Output for tracing.
 
-    - [`Stdout] will enable tracing and print events on stdout
+val subscriber : out:[< output ] -> unit -> Trace_subscriber.t
-    - [`Stderr] will enable tracing and print events on stderr
+
-    - [`File "foo"] will enable tracing and print events into file named "foo"
+val collector : out:[< output ] -> unit -> Trace_core.collector
-*)
+(** Make a collector that writes into the given output. See {!setup} for more
+    details. *)
 
 val setup : ?out:[ output | `Env ] -> unit -> unit
 (** [setup ()] installs the collector depending on [out].
@@ -32,12 +33,10 @@ val setup : ?out:[ output | `Env ] -> unit -> unit
       - [`Env] will enable tracing if the environment variable "TRACE" is set.
 
     - If it's set to "1", then the file is "trace.fxt".
-    - If it's set to "stdout", then logging happens on stdout (since 0.2)
-    - If it's set to "stderr", then logging happens on stdout (since 0.2)
     - Otherwise, if it's set to a non empty string, the value is taken to be the
       file path into which to write. *)
 
-val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
+val with_setup : ?out:[< output | `Env > `Env ] -> unit -> (unit -> 'a) -> 'a
 (** [with_setup () f] (optionally) sets a collector up, calls [f()], and makes
     sure to shutdown before exiting. *)
 
@@ -45,6 +44,9 @@ val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
 
 module Internal_ : sig
   val on_tracing_error : (string -> unit) ref
+
+  val mock_all_ : unit -> unit
+  (** use fake, deterministic timestamps, TID, PID *)
 end
 
 (**/**)

src/fuchsia/write/buf_pool.ml

@@ -1,58 +0,0 @@
-open struct
-  module A = Trace_core.Internal_.Atomic_
-
-  exception Got_buf of Buf.t
-end
-
-module List_with_len = struct
-  type +'a t =
-    | Nil
-    | Cons of int * 'a * 'a t
-
-  let empty : _ t = Nil
-
-  let[@inline] len = function
-    | Nil -> 0
-    | Cons (i, _, _) -> i
-
-  let[@inline] cons x self = Cons (len self + 1, x, self)
-end
-
-type t = {
-  max_len: int;
-  buf_size: int;
-  bufs: Buf.t List_with_len.t A.t;
-}
-
-let create ?(max_len = 64) ?(buf_size = 1 lsl 16) () : t =
-  let buf_size = min (1 lsl 22) (max buf_size (1 lsl 15)) in
-  { max_len; buf_size; bufs = A.make List_with_len.empty }
-
-let alloc (self : t) : Buf.t =
-  try
-    while
-      match A.get self.bufs with
-      | Nil -> false
-      | Cons (_, buf, tl) as old ->
-        if A.compare_and_set self.bufs old tl then
-          raise (Got_buf buf)
-        else
-          false
-    do
-      ()
-    done;
-    Buf.create self.buf_size
-  with Got_buf b -> b
-
-let recycle (self : t) (buf : Buf.t) : unit =
-  Buf.clear buf;
-  try
-    while
-      match A.get self.bufs with
-      | Cons (i, _, _) when i >= self.max_len -> raise Exit
-      | old ->
-        not (A.compare_and_set self.bufs old (List_with_len.cons buf old))
-    do
-      ()
-    done
-  with Exit -> () (* do not recycle *)

src/fuchsia/write/dune

@@ -1,9 +0,0 @@
-(library
- (name trace_fuchsia_write)
- (public_name trace-fuchsia.write)
- (synopsis "Serialization part of trace-fuchsia")
- (ocamlopt_flags
-  :standard
-  ;-dlambda
-  )
- (libraries trace.core threads))

src/fuchsia/write/output.ml

@@ -1,65 +0,0 @@
-type t = {
-  mutable buf: Buf.t;
-  mutable send_buf: Buf.t -> unit;
-  buf_pool: Buf_pool.t;
-}
-
-let create ~(buf_pool : Buf_pool.t) ~send_buf () : t =
-  let buf_size = buf_pool.buf_size in
-  let buf = Buf.create buf_size in
-  { buf; send_buf; buf_pool }
-
-open struct
-  (* NOTE: there is a potential race condition if an output is
-     flushed from the main thread upon closing, while
-     the local thread is blissfully writing new records to it
-     as we're winding down the collector. This is trying to reduce
-     the likelyhood of a race happening. *)
-  let[@poll error] replace_buf_ (self : t) (new_buf : Buf.t) : Buf.t =
-    let old_buf = self.buf in
-    self.buf <- new_buf;
-    old_buf
-
-  let flush_ (self : t) : unit =
-    let new_buf = Buf_pool.alloc self.buf_pool in
-    let old_buf = replace_buf_ self new_buf in
-    self.send_buf old_buf
-
-  let[@inline never] cycle_buf (self : t) ~available : Buf.t =
-    flush_ self;
-    let buf = self.buf in
-
-    if Buf.available buf < available then (
-      let msg =
-        Printf.sprintf
-          "fuchsia: buffer is too small (available: %d bytes, needed: %d bytes)"
-          (Buf.available buf) available
-      in
-      failwith msg
-    );
-    buf
-end
-
-let[@inline] flush (self : t) : unit = if Buf.size self.buf > 0 then flush_ self
-
-(** Maximum size available, in words, for a single message *)
-let[@inline] max_size_word (self : t) : int = self.buf_pool.buf_size lsr 3
-
-(** Obtain a buffer with at least [available] bytes *)
-let[@inline] get_buf (self : t) ~(available_word : int) : Buf.t =
-  let available = available_word lsl 3 in
-  if Buf.available self.buf >= available then
-    self.buf
-  else
-    cycle_buf self ~available
-
-let into_buffer ~buf_pool (buffer : Buffer.t) : t =
-  let send_buf (buf : Buf.t) =
-    Buffer.add_subbytes buffer buf.buf 0 buf.offset
-  in
-  create ~buf_pool ~send_buf ()
-
-let dispose (self : t) : unit =
-  flush_ self;
-  Buf_pool.recycle self.buf_pool self.buf;
-  self.buf <- Buf.empty

src/fuchsia/writer.ml

@@ -2,14 +2,10 @@
 
     Reference: https://fuchsia.dev/fuchsia-src/reference/tracing/trace-format *)
 
+open Common_
 module Util = Util
-module Buf = Buf
-module Output = Output
-module Buf_pool = Buf_pool
 
 open struct
-  let spf = Printf.sprintf
-
   let[@inline] int64_of_trace_id_ (id : Trace_core.trace_id) : int64 =
     if id == Trace_core.Collector.dummy_trace_id then
       0L
@@ -19,7 +15,27 @@ end
 
 open Util
 
-type user_data = Trace_core.user_data
+type user_data = Sub.user_data =
+  | U_bool of bool
+  | U_float of float
+  | U_int of int
+  | U_none
+  | U_string of string
+
+type arg =
+  | A_bool of bool
+  | A_float of float
+  | A_int of int
+  | A_none
+  | A_string of string
+  | A_kid of int64
+
+(* NOTE: only works because [user_data] is a prefix of [arg] and is immutable *)
+let arg_of_user_data : user_data -> arg = Obj.magic
+
+(* NOTE: only works because [user_data] is a prefix of [arg] and is immutable *)
+let args_of_user_data : (string * user_data) list -> (string * arg) list =
+  Obj.magic
 
 module I64 = struct
   include Int64
@@ -111,8 +127,8 @@ module Metadata = struct
     let value = 0x0016547846040010L
     let size_word = 1
 
-    let encode (out : Output.t) =
+    let encode (bufs : Buf_chain.t) =
-      let buf = Output.get_buf out ~available_word:size_word in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size_word in
       Buf.add_i64 buf value
   end
 
@@ -122,8 +138,8 @@ module Metadata = struct
     (** Default: 1 tick = 1 ns *)
     let default_ticks_per_sec = 1_000_000_000L
 
-    let encode (out : Output.t) ~ticks_per_secs () : unit =
+    let encode (bufs : Buf_chain.t) ~ticks_per_secs () : unit =
-      let buf = Output.get_buf out ~available_word:size_word in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size_word in
       let hd = I64.(1L lor (of_int size_word lsl 4)) in
       Buf.add_i64 buf hd;
       Buf.add_i64 buf ticks_per_secs
@@ -132,10 +148,10 @@ module Metadata = struct
   module Provider_info = struct
     let size_word ~name () = 1 + str_len_word name
 
-    let encode (out : Output.t) ~(id : int) ~name () : unit =
+    let encode (bufs : Buf_chain.t) ~(id : int) ~name () : unit =
       let name = truncate_string name in
       let size = size_word ~name () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
       let hd =
         I64.(
           (of_int size lsl 4)
@@ -152,29 +168,29 @@ module Metadata = struct
 end
 
 module Argument = struct
-  type 'a t = string * ([< user_data | `Kid of int ] as 'a)
+  type t = string * arg
 
-  let check_valid_ : _ t -> unit = function
+  let check_valid_ : t -> unit = function
-    | _, `String s -> assert (String.length s < max_str_len)
+    | _, A_string s -> assert (String.length s < max_str_len)
     | _ -> ()
 
   let[@inline] is_i32_ (i : int) : bool = Int32.(to_int (of_int i) = i)
 
-  let size_word (self : _ t) =
+  let size_word (self : t) =
     let name, data = self in
     match data with
-    | `None | `Bool _ -> 1 + str_len_word name
+    | A_none | A_bool _ -> 1 + str_len_word name
-    | `Int i when is_i32_ i -> 1 + str_len_word name
+    | A_int i when is_i32_ i -> 1 + str_len_word name
-    | `Int _ -> (* int64 *) 2 + str_len_word name
+    | A_int _ -> (* int64 *) 2 + str_len_word name
-    | `Float _ -> 2 + str_len_word name
+    | A_float _ -> 2 + str_len_word name
-    | `String s -> 1 + str_len_word_maybe_too_big s + str_len_word name
+    | A_string s -> 1 + str_len_word_maybe_too_big s + str_len_word name
-    | `Kid _ -> 2 + str_len_word name
+    | A_kid _ -> 2 + str_len_word name
 
   open struct
     external int_of_bool : bool -> int = "%identity"
   end
 
-  let encode (buf : Buf.t) (self : _ t) : unit =
+  let encode (buf : Buf.t) (self : t) : unit =
     let name, data = self in
     let name = truncate_string name in
     let size = size_word self in
@@ -187,26 +203,26 @@ module Argument = struct
     in
 
     match data with
-    | `None ->
+    | A_none ->
       let hd = hd_arg_size in
       Buf.add_i64 buf hd;
       Buf.add_string buf name
-    | `Int i when is_i32_ i ->
+    | A_int i when is_i32_ i ->
       let hd = I64.(1L lor hd_arg_size lor (of_int i lsl 32)) in
       Buf.add_i64 buf hd;
       Buf.add_string buf name
-    | `Int i ->
+    | A_int i ->
       (* int64 *)
       let hd = I64.(3L lor hd_arg_size) in
       Buf.add_i64 buf hd;
       Buf.add_string buf name;
       Buf.add_i64 buf (I64.of_int i)
-    | `Float f ->
+    | A_float f ->
       let hd = I64.(5L lor hd_arg_size) in
       Buf.add_i64 buf hd;
       Buf.add_string buf name;
       Buf.add_i64 buf (I64.bits_of_float f)
-    | `String s ->
+    | A_string s ->
       let s = truncate_string s in
       let hd =
         I64.(
@@ -216,35 +232,35 @@ module Argument = struct
       Buf.add_i64 buf hd;
       Buf.add_string buf name;
       Buf.add_string buf s
-    | `Bool b ->
+    | A_bool b ->
       let hd = I64.(9L lor hd_arg_size lor (of_int (int_of_bool b) lsl 16)) in
       Buf.add_i64 buf hd;
       Buf.add_string buf name
-    | `Kid kid ->
+    | A_kid kid ->
       (* int64 *)
       let hd = I64.(8L lor hd_arg_size) in
       Buf.add_i64 buf hd;
       Buf.add_string buf name;
-      Buf.add_i64 buf (I64.of_int kid)
+      Buf.add_i64 buf kid
 end
 
 module Arguments = struct
-  type 'a t = 'a Argument.t list
+  type t = Argument.t list
 
-  let[@inline] len (self : _ t) : int =
+  let[@inline] len (self : t) : int =
     match self with
     | [] -> 0
     | [ _ ] -> 1
     | _ :: _ :: tl -> 2 + List.length tl
 
-  let check_valid (self : _ t) =
+  let check_valid (self : t) =
     let len = len self in
     if len > 15 then
       invalid_arg (spf "fuchsia: can have at most 15 args, got %d" len);
     List.iter Argument.check_valid_ self;
     ()
 
-  let[@inline] size_word (self : _ t) =
+  let[@inline] size_word (self : t) =
     match self with
     | [] -> 0
     | [ a ] -> Argument.size_word a
@@ -254,7 +270,7 @@ module Arguments = struct
         (Argument.size_word a + Argument.size_word b)
         tl
 
-  let[@inline] encode (buf : Buf.t) (self : _ t) =
+  let[@inline] encode (buf : Buf.t) (self : t) =
     let rec aux buf l =
       match l with
       | [] -> ()
@@ -276,11 +292,11 @@ module Thread_record = struct
   let size_word : int = 3
 
   (** Record that [Thread_ref.ref as_ref] represents the pair [pid, tid] *)
-  let encode (out : Output.t) ~as_ref ~pid ~tid () : unit =
+  let encode (bufs : Buf_chain.t) ~as_ref ~pid ~tid () : unit =
     if as_ref <= 0 || as_ref > 255 then
       invalid_arg "fuchsia: thread_record: invalid ref";
 
-    let buf = Output.get_buf out ~available_word:size_word in
+    let@ buf = Buf_chain.with_buf bufs ~available_word:size_word in
 
     let hd = I64.(3L lor (of_int size_word lsl 4) lor (of_int as_ref lsl 16)) in
     Buf.add_i64 buf hd;
@@ -296,11 +312,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args ()
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args
-        : unit =
+        () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       (* set category = 0 *)
       let hd =
@@ -331,11 +347,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args + 1 (* counter id *)
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args ()
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args
-        : unit =
+        () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       let hd =
         I64.(
@@ -368,11 +384,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args ()
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args
-        : unit =
+        () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       let hd =
         I64.(
@@ -403,11 +419,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args ()
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns ~args
-        : unit =
+        () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       let hd =
         I64.(
@@ -438,11 +454,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args + 1 (* end timestamp *)
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
         ~end_time_ns ~args () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       (* set category = 0 *)
       let hd =
@@ -475,11 +491,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args + 1 (* async id *)
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
         ~(async_id : Trace_core.trace_id) ~args () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       let hd =
         I64.(
@@ -511,11 +527,11 @@ module Event = struct
       1 + Thread_ref.size_word t_ref + 1 (* timestamp *) + str_len_word name
       + Arguments.size_word args + 1 (* async id *)
 
-    let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
+    let encode (bufs : Buf_chain.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
         ~(async_id : Trace_core.trace_id) ~args () : unit =
       let name = truncate_string name in
       let size = size_word ~name ~t_ref ~args () in
-      let buf = Output.get_buf out ~available_word:size in
+      let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
       let hd =
         I64.(
@@ -556,10 +572,11 @@ module Kernel_object = struct
   let ty_process : ty = 1
   let ty_thread : ty = 2
 
-  let encode (out : Output.t) ~name ~(ty : ty) ~(kid : int) ~args () : unit =
+  let encode (bufs : Buf_chain.t) ~name ~(ty : ty) ~(kid : int) ~args () : unit
+      =
     let name = truncate_string name in
     let size = size_word ~name ~args () in
-    let buf = Output.get_buf out ~available_word:size in
+    let@ buf = Buf_chain.with_buf bufs ~available_word:size in
 
     let hd =
       I64.(

src/subscriber/callbacks.ml

@@ -29,16 +29,16 @@ module type S = sig
   type st
   (** Type of the state passed to every callback. *)
 
-  val on_init : st -> time_ns:float -> unit
+  val on_init : st -> time_ns:int64 -> unit
   (** Called when the subscriber is initialized in a collector *)
 
-  val on_shutdown : st -> time_ns:float -> unit
+  val on_shutdown : st -> time_ns:int64 -> unit
   (** Called when the collector is shutdown *)
 
-  val on_name_thread : st -> time_ns:float -> tid:int -> name:string -> unit
+  val on_name_thread : st -> time_ns:int64 -> tid:int -> name:string -> unit
   (** Current thread is being named *)
 
-  val on_name_process : st -> time_ns:float -> tid:int -> name:string -> unit
+  val on_name_process : st -> time_ns:int64 -> tid:int -> name:string -> unit
   (** Current process is being named *)
 
   val on_enter_span :
@@ -46,7 +46,7 @@ module type S = sig
     __FUNCTION__:string option ->
     __FILE__:string ->
     __LINE__:int ->
-    time_ns:float ->
+    time_ns:int64 ->
     tid:int ->
     data:(string * user_data) list ->
     name:string ->
@@ -54,7 +54,7 @@ module type S = sig
     unit
   (** Enter a regular (sync) span *)
 
-  val on_exit_span : st -> time_ns:float -> tid:int -> span -> unit
+  val on_exit_span : st -> time_ns:int64 -> tid:int -> span -> unit
   (** Exit a span. This and [on_enter_span] must follow strict stack discipline
   *)
 
@@ -63,7 +63,7 @@ module type S = sig
 
   val on_message :
     st ->
-    time_ns:float ->
+    time_ns:int64 ->
     tid:int ->
     span:span option ->
     data:(string * user_data) list ->
@@ -73,7 +73,7 @@ module type S = sig
 
   val on_counter :
     st ->
-    time_ns:float ->
+    time_ns:int64 ->
     tid:int ->
     data:(string * user_data) list ->
     name:string ->
@@ -86,7 +86,7 @@ module type S = sig
     __FUNCTION__:string option ->
     __FILE__:string ->
     __LINE__:int ->
-    time_ns:float ->
+    time_ns:int64 ->
     tid:int ->
     parent:span option ->
     data:(string * user_data) list ->
@@ -99,7 +99,7 @@ module type S = sig
 
   val on_exit_manual_span :
     st ->
-    time_ns:float ->
+    time_ns:int64 ->
     tid:int ->
     name:string ->
     data:(string * user_data) list ->
@@ -110,7 +110,7 @@ module type S = sig
   (** Exit a manual span *)
 
   val on_extension_event :
-    st -> time_ns:float -> tid:int -> extension_event -> unit
+    st -> time_ns:int64 -> tid:int -> extension_event -> unit
   (** Extension event
       @since 0.8 *)
 end
@@ -121,7 +121,16 @@ type 'st t = (module S with type st = 'st)
 
 (** Dummy callbacks. It can be useful to reuse some of these functions in a real
     subscriber that doesn't want to handle {b all} events, but only some of
-    them. *)
+    them.
+
+    To write a subscriber that only supports some callbacks, this can be handy:
+    {[
+      module My_callbacks = struct
+      type st = my_own_state
+      include Callbacks.Dummy
+      let on_counter (st:st) ~time_ns ~tid ~data ~name v : unit = ...
+      end
+    ]} *)
 module Dummy = struct
   let on_init _ ~time_ns:_ = ()
   let on_shutdown _ ~time_ns:_ = ()

src/subscriber/dune

@@ -1,6 +1,7 @@
 (library
  (name trace_subscriber)
  (public_name trace.subscriber)
+ (private_modules time_ thread_ tbl_)
  (libraries
   (re_export trace.core)
   (select
@@ -8,6 +9,12 @@
    from
    (threads -> thread_.real.ml)
    (-> thread_.dummy.ml))
+  (select
+   tbl_.ml
+   from
+   (picos_aux.htbl -> tbl_.picos.ml)
+   (threads -> tbl_.thread.ml)
+   (-> tbl_.basic.ml))
   (select
    time_.ml
    from

src/subscriber/span_tbl.ml

@@ -0,0 +1 @@
+include Tbl_

src/subscriber/span_tbl.mli

@@ -0,0 +1,21 @@
+(** A table that can be used to remember information about spans.
+
+    This is convenient when we want to rememner information from a span begin,
+    when dealing with the corresponding span end.
+
+    {b NOTE}: this is thread safe when threads are enabled. *)
+
+open Trace_core
+
+type 'v t
+
+val create : unit -> 'v t
+val add : 'v t -> span -> 'v -> unit
+
+val find_exn : 'v t -> span -> 'v
+(** @raise Not_found if information isn't found *)
+
+val remove : _ t -> span -> unit
+(** Remove the span if present *)
+
+val to_list : 'v t -> (span * 'v) list

src/subscriber/tbl_.basic.ml

@@ -0,0 +1,13 @@
+module T = Hashtbl.Make (struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end)
+
+type 'v t = 'v T.t
+
+let create () : _ t = T.create 32
+let find_exn = T.find
+let remove = T.remove
+let add = T.replace
+let to_list self : _ list = T.fold (fun k v l -> (k, v) :: l) self []

src/subscriber/tbl_.mli

@@ -0,0 +1,7 @@
+type 'v t
+
+val create : unit -> 'v t
+val add : 'v t -> int64 -> 'v -> unit
+val find_exn : 'v t -> int64 -> 'v
+val remove : _ t -> int64 -> unit
+val to_list : 'v t -> (int64 * 'v) list

src/subscriber/tbl_.picos.ml

@@ -0,0 +1,18 @@
+module H = Picos_aux_htbl
+
+module Key = struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end
+
+type 'v t = (int64, 'v) H.t
+
+let create () : _ t = H.create ~hashed_type:(module Key) ()
+let find_exn = H.find_exn
+let[@inline] remove self k = ignore (H.try_remove self k : bool)
+
+let[@inline] add self k v =
+  if not (H.try_add self k v) then ignore (H.try_set self k v)
+
+let[@inline] to_list self = H.to_seq self |> List.of_seq

src/subscriber/tbl_.thread.ml

@@ -0,0 +1,38 @@
+module T = Hashtbl.Make (struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end)
+
+type 'v t = {
+  tbl: 'v T.t;
+  lock: Mutex.t;
+}
+
+let create () : _ t = { tbl = T.create 32; lock = Mutex.create () }
+
+let find_exn self k =
+  Mutex.lock self.lock;
+  try
+    let v = T.find self.tbl k in
+    Mutex.unlock self.lock;
+    v
+  with e ->
+    Mutex.unlock self.lock;
+    raise e
+
+let remove self k =
+  Mutex.lock self.lock;
+  T.remove self.tbl k;
+  Mutex.unlock self.lock
+
+let add self k v =
+  Mutex.lock self.lock;
+  T.replace self.tbl k v;
+  Mutex.unlock self.lock
+
+let to_list self : _ list =
+  Mutex.lock self.lock;
+  let l = T.fold (fun k v l -> (k, v) :: l) self.tbl [] in
+  Mutex.unlock self.lock;
+  l

src/subscriber/time_.dummy.ml

@@ -1 +1 @@
-let[@inline] get_time_ns () : float = 0.
+let[@inline] get_time_ns () : int64 = 0L

src/subscriber/time_.mli

@@ -1 +1 @@
-val get_time_ns : unit -> float
+val get_time_ns : unit -> int64

src/subscriber/time_.mtime.ml

@@ -1,3 +1,3 @@
-let[@inline] get_time_ns () : float =
+let[@inline] get_time_ns () : int64 =
   let t = Mtime_clock.now () in
-  Int64.to_float (Mtime.to_uint64_ns t)
+  Mtime.to_uint64_ns t

src/subscriber/time_.unix.ml

@@ -1,3 +1,3 @@
-let[@inline] get_time_ns () : float =
+let[@inline] get_time_ns () : int64 =
   let t = Unix.gettimeofday () in
-  t *. 1e9
+  Int64.of_float (t *. 1e9)

src/subscriber/trace_subscriber.ml

@@ -1,24 +1,28 @@
 open Trace_core
 module Callbacks = Callbacks
 module Subscriber = Subscriber
+module Span_tbl = Span_tbl
 include Types
 
 type t = Subscriber.t
 
 module Private_ = struct
-  let get_now_ns_ = ref None
+  let mock = ref false
-  let get_tid_ = ref None
+  let get_now_ns_ = ref Time_.get_time_ns
+  let get_tid_ = ref Thread_.get_tid
 
   (** Now, in nanoseconds *)
-  let[@inline] now_ns () : float =
+  let[@inline] now_ns () : int64 =
-    match !get_now_ns_ with
+    if !mock then
-    | Some f -> f ()
+      !get_now_ns_ ()
-    | None -> Time_.get_time_ns ()
+    else
+      Time_.get_time_ns ()
 
   let[@inline] tid_ () : int =
-    match !get_tid_ with
+    if !mock then
-    | Some f -> f ()
+      !get_tid_ ()
-    | None -> Thread_.get_tid ()
+    else
+      Thread_.get_tid ()
 end
 
 open struct

src/subscriber/trace_subscriber.mli

@@ -4,10 +4,14 @@
     trace event. It also defines a collector that needs to be installed for the
     subscriber(s) to be called.
 
+    Thanks to {!Subscriber.tee_l} it's possible to combine multiple subscribers
+    into a single collector.
+
     @since 0.8 *)
 
 module Callbacks = Callbacks
 module Subscriber = Subscriber
+module Span_tbl = Span_tbl
 
 include module type of struct
   include Types
@@ -31,13 +35,17 @@ val collector : t -> Trace_core.collector
 (**/**)
 
 module Private_ : sig
-  val get_now_ns_ : (unit -> float) option ref
+  val mock : bool ref
+  (** Global mock flag. If enable, all timestamps, tid, etc should be faked. *)
+
+  val get_now_ns_ : (unit -> int64) ref
   (** The callback used to get the current timestamp *)
 
-  val get_tid_ : (unit -> int) option ref
+  val get_tid_ : (unit -> int) ref
   (** The callback used to get the current thread's id *)
 
-  val now_ns : unit -> float
+  val now_ns : unit -> int64
+  (** Get the current timestamp, or a mock version *)
 end
 
 (**/**)

src/tef-tldrs/trace_tef_tldrs.ml

@@ -22,7 +22,7 @@ let get_unix_socket () =
 
 type as_client = {
   trace_id: string;
-  socket: string;
+  socket: string;  (** Unix socket address *)
   emit_tef_at_exit: string option;
       (** For parent, ask daemon to emit traces here *)
 }

src/tef/common_.ml

@@ -0,0 +1,4 @@
+module Sub = Trace_subscriber
+module A = Trace_core.Internal_.Atomic_
+
+let ( let@ ) = ( @@ )

src/tef/exporter.ml

@@ -0,0 +1,85 @@
+(** An exporter, takes JSON objects and writes them somewhere *)
+
+open Common_
+
+type t = {
+  on_json: Buffer.t -> unit;
+      (** Takes a buffer and writes it somewhere. The buffer is only valid
+          during this call and must not be stored. *)
+  flush: unit -> unit;  (** Force write *)
+  close: unit -> unit;  (** Close underlying resources *)
+}
+(** An exporter, takes JSON objects and writes them somewhere.
+
+    This should be thread-safe if used in a threaded environment. *)
+
+open struct
+  let with_lock lock f =
+    Mutex.lock lock;
+    try
+      let res = f () in
+      Mutex.unlock lock;
+      res
+    with e ->
+      let bt = Printexc.get_raw_backtrace () in
+      Mutex.unlock lock;
+      Printexc.raise_with_backtrace e bt
+end
+
+(** Export to the channel
+    @param jsonl
+      if true, export as a JSON object per line, otherwise export as a single
+      big JSON array.
+    @param close_channel if true, closing the exporter will close the channel *)
+let of_out_channel ~close_channel ~jsonl oc : t =
+  let lock = Mutex.create () in
+  let first = ref true in
+  let closed = ref false in
+  let flush () =
+    let@ () = with_lock lock in
+    flush oc
+  in
+  let close () =
+    let@ () = with_lock lock in
+    if not !closed then (
+      closed := true;
+      if not jsonl then output_char oc ']';
+      if close_channel then close_out_noerr oc
+    )
+  in
+  let on_json buf =
+    let@ () = with_lock lock in
+    if not jsonl then
+      if !first then (
+        if not jsonl then output_char oc '[';
+        first := false
+      ) else
+        output_string oc ",\n";
+    Buffer.output_buffer oc buf;
+    if jsonl then output_char oc '\n'
+  in
+  { flush; close; on_json }
+
+let of_buffer ~jsonl (buf : Buffer.t) : t =
+  let lock = Mutex.create () in
+  let first = ref true in
+  let closed = ref false in
+  let close () =
+    let@ () = with_lock lock in
+    if not !closed then (
+      closed := true;
+      if not jsonl then Buffer.add_char buf ']'
+    )
+  in
+  let on_json json =
+    let@ () = with_lock lock in
+    if not jsonl then
+      if !first then (
+        if not jsonl then Buffer.add_char buf '[';
+        first := false
+      ) else
+        Buffer.add_string buf ",\n";
+    Buffer.add_buffer buf json;
+    if jsonl then Buffer.add_char buf '\n'
+  in
+  { flush = ignore; close; on_json }

src/tef/subscriber.ml

@@ -0,0 +1,175 @@
+open Common_
+open Trace_core
+open Trace_private_util
+module Span_tbl = Sub.Span_tbl
+
+module Buf_pool = struct
+  type t = Buffer.t Rpool.t
+
+  let create ?(max_size = 32) ?(buf_size = 256) () : t =
+    Rpool.create ~max_size ~clear:Buffer.reset
+      ~create:(fun () -> Buffer.create buf_size)
+      ()
+end
+
+open struct
+  let[@inline] time_us_of_time_ns (t : int64) : float =
+    Int64.div t 1_000L |> Int64.to_float
+
+  let[@inline] int64_of_trace_id_ (id : Trace_core.trace_id) : int64 =
+    if id == Trace_core.Collector.dummy_trace_id then
+      0L
+    else
+      Bytes.get_int64_le (Bytes.unsafe_of_string id) 0
+end
+
+let on_tracing_error = ref (fun s -> Printf.eprintf "%s\n%!" s)
+
+type span_info = {
+  tid: int;
+  name: string;
+  start_us: float;
+  mutable data: (string * Sub.user_data) list;
+      (* NOTE: thread safety: this is supposed to only be modified by the thread
+that's running this (synchronous, stack-abiding) span. *)
+}
+(** Information we store about a span begin event, to emit a complete event when
+    we meet the corresponding span end event *)
+
+type t = {
+  active: bool A.t;
+  pid: int;
+  spans: span_info Span_tbl.t;
+  buf_pool: Buf_pool.t;
+  exporter: Exporter.t;
+}
+(** Subscriber state *)
+
+open struct
+  let print_non_closed_spans_warning spans =
+    let module Str_set = Set.Make (String) in
+    let spans = Span_tbl.to_list spans in
+    if spans <> [] then (
+      !on_tracing_error
+      @@ Printf.sprintf "trace-tef: warning: %d spans were not closed"
+           (List.length spans);
+      let names =
+        List.fold_left
+          (fun set (_, span) -> Str_set.add span.name set)
+          Str_set.empty spans
+      in
+      Str_set.iter
+        (fun name ->
+          !on_tracing_error @@ Printf.sprintf "  span %S was not closed" name)
+        names;
+      flush stderr
+    )
+end
+
+let close (self : t) : unit =
+  if A.exchange self.active false then (
+    print_non_closed_spans_warning self.spans;
+    self.exporter.close ()
+  )
+
+let[@inline] active self = A.get self.active
+let[@inline] flush (self : t) : unit = self.exporter.flush ()
+
+let create ?(buf_pool = Buf_pool.create ()) ~pid ~exporter () : t =
+  { active = A.make true; exporter; buf_pool; pid; spans = Span_tbl.create () }
+
+module Callbacks = struct
+  type st = t
+
+  let on_init _ ~time_ns:_ = ()
+  let on_shutdown (self : st) ~time_ns:_ = close self
+
+  let on_name_process (self : st) ~time_ns:_ ~tid:_ ~name : unit =
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_name_process ~pid:self.pid ~name buf;
+    self.exporter.on_json buf
+
+  let on_name_thread (self : st) ~time_ns:_ ~tid ~name : unit =
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_name_thread buf ~pid:self.pid ~tid ~name;
+    self.exporter.on_json buf
+
+  (* add function name, if provided, to the metadata *)
+  let add_fun_name_ fun_name data : _ list =
+    match fun_name with
+    | None -> data
+    | Some f -> ("function", Sub.U_string f) :: data
+
+  let[@inline] on_enter_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
+      ~__LINE__:_ ~time_ns ~tid ~data ~name span : unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+    let data = add_fun_name_ fun_name data in
+    let info = { tid; name; start_us = time_us; data } in
+    (* save the span so we find it at exit *)
+    Span_tbl.add self.spans span info
+
+  let on_exit_span (self : st) ~time_ns ~tid:_ span : unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+
+    match Span_tbl.find_exn self.spans span with
+    | exception Not_found ->
+      !on_tracing_error
+        (Printf.sprintf "trace-tef: error: cannot find span %Ld" span)
+    | { tid; name; start_us; data } ->
+      Span_tbl.remove self.spans span;
+      let@ buf = Rpool.with_ self.buf_pool in
+      Writer.emit_duration_event buf ~pid:self.pid ~tid ~name ~start:start_us
+        ~end_:time_us ~args:data;
+
+      self.exporter.on_json buf
+
+  let on_add_data (self : st) ~data span =
+    if data <> [] then (
+      try
+        let info = Span_tbl.find_exn self.spans span in
+        info.data <- List.rev_append data info.data
+      with Not_found ->
+        !on_tracing_error
+          (Printf.sprintf "trace-tef: error: cannot find span %Ld" span)
+    )
+
+  let on_message (self : st) ~time_ns ~tid ~span:_ ~data msg : unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_instant_event buf ~pid:self.pid ~tid ~name:msg ~ts:time_us
+      ~args:data;
+    self.exporter.on_json buf
+
+  let on_counter (self : st) ~time_ns ~tid ~data:_ ~name n : unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_counter buf ~pid:self.pid ~name ~tid ~ts:time_us n;
+    self.exporter.on_json buf
+
+  let on_enter_manual_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
+      ~__LINE__:_ ~time_ns ~tid ~parent:_ ~data ~name ~flavor ~trace_id _span :
+      unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+
+    let data = add_fun_name_ fun_name data in
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_manual_begin buf ~pid:self.pid ~tid ~name
+      ~id:(int64_of_trace_id_ trace_id)
+      ~ts:time_us ~args:data ~flavor;
+    self.exporter.on_json buf
+
+  let on_exit_manual_span (self : st) ~time_ns ~tid ~name ~data ~flavor
+      ~trace_id (_ : span) : unit =
+    let time_us = time_us_of_time_ns @@ time_ns in
+
+    let@ buf = Rpool.with_ self.buf_pool in
+    Writer.emit_manual_end buf ~pid:self.pid ~tid ~name
+      ~id:(int64_of_trace_id_ trace_id)
+      ~ts:time_us ~flavor ~args:data;
+    self.exporter.on_json buf
+
+  let on_extension_event _ ~time_ns:_ ~tid:_ _ev = ()
+end
+
+let subscriber (self : t) : Sub.t =
+  Sub.Subscriber.Sub { st = self; callbacks = (module Callbacks) }

src/tef/subscriber.mli

@@ -0,0 +1,28 @@
+open Common_
+
+module Buf_pool : sig
+  type t
+
+  val create : ?max_size:int -> ?buf_size:int -> unit -> t
+end
+
+type t
+(** Main subscriber state. *)
+
+val create : ?buf_pool:Buf_pool.t -> pid:int -> exporter:Exporter.t -> unit -> t
+(** Create a subscriber state. *)
+
+val flush : t -> unit
+val close : t -> unit
+val active : t -> bool
+
+module Callbacks : Sub.Callbacks.S with type st = t
+
+val subscriber : t -> Sub.t
+(** Subscriber that writes json into this writer *)
+
+(**/**)
+
+val on_tracing_error : (string -> unit) ref
+
+(**/**)

src/tef/trace_tef.ml

@@ -1,214 +1,7 @@
 open Trace_core
-open Trace_private_util
+module Subscriber = Subscriber
-open Event
+module Exporter = Exporter
-module Sub = Trace_subscriber
+module Writer = Writer
-module A = Trace_core.Internal_.Atomic_
-
-let on_tracing_error = ref (fun s -> Printf.eprintf "trace-tef error: %s\n%!" s)
-
-let[@inline] int64_of_trace_id_ (id : Trace_core.trace_id) : int64 =
-  if id == Trace_core.Collector.dummy_trace_id then
-    0L
-  else
-    Bytes.get_int64_le (Bytes.unsafe_of_string id) 0
-
-module Mock_ = struct
-  let enabled = ref false
-  let now = ref 0
-
-  (* used to mock timing *)
-  let get_now_ns () : float =
-    let x = !now in
-    incr now;
-    float_of_int x *. 1000.
-
-  let get_tid_ () : int = 3
-end
-
-module Span_tbl = Hashtbl.Make (struct
-  include Int64
-
-  let hash : t -> int = Hashtbl.hash
-end)
-
-type span_info = {
-  tid: int;
-  name: string;
-  start_us: float;
-  mutable data: (string * Sub.user_data) list;
-}
-
-(** Writer: knows how to write entries to a file in TEF format *)
-module Writer = struct
-  type t = {
-    oc: out_channel;
-    jsonl: bool;  (** JSONL mode, one json event per line *)
-    mutable first: bool;  (** first event? useful in json mode *)
-    buf: Buffer.t;  (** Buffer to write into *)
-    must_close: bool;  (** Do we have to close the underlying channel [oc]? *)
-    pid: int;
-  }
-  (** A writer to a [out_channel]. It writes JSON entries in an array and closes
-      the array at the end. *)
-
-  let create ~(mode : [ `Single | `Jsonl ]) ~out () : t =
-    let jsonl = mode = `Jsonl in
-    let oc, must_close =
-      match out with
-      | `Stdout -> stdout, false
-      | `Stderr -> stderr, false
-      | `File path -> open_out path, true
-      | `File_append path ->
-        open_out_gen [ Open_creat; Open_wronly; Open_append ] 0o644 path, true
-      | `Output oc -> oc, false
-    in
-    let pid =
-      if !Mock_.enabled then
-        2
-      else
-        Unix.getpid ()
-    in
-    if not jsonl then output_char oc '[';
-    { oc; jsonl; first = true; pid; must_close; buf = Buffer.create 2_048 }
-
-  let close (self : t) : unit =
-    if self.jsonl then
-      output_char self.oc '\n'
-    else
-      output_char self.oc ']';
-    flush self.oc;
-    if self.must_close then close_out self.oc
-
-  let with_ ~mode ~out f =
-    let writer = create ~mode ~out () in
-    Fun.protect ~finally:(fun () -> close writer) (fun () -> f writer)
-
-  let[@inline] flush (self : t) : unit = flush self.oc
-
-  (** Emit "," if we need, and get the buffer ready *)
-  let emit_sep_and_start_ (self : t) =
-    Buffer.reset self.buf;
-    if self.jsonl then
-      Buffer.add_char self.buf '\n'
-    else if self.first then
-      self.first <- false
-    else
-      Buffer.add_string self.buf ",\n"
-
-  let char = Buffer.add_char
-  let raw_string = Buffer.add_string
-
-  let str_val (buf : Buffer.t) (s : string) =
-    char buf '"';
-    let encode_char c =
-      match c with
-      | '"' -> raw_string buf {|\"|}
-      | '\\' -> raw_string buf {|\\|}
-      | '\n' -> raw_string buf {|\n|}
-      | '\b' -> raw_string buf {|\b|}
-      | '\r' -> raw_string buf {|\r|}
-      | '\t' -> raw_string buf {|\t|}
-      | _ when Char.code c <= 0x1f ->
-        raw_string buf {|\u00|};
-        Printf.bprintf buf "%02x" (Char.code c)
-      | c -> char buf c
-    in
-    String.iter encode_char s;
-    char buf '"'
-
-  let pp_user_data_ (out : Buffer.t) : Sub.user_data -> unit = function
-    | U_none -> raw_string out "null"
-    | U_int i -> Printf.bprintf out "%d" i
-    | U_bool b -> Printf.bprintf out "%b" b
-    | U_string s -> str_val out s
-    | U_float f -> Printf.bprintf out "%g" f
-
-  (* emit args, if not empty. [ppv] is used to print values. *)
-  let emit_args_o_ ppv (out : Buffer.t) args : unit =
-    if args <> [] then (
-      Printf.bprintf out {json|,"args": {|json};
-      List.iteri
-        (fun i (n, value) ->
-          if i > 0 then raw_string out ",";
-          Printf.bprintf out {json|"%s":%a|json} n ppv value)
-        args;
-      char out '}'
-    )
-
-  let emit_duration_event ~tid ~name ~start ~end_ ~args (self : t) : unit =
-    let dur = end_ -. start in
-    let ts = start in
-
-    emit_sep_and_start_ self;
-
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"cat":"","tid": %d,"dur": %.2f,"ts": %.2f,"name":%a,"ph":"X"%a}|json}
-      self.pid tid dur ts str_val name
-      (emit_args_o_ pp_user_data_)
-      args;
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_manual_begin ~tid ~name ~(id : trace_id) ~ts ~args
-      ~(flavor : Sub.flavor option) (self : t) : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"cat":"trace","id":%Ld,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
-      self.pid (int64_of_trace_id_ id) tid ts str_val name
-      (match flavor with
-      | None | Some Async -> 'b'
-      | Some Sync -> 'B')
-      (emit_args_o_ pp_user_data_)
-      args;
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_manual_end ~tid ~name ~(id : trace_id) ~ts
-      ~(flavor : Sub.flavor option) ~args (self : t) : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"cat":"trace","id":%Ld,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
-      self.pid (int64_of_trace_id_ id) tid ts str_val name
-      (match flavor with
-      | None | Some Async -> 'e'
-      | Some Sync -> 'E')
-      (emit_args_o_ pp_user_data_)
-      args;
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_instant_event ~tid ~name ~ts ~args (self : t) : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"cat":"","tid": %d,"ts": %.2f,"name":%a,"ph":"I"%a}|json}
-      self.pid tid ts str_val name
-      (emit_args_o_ pp_user_data_)
-      args;
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_name_thread ~tid ~name (self : t) : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"tid": %d,"name":"thread_name","ph":"M"%a}|json} self.pid
-      tid
-      (emit_args_o_ pp_user_data_)
-      [ "name", U_string name ];
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_name_process ~name (self : t) : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"name":"process_name","ph":"M"%a}|json} self.pid
-      (emit_args_o_ pp_user_data_)
-      [ "name", U_string name ];
-    Buffer.output_buffer self.oc self.buf
-
-  let emit_counter ~name ~tid ~ts (self : t) f : unit =
-    emit_sep_and_start_ self;
-    Printf.bprintf self.buf
-      {json|{"pid":%d,"tid":%d,"ts":%.2f,"name":"c","ph":"C"%a}|json} self.pid
-      tid ts
-      (emit_args_o_ pp_user_data_)
-      [ name, U_float f ];
-    Buffer.output_buffer self.oc self.buf
-end
 
 let block_signals () =
   try
@@ -226,97 +19,14 @@ let block_signals () =
         : _ list)
   with _ -> ()
 
-let print_non_closed_spans_warning spans =
-  let module Str_set = Set.Make (String) in
-  Printf.eprintf "trace-tef: warning: %d spans were not closed\n"
-    (Span_tbl.length spans);
-  let names = ref Str_set.empty in
-  Span_tbl.iter (fun _ span -> names := Str_set.add span.name !names) spans;
-  Str_set.iter
-    (fun name -> Printf.eprintf "  span %S was not closed\n" name)
-    !names;
-  flush stderr
-
-(** Background thread, takes events from the queue, puts them in context using
-    local state, and writes fully resolved TEF events to [out]. *)
-let bg_thread ~mode ~out (events : Event.t B_queue.t) : unit =
-  block_signals ();
-
-  (* open a writer to [out] *)
-  Writer.with_ ~mode ~out @@ fun writer ->
-  (* local state, to keep track of span information and implicit stack context *)
-  let spans : span_info Span_tbl.t = Span_tbl.create 32 in
-
-  (* add function name, if provided, to the metadata *)
-  let add_fun_name_ fun_name data : _ list =
-    match fun_name with
-    | None -> data
-    | Some f -> ("function", Sub.U_string f) :: data
-  in
-
-  (* how to deal with an event *)
-  let handle_ev (ev : Event.t) : unit =
-    match ev with
-    | E_tick -> Writer.flush writer
-    | E_message { tid; msg; time_us; data } ->
-      Writer.emit_instant_event ~tid ~name:msg ~ts:time_us ~args:data writer
-    | E_define_span { tid; name; id; time_us; fun_name; data } ->
-      let data = add_fun_name_ fun_name data in
-      let info = { tid; name; start_us = time_us; data } in
-      (* save the span so we find it at exit *)
-      Span_tbl.add spans id info
-    | E_exit_span { id; time_us = stop_us } ->
-      (match Span_tbl.find_opt spans id with
-      | None -> !on_tracing_error (Printf.sprintf "cannot find span %Ld" id)
-      | Some { tid; name; start_us; data } ->
-        Span_tbl.remove spans id;
-        Writer.emit_duration_event ~tid ~name ~start:start_us ~end_:stop_us
-          ~args:data writer)
-    | E_add_data { id; data } ->
-      (match Span_tbl.find_opt spans id with
-      | None -> !on_tracing_error (Printf.sprintf "cannot find span %Ld" id)
-      | Some info -> info.data <- List.rev_append data info.data)
-    | E_enter_manual_span { tid; time_us; name; id; data; fun_name; flavor } ->
-      let data = add_fun_name_ fun_name data in
-      Writer.emit_manual_begin ~tid ~name ~id ~ts:time_us ~args:data ~flavor
-        writer
-    | E_exit_manual_span { tid; time_us; name; id; flavor; data } ->
-      Writer.emit_manual_end ~tid ~name ~id ~ts:time_us ~flavor ~args:data
-        writer
-    | E_counter { tid; name; time_us; n } ->
-      Writer.emit_counter ~name ~tid ~ts:time_us writer n
-    | E_name_process { name } -> Writer.emit_name_process ~name writer
-    | E_name_thread { tid; name } -> Writer.emit_name_thread ~tid ~name writer
-  in
-
-  try
-    while true do
-      (* get all the events in the incoming blocking queue, in
-         one single critical section. *)
-      let local = B_queue.pop_all events in
-      List.iter handle_ev local
-    done
-  with B_queue.Closed ->
-    (* write a message about us closing *)
-    Writer.emit_instant_event ~name:"tef-worker.exit"
-      ~tid:(Thread.id @@ Thread.self ())
-      ~ts:(Sub.Private_.now_ns () *. 1e-3)
-      ~args:[] writer;
-
-    (* warn if app didn't close all spans *)
-    if Span_tbl.length spans > 0 then print_non_closed_spans_warning spans;
-    ()
-
 (** Thread that simply regularly "ticks", sending events to the background
     thread so it has a chance to write to the file *)
-let tick_thread events : unit =
+let tick_thread (sub : Subscriber.t) : unit =
   block_signals ();
-  try
+  while Subscriber.active sub do
-    while true do
+    Thread.delay 0.5;
-      Thread.delay 0.5;
+    Subscriber.flush sub
-      B_queue.push events E_tick
+  done
-    done
-  with B_queue.Closed -> ()
 
 type output =
   [ `Stdout
@@ -324,91 +34,45 @@ type output =
   | `File of string
   ]
 
-module Internal_st = struct
+let subscriber_ ~finally ~out ~(mode : [ `Single | `Jsonl ]) () :
-  type t = {
+    Trace_subscriber.t =
-    active: bool A.t;
+  let jsonl = mode = `Jsonl in
-    events: Event.t B_queue.t;
+  let oc, must_close =
-    t_write: Thread.t;
+    match out with
-  }
+    | `Stdout -> stdout, false
-end
+    | `Stderr -> stderr, false
-
+    | `File path -> open_out path, true
-let subscriber_ ~finally ~out ~(mode : [ `Single | `Jsonl ]) () : Sub.t =
+    | `File_append path ->
-  let module M : Sub.Callbacks.S with type st = Internal_st.t = struct
+      open_out_gen [ Open_creat; Open_wronly; Open_append ] 0o644 path, true
-    type st = Internal_st.t
+    | `Output oc -> oc, false
-
+  in
-    let on_init _ ~time_ns:_ = ()
+  let pid =
-
+    if !Trace_subscriber.Private_.mock then
-    let on_shutdown (self : st) ~time_ns:_ =
+      2
-      if A.exchange self.active false then (
+    else
-        B_queue.close self.events;
+      Unix.getpid ()
-        (* wait for writer thread to be done. The writer thread will exit
-           after processing remaining events because the queue is now closed *)
-        Thread.join self.t_write
-      )
-
-    let on_name_process (self : st) ~time_ns:_ ~tid:_ ~name : unit =
-      B_queue.push self.events @@ E_name_process { name }
-
-    let on_name_thread (self : st) ~time_ns:_ ~tid ~name : unit =
-      B_queue.push self.events @@ E_name_thread { tid; name }
-
-    let[@inline] on_enter_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
-        ~__LINE__:_ ~time_ns ~tid ~data ~name span : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events
-      @@ E_define_span { tid; name; time_us; id = span; fun_name; data }
-
-    let on_exit_span (self : st) ~time_ns ~tid:_ span : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events @@ E_exit_span { id = span; time_us }
-
-    let on_add_data (self : st) ~data span =
-      if data <> [] then
-        B_queue.push self.events @@ E_add_data { id = span; data }
-
-    let on_message (self : st) ~time_ns ~tid ~span:_ ~data msg : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events @@ E_message { tid; time_us; msg; data }
-
-    let on_counter (self : st) ~time_ns ~tid ~data:_ ~name f : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events @@ E_counter { name; n = f; time_us; tid }
-
-    let on_enter_manual_span (self : st) ~__FUNCTION__:fun_name ~__FILE__:_
-        ~__LINE__:_ ~time_ns ~tid ~parent:_ ~data ~name ~flavor ~trace_id _span
-        : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events
-      @@ E_enter_manual_span
-           { id = trace_id; time_us; tid; data; name; fun_name; flavor }
-
-    let on_exit_manual_span (self : st) ~time_ns ~tid ~name ~data ~flavor
-        ~trace_id (_ : span) : unit =
-      let time_us = time_ns *. 1e-3 in
-      B_queue.push self.events
-      @@ E_exit_manual_span { tid; id = trace_id; name; time_us; data; flavor }
-
-    let on_extension_event _ ~time_ns:_ ~tid:_ _ev = ()
-  end in
-  let events = B_queue.create () in
-  let t_write =
-    Thread.create
-      (fun () -> Fun.protect ~finally @@ fun () -> bg_thread ~mode ~out events)
-      ()
   in
 
-  (* ticker thread, regularly sends a message to the writer thread.
+  let exporter = Exporter.of_out_channel oc ~jsonl ~close_channel:must_close in
-      no need to join it. *)
+  let exporter =
-  let _t_tick : Thread.t = Thread.create (fun () -> tick_thread events) () in
+    {
-  let st : Internal_st.t = { active = A.make true; events; t_write } in
+      exporter with
-  Sub.Subscriber.Sub { st; callbacks = (module M) }
+      close =
+        (fun () ->
+          exporter.close ();
+          finally ());
+    }
+  in
+  let sub = Subscriber.create ~pid ~exporter () in
+  let _t_tick : Thread.t = Thread.create tick_thread sub in
+  Subscriber.subscriber sub
 
 let collector_ ~(finally : unit -> unit) ~(mode : [ `Single | `Jsonl ]) ~out ()
     : collector =
   let sub = subscriber_ ~finally ~mode ~out () in
-  Sub.collector sub
+  Trace_subscriber.collector sub
 
-let[@inline] subscriber ~out () : Sub.t =
+let[@inline] subscriber ~out () : Trace_subscriber.t =
   subscriber_ ~finally:ignore ~mode:`Single ~out ()
 
 let[@inline] collector ~out () : collector =
@@ -436,14 +100,26 @@ let with_setup ?out () f =
   setup ?out ();
   Fun.protect ~finally:Trace_core.shutdown f
 
+module Mock_ = struct
+  let now = ref 0
+
+  (* used to mock timing *)
+  let get_now_ns () : int64 =
+    let x = !now in
+    incr now;
+    Int64.(mul (of_int x) 1000L)
+
+  let get_tid_ () : int = 3
+end
+
 module Private_ = struct
   let mock_all_ () =
-    Mock_.enabled := true;
+    Trace_subscriber.Private_.mock := true;
-    Sub.Private_.get_now_ns_ := Some Mock_.get_now_ns;
+    Trace_subscriber.Private_.get_now_ns_ := Mock_.get_now_ns;
-    Sub.Private_.get_tid_ := Some Mock_.get_tid_;
+    Trace_subscriber.Private_.get_tid_ := Mock_.get_tid_;
     ()
 
-  let on_tracing_error = on_tracing_error
+  let on_tracing_error = Subscriber.on_tracing_error
 
   let subscriber_jsonl ~finally ~out () =
     subscriber_ ~finally ~mode:`Jsonl ~out ()

src/tef/trace_tef.mli

@@ -1,3 +1,7 @@
+module Subscriber = Subscriber
+module Exporter = Exporter
+module Writer = Writer
+
 type output =
   [ `Stdout
   | `Stderr

src/tef/writer.ml

@@ -0,0 +1,97 @@
+open Common_
+
+let char = Buffer.add_char
+let raw_string = Buffer.add_string
+
+let str_val (buf : Buffer.t) (s : string) =
+  char buf '"';
+  let encode_char c =
+    match c with
+    | '"' -> raw_string buf {|\"|}
+    | '\\' -> raw_string buf {|\\|}
+    | '\n' -> raw_string buf {|\n|}
+    | '\b' -> raw_string buf {|\b|}
+    | '\r' -> raw_string buf {|\r|}
+    | '\t' -> raw_string buf {|\t|}
+    | _ when Char.code c <= 0x1f ->
+      raw_string buf {|\u00|};
+      Printf.bprintf buf "%02x" (Char.code c)
+    | c -> char buf c
+  in
+  String.iter encode_char s;
+  char buf '"'
+
+let pp_user_data_ (out : Buffer.t) : Sub.user_data -> unit = function
+  | U_none -> raw_string out "null"
+  | U_int i -> Printf.bprintf out "%d" i
+  | U_bool b -> Printf.bprintf out "%b" b
+  | U_string s -> str_val out s
+  | U_float f -> Printf.bprintf out "%g" f
+
+(* emit args, if not empty. [ppv] is used to print values. *)
+let emit_args_o_ ppv (out : Buffer.t) args : unit =
+  if args <> [] then (
+    Printf.bprintf out {json|,"args": {|json};
+    List.iteri
+      (fun i (n, value) ->
+        if i > 0 then raw_string out ",";
+        Printf.bprintf out {json|"%s":%a|json} n ppv value)
+      args;
+    char out '}'
+  )
+
+let emit_duration_event ~pid ~tid ~name ~start ~end_ ~args buf : unit =
+  let dur = end_ -. start in
+  let ts = start in
+
+  Printf.bprintf buf
+    {json|{"pid":%d,"cat":"","tid": %d,"dur": %.2f,"ts": %.2f,"name":%a,"ph":"X"%a}|json}
+    pid tid dur ts str_val name
+    (emit_args_o_ pp_user_data_)
+    args
+
+let emit_manual_begin ~pid ~tid ~name ~(id : int64) ~ts ~args
+    ~(flavor : Sub.flavor option) buf : unit =
+  Printf.bprintf buf
+    {json|{"pid":%d,"cat":"trace","id":%Ld,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
+    pid id tid ts str_val name
+    (match flavor with
+    | None | Some Async -> 'b'
+    | Some Sync -> 'B')
+    (emit_args_o_ pp_user_data_)
+    args
+
+let emit_manual_end ~pid ~tid ~name ~(id : int64) ~ts
+    ~(flavor : Sub.flavor option) ~args buf : unit =
+  Printf.bprintf buf
+    {json|{"pid":%d,"cat":"trace","id":%Ld,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
+    pid id tid ts str_val name
+    (match flavor with
+    | None | Some Async -> 'e'
+    | Some Sync -> 'E')
+    (emit_args_o_ pp_user_data_)
+    args
+
+let emit_instant_event ~pid ~tid ~name ~ts ~args buf : unit =
+  Printf.bprintf buf
+    {json|{"pid":%d,"cat":"","tid": %d,"ts": %.2f,"name":%a,"ph":"I"%a}|json}
+    pid tid ts str_val name
+    (emit_args_o_ pp_user_data_)
+    args
+
+let emit_name_thread ~pid ~tid ~name buf : unit =
+  Printf.bprintf buf
+    {json|{"pid":%d,"tid": %d,"name":"thread_name","ph":"M"%a}|json} pid tid
+    (emit_args_o_ pp_user_data_)
+    [ "name", U_string name ]
+
+let emit_name_process ~pid ~name buf : unit =
+  Printf.bprintf buf {json|{"pid":%d,"name":"process_name","ph":"M"%a}|json} pid
+    (emit_args_o_ pp_user_data_)
+    [ "name", U_string name ]
+
+let emit_counter ~pid ~tid ~name ~ts buf f : unit =
+  Printf.bprintf buf
+    {json|{"pid":%d,"tid":%d,"ts":%.2f,"name":"c","ph":"C"%a}|json} pid tid ts
+    (emit_args_o_ pp_user_data_)
+    [ name, U_float f ]

src/tef/writer.mli

@@ -0,0 +1,54 @@
+(** Write JSON events to a buffer.
+
+    This is the part of the code that knows how to emit TEF-compliant JSON from
+    raw event data. *)
+
+open Common_
+open Trace_core
+
+val emit_duration_event :
+  pid:int ->
+  tid:int ->
+  name:string ->
+  start:float ->
+  end_:float ->
+  args:(string * Sub.user_data) list ->
+  Buffer.t ->
+  unit
+
+val emit_manual_begin :
+  pid:int ->
+  tid:int ->
+  name:string ->
+  id:span ->
+  ts:float ->
+  args:(string * Sub.user_data) list ->
+  flavor:Sub.flavor option ->
+  Buffer.t ->
+  unit
+
+val emit_manual_end :
+  pid:int ->
+  tid:int ->
+  name:string ->
+  id:span ->
+  ts:float ->
+  flavor:Sub.flavor option ->
+  args:(string * Sub.user_data) list ->
+  Buffer.t ->
+  unit
+
+val emit_instant_event :
+  pid:int ->
+  tid:int ->
+  name:string ->
+  ts:float ->
+  args:(string * Sub.user_data) list ->
+  Buffer.t ->
+  unit
+
+val emit_name_thread : pid:int -> tid:int -> name:string -> Buffer.t -> unit
+val emit_name_process : pid:int -> name:string -> Buffer.t -> unit
+
+val emit_counter :
+  pid:int -> tid:int -> name:string -> ts:float -> Buffer.t -> float -> unit

src/util/b_queue.ml

@@ -1,65 +0,0 @@
-module A = Trace_core.Internal_.Atomic_
-
-type 'a t = {
-  mutex: Mutex.t;
-  cond: Condition.t;
-  q: 'a Mpsc_bag.t;
-  mutable closed: bool;
-  consumer_waiting: bool A.t;
-}
-
-exception Closed
-
-let create () : _ t =
-  {
-    mutex = Mutex.create ();
-    cond = Condition.create ();
-    q = Mpsc_bag.create ();
-    closed = false;
-    consumer_waiting = A.make false;
-  }
-
-let close (self : _ t) =
-  Mutex.lock self.mutex;
-  if not self.closed then (
-    self.closed <- true;
-    Condition.broadcast self.cond (* awake waiters so they fail  *)
-  );
-  Mutex.unlock self.mutex
-
-let push (self : _ t) x : unit =
-  if self.closed then raise Closed;
-  Mpsc_bag.add self.q x;
-  if self.closed then raise Closed;
-  if A.get self.consumer_waiting then (
-    (* wakeup consumer *)
-    Mutex.lock self.mutex;
-    Condition.broadcast self.cond;
-    Mutex.unlock self.mutex
-  )
-
-let rec pop_all (self : 'a t) : 'a list =
-  match Mpsc_bag.pop_all self.q with
-  | Some l -> l
-  | None ->
-    if self.closed then raise Closed;
-    Mutex.lock self.mutex;
-    A.set self.consumer_waiting true;
-    (* check again, a producer might have pushed an element since we
-       last checked. However if we still find
-       nothing, because this comes after [consumer_waiting:=true],
-       any producer arriving after that will know to wake us up. *)
-    (match Mpsc_bag.pop_all self.q with
-    | Some l ->
-      A.set self.consumer_waiting false;
-      Mutex.unlock self.mutex;
-      l
-    | None ->
-      if self.closed then (
-        Mutex.unlock self.mutex;
-        raise Closed
-      );
-      Condition.wait self.cond self.mutex;
-      A.set self.consumer_waiting false;
-      Mutex.unlock self.mutex;
-      pop_all self)

src/util/b_queue.mli

@@ -1,18 +0,0 @@
-(** Basic Blocking Queue *)
-
-type 'a t
-
-val create : unit -> _ t
-
-exception Closed
-
-val push : 'a t -> 'a -> unit
-(** [push q x] pushes [x] into [q], and returns [()].
-    @raise Closed if [close q] was previously called.*)
-
-val pop_all : 'a t -> 'a list
-(** [pop_all bq] returns all items presently in [bq], in the same order, and
-    clears [bq]. It blocks if no element is in [bq]. *)
-
-val close : _ t -> unit
-(** Close the queue, meaning there won't be any more [push] allowed. *)

src/util/mpsc_bag.ml

@@ -1,32 +0,0 @@
-module A = Trace_core.Internal_.Atomic_
-
-type 'a t = { bag: 'a list A.t } [@@unboxed]
-
-let create () =
-  let bag = A.make [] in
-  { bag }
-
-module Backoff = struct
-  type t = int
-
-  let default = 2
-
-  let once (b : t) : t =
-    for _i = 1 to b do
-      Domain_util.cpu_relax ()
-    done;
-    min (b * 2) 256
-end
-
-let rec add backoff t x =
-  let before = A.get t.bag in
-  let after = x :: before in
-  if not (A.compare_and_set t.bag before after) then
-    add (Backoff.once backoff) t x
-
-let[@inline] add t x = add Backoff.default t x
-
-let[@inline] pop_all t : _ list option =
-  match A.exchange t.bag [] with
-  | [] -> None
-  | l -> Some (List.rev l)

src/util/mpsc_bag.mli

@@ -1,11 +0,0 @@
-(** A multi-producer, single-consumer bag *)
-
-type 'a t
-
-val create : unit -> 'a t
-
-val add : 'a t -> 'a -> unit
-(** [add q x] adds [x] in the bag. *)
-
-val pop_all : 'a t -> 'a list option
-(** Return all current items in the insertion order. *)

src/util/rpool.ml

@@ -0,0 +1,67 @@
+open struct
+  module A = Trace_core.Internal_.Atomic_
+end
+
+module List_with_len = struct
+  type +'a t =
+    | Nil
+    | Cons of int * 'a * 'a t
+
+  let empty : _ t = Nil
+
+  let[@inline] len = function
+    | Nil -> 0
+    | Cons (i, _, _) -> i
+
+  let[@inline] cons x self = Cons (len self + 1, x, self)
+end
+
+type 'a t = {
+  max_size: int;
+  create: unit -> 'a;
+  clear: 'a -> unit;
+  cached: 'a List_with_len.t A.t;
+}
+
+let create ~max_size ~create ~clear () : _ t =
+  { max_size; create; clear; cached = A.make List_with_len.empty }
+
+let alloc (type a) (self : a t) : a =
+  let module M = struct
+    exception Found of a
+  end in
+  try
+    while
+      match A.get self.cached with
+      | Nil -> false
+      | Cons (_, x, tl) as old ->
+        if A.compare_and_set self.cached old tl then
+          raise_notrace (M.Found x)
+        else
+          true
+    do
+      ()
+    done;
+    self.create ()
+  with M.Found x -> x
+
+let recycle (self : 'a t) (x : 'a) : unit =
+  self.clear x;
+  while
+    match A.get self.cached with
+    | Cons (i, _, _) when i >= self.max_size -> false (* drop buf *)
+    | old -> not (A.compare_and_set self.cached old (List_with_len.cons x old))
+  do
+    ()
+  done
+
+let with_ (self : 'a t) f =
+  let x = alloc self in
+  try
+    let res = f x in
+    recycle self x;
+    res
+  with e ->
+    let bt = Printexc.get_raw_backtrace () in
+    recycle self x;
+    Printexc.raise_with_backtrace e bt

src/util/rpool.mli

@@ -0,0 +1,10 @@
+(** A resource pool (for buffers) *)
+
+type 'a t
+
+val create :
+  max_size:int -> create:(unit -> 'a) -> clear:('a -> unit) -> unit -> 'a t
+
+val alloc : 'a t -> 'a
+val recycle : 'a t -> 'a -> unit
+val with_ : 'a t -> ('a -> 'b) -> 'b

test/fuchsia/dune

@@ -0,0 +1,5 @@
+(test
+ (name t1)
+ (package trace-fuchsia)
+ (modules t1)
+ (libraries trace trace-fuchsia))

test/fuchsia/t1.ml

@@ -0,0 +1,75 @@
+let run () =
+  Trace.set_process_name "main";
+  Trace.set_thread_name "t1";
+
+  let n = ref 0 in
+
+  for _i = 1 to 50 do
+    Trace.with_span ~__FILE__ ~__LINE__ "outer.loop" @@ fun _sp ->
+    let pseudo_async_sp =
+      Trace.enter_manual_span ~parent:None ~__FILE__ ~__LINE__ "fake_sleep"
+    in
+
+    for _j = 2 to 5 do
+      incr n;
+      Trace.with_span ~__FILE__ ~__LINE__ "inner.loop" @@ fun _sp ->
+      Trace.messagef (fun k -> k "hello %d %d" _i _j);
+      Trace.message "world";
+      Trace.counter_int "n" !n;
+
+      Trace.add_data_to_span _sp [ "i", `Int _i ];
+
+      if _j = 2 then (
+        Trace.add_data_to_span _sp [ "j", `Int _j ];
+        let _sp =
+          Trace.enter_manual_span
+            ~parent:(Some (Trace.ctx_of_span pseudo_async_sp))
+            ~flavor:
+              (if _i mod 3 = 0 then
+                 `Sync
+               else
+                 `Async)
+            ~__FILE__ ~__LINE__ "sub-sleep"
+        in
+
+        (* fake micro sleep *)
+        Thread.delay 0.005;
+        Trace.exit_manual_span _sp
+      ) else if _j = 3 then (
+        (* pretend some task finished. Note that this is not well scoped wrt other spans. *)
+        Trace.add_data_to_manual_span pseudo_async_sp [ "slept", `Bool true ];
+        Trace.exit_manual_span pseudo_async_sp
+      )
+    done
+  done
+
+let to_hex (s : string) : string =
+  let i_to_hex (i : int) =
+    if i < 10 then
+      Char.chr (i + Char.code '0')
+    else
+      Char.chr (i - 10 + Char.code 'a')
+  in
+
+  let res = Bytes.create (2 * String.length s) in
+  for i = 0 to String.length s - 1 do
+    let n = Char.code (String.get s i) in
+    Bytes.set res (2 * i) (i_to_hex ((n land 0xf0) lsr 4));
+    Bytes.set res ((2 * i) + 1) (i_to_hex (n land 0x0f))
+  done;
+  Bytes.unsafe_to_string res
+
+let () =
+  Trace_fuchsia.Internal_.mock_all_ ();
+  let buf = Buffer.create 32 in
+  let exporter = Trace_fuchsia.Exporter.of_buffer buf in
+  Trace_fuchsia.with_setup ~out:(`Exporter exporter) () run;
+  exporter.close ();
+
+  let data = Buffer.contents buf in
+  (let oc = open_out_bin "t1.fxt" in
+   output_string oc data;
+   close_out_noerr oc);
+
+  print_endline (to_hex data);
+  flush stdout

test/fuchsia/write/dune

@@ -1,4 +1,4 @@
 (tests
  (names t1 t2)
  (package trace-fuchsia)
- (libraries trace-fuchsia.write))
+ (libraries trace-fuchsia))

test/fuchsia/write/t1.ml

@@ -1,4 +1,4 @@
-open Trace_fuchsia_write
+open Trace_fuchsia
 
 module Str_ = struct
   open String
@@ -39,14 +39,14 @@ module Str_ = struct
 end
 
 let () =
-  let l = List.init 100 (fun i -> Util.round_to_word i) in
+  let l = List.init 100 (fun i -> Writer.Util.round_to_word i) in
   assert (List.for_all (fun x -> x mod 8 = 0) l)
 
 let () =
-  assert (Str_ref.inline 0 = 0b0000_0000_0000_0000);
+  assert (Writer.Str_ref.inline 0 = 0b0000_0000_0000_0000);
-  assert (Str_ref.inline 1 = 0b1000_0000_0000_0001);
+  assert (Writer.Str_ref.inline 1 = 0b1000_0000_0000_0001);
-  assert (Str_ref.inline 6 = 0b1000_0000_0000_0110);
+  assert (Writer.Str_ref.inline 6 = 0b1000_0000_0000_0110);
-  assert (Str_ref.inline 31999 = 0b1111_1100_1111_1111);
+  assert (Writer.Str_ref.inline 31999 = 0b1111_1100_1111_1111);
   ()
 
 let () =

test/fuchsia/write/t2.ml

@@ -1,4 +1,5 @@
-open Trace_fuchsia_write
+open Trace_fuchsia
+open Trace_fuchsia.Writer
 
 let pf = Printf.printf
 
@@ -40,24 +41,27 @@ module Str_ = struct
     Bytes.unsafe_to_string res
 end
 
-let with_buf_output (f : Output.t -> unit) : string =
+let with_buf_chain (f : Buf_chain.t -> unit) : string =
   let buf_pool = Buf_pool.create () in
   let buffer = Buffer.create 32 in
-  let out = Output.into_buffer ~buf_pool buffer in
+  let buf_chain = Buf_chain.create ~sharded:true ~buf_pool () in
-  f out;
+  f buf_chain;
-  Output.flush out;
+
+  Buf_chain.ready_all_non_empty buf_chain;
+  let exp = Exporter.of_buffer buffer in
+  Buf_chain.pop_ready buf_chain ~f:exp.write_bufs;
   Buffer.contents buffer
 
 let () = pf "first trace\n"
 
 let () =
   let str =
-    with_buf_output (fun out ->
+    with_buf_chain (fun bufs ->
-        Metadata.Magic_record.encode out;
+        Metadata.Magic_record.encode bufs;
-        Thread_record.encode out ~as_ref:5 ~pid:1 ~tid:86 ();
+        Thread_record.encode bufs ~as_ref:5 ~pid:1 ~tid:86 ();
-        Event.Instant.encode out ~name:"hello" ~time_ns:1234_5678L
+        Event.Instant.encode bufs ~name:"hello" ~time_ns:1234_5678L
           ~t_ref:(Thread_ref.Ref 5)
-          ~args:[ "x", `Int 42 ]
+          ~args:[ "x", A_int 42 ]
           ())
   in
   pf "%s\n" (Str_.to_hex str)
@@ -66,21 +70,21 @@ let () = pf "second trace\n"
 
 let () =
   let str =
-    with_buf_output (fun out ->
+    with_buf_chain (fun bufs ->
-        Metadata.Magic_record.encode out;
+        Metadata.Magic_record.encode bufs;
         Metadata.Initialization_record.(
-          encode out ~ticks_per_secs:default_ticks_per_sec ());
+          encode bufs ~ticks_per_secs:default_ticks_per_sec ());
-        Thread_record.encode out ~as_ref:5 ~pid:1 ~tid:86 ();
+        Thread_record.encode bufs ~as_ref:5 ~pid:1 ~tid:86 ();
-        Metadata.Provider_info.encode out ~id:1 ~name:"ocaml-trace" ();
+        Metadata.Provider_info.encode bufs ~id:1 ~name:"ocaml-trace" ();
-        Event.Duration_complete.encode out ~name:"outer"
+        Event.Duration_complete.encode bufs ~name:"outer"
           ~t_ref:(Thread_ref.Ref 5) ~time_ns:100_000L ~end_time_ns:5_000_000L
           ~args:[] ();
-        Event.Duration_complete.encode out ~name:"inner"
+        Event.Duration_complete.encode bufs ~name:"inner"
           ~t_ref:(Thread_ref.Ref 5) ~time_ns:180_000L ~end_time_ns:4_500_000L
           ~args:[] ();
-        Event.Instant.encode out ~name:"hello" ~time_ns:1_234_567L
+        Event.Instant.encode bufs ~name:"hello" ~time_ns:1_234_567L
           ~t_ref:(Thread_ref.Ref 5)
-          ~args:[ "x", `Int 42 ]
+          ~args:[ "x", A_int 42 ]
           ())
   in
   (let oc = open_out "foo.fxt" in

test/t1.expected

@@ -1049,5 +1049,4 @@
 {"pid":2,"cat":"","tid": 3,"ts": 1299.00,"name":"world","ph":"I"},
 {"pid":2,"tid":3,"ts":1300.00,"name":"c","ph":"C","args": {"n":200}},
 {"pid":2,"cat":"","tid": 3,"dur": 4.00,"ts": 1297.00,"name":"inner.loop","ph":"X","args": {"i":50}},
-{"pid":2,"cat":"","tid": 3,"dur": 25.00,"ts": 1277.00,"name":"outer.loop","ph":"X"},
+{"pid":2,"cat":"","tid": 3,"dur": 25.00,"ts": 1277.00,"name":"outer.loop","ph":"X"}]
-{"pid":2,"cat":"","tid": 1,"ts": 1304.00,"name":"tef-worker.exit","ph":"I"}]

test/t2.expected

@@ -929,5 +929,4 @@
 {"pid":2,"cat":"","tid": 3,"dur": 217.00,"ts": 1642.00,"name":"Dune__exe__T2.fib2","ph":"X"},
 {"pid":2,"cat":"","tid": 3,"dur": 353.00,"ts": 1507.00,"name":"Dune__exe__T2.fib2","ph":"X"},
 {"pid":2,"cat":"","tid": 3,"dur": 573.00,"ts": 1288.00,"name":"Dune__exe__T2.fib2","ph":"X"},
-{"pid":2,"cat":"","tid": 3,"dur": 929.00,"ts": 933.00,"name":"Dune__exe__T2.fib2","ph":"X"},
+{"pid":2,"cat":"","tid": 3,"dur": 929.00,"ts": 933.00,"name":"Dune__exe__T2.fib2","ph":"X"}]
-{"pid":2,"cat":"","tid": 1,"ts": 1864.00,"name":"tef-worker.exit","ph":"I"}]

trace.opam

@@ -17,6 +17,7 @@ depends: [
 depopts: [
   "hmap"
   "unix"
+  "picos_aux" {>= "0.6"}
   "mtime" {>= "2.0"}
 ]
 build: [