subscriber: restore trace_id

we need it for async backends, clearly

src/subscriber/callbacks.ml

@@ -29,6 +29,9 @@ type 'st t = {
   new_span_id: 'st -> Span_sub.span_id;
       (** How to generate a new span ID?
           @since NEXT_RELEASE *)
+  new_trace_id: 'st -> Span_sub.trace_id;
+      (** How to generate a new trace ID?
+          @since NEXT_RELEASE *)
   on_shutdown: 'st -> time_ns:int64 -> unit;
       (** Called when the collector is shutdown *)
   on_enter_span: 'st -> Span_sub.t -> unit;  (** Enter a span *)
@@ -78,6 +81,7 @@ type 'st t = {
 module Dummy = struct
   let on_init _ ~time_ns:_ = ()
   let new_span_id _ = Int64.min_int
+  let new_trace_id _ = -1L
   let on_shutdown _ ~time_ns:_ = ()
   let on_message _ ~time_ns:_ ~tid:_ ~span:_ ~params:_ ~data:_ _msg = ()
   let on_counter _ ~time_ns:_ ~tid:_ ~params:_ ~data:_ ~name:_ _v = ()
@@ -88,14 +92,16 @@ end
 
 (** Build a set of callbacks.
     @since NEXT_RELEASE *)
-let make ~new_span_id ?(on_init = Dummy.on_init)
+let make ~new_span_id ?(new_trace_id = Dummy.new_trace_id)
-    ?(on_enter_span = Dummy.on_enter_span) ?(on_exit_span = Dummy.on_exit_span)
+    ?(on_init = Dummy.on_init) ?(on_enter_span = Dummy.on_enter_span)
-    ?(on_message = Dummy.on_message) ?(on_counter = Dummy.on_counter)
+    ?(on_exit_span = Dummy.on_exit_span) ?(on_message = Dummy.on_message)
+    ?(on_counter = Dummy.on_counter)
     ?(on_extension_event = Dummy.on_extension_event)
     ?(on_shutdown = Dummy.on_shutdown) () : _ t =
   {
     on_init;
     new_span_id;
+    new_trace_id;
     on_enter_span;
     on_exit_span;
     on_message;

src/subscriber/span_sub.ml

@@ -9,6 +9,9 @@ open Trace_core
 type span_id = int64
 (** Unique ID *)
 
+type trace_id = int64
+(** Unique trace ID *)
+
 let dummy_span_id = Int64.min_int
 
 type span_flavor =
@@ -25,6 +28,7 @@ type t = {
   time_ns: int64;  (** Time the span was entered. *)
   mutable time_exit_ns: int64;
       (** Time the span was exited. Set at exit, [Int64.max_int] otherwise *)
+  trace_id: trace_id;
   tid: int;  (** Thread in which span was created *)
   parent: parent;
   flavor: span_flavor;

src/subscriber/subscriber.ml

@@ -23,6 +23,10 @@ open struct
     let (Sub { st; callbacks = cb }) = Array.get st 0 in
     cb.new_span_id st
 
+  let new_trace_id st =
+    let (Sub { st; callbacks = cb }) = Array.get st 0 in
+    cb.new_trace_id st
+
   let on_init st ~time_ns =
     for i = 0 to Array.length st - 1 do
       let (Sub { st; callbacks = cb }) = Array.get st i in
@@ -69,6 +73,7 @@ open struct
     {
       Callbacks.on_init;
       new_span_id;
+      new_trace_id;
       on_enter_span;
       on_exit_span;
       on_message;

src/subscriber/trace_subscriber.ml

@@ -36,6 +36,12 @@ open struct
     let tid = tid_ () in
     let time_ns = now_ns () in
 
+    let trace_id =
+      match parent with
+      | P_some (Span_sub span) -> span.trace_id
+      | _ -> cb.new_trace_id st
+    in
+
     let flavor = ref `Sync in
     List.iter
       (function
@@ -53,6 +59,7 @@ open struct
         __LINE__;
         data;
         parent;
+        trace_id;
         flavor = !flavor;
         params;
         time_ns;
@@ -125,9 +132,16 @@ end
 let collector (self : Subscriber.t) : collector =
   Collector.C_some (self, coll_cbs)
 
-module Span_generator = struct
+module Span_id_generator = struct
   type t = int A.t
 
   let create () = A.make 0
-  let[@inline] mk_span self = A.fetch_and_add self 1 |> Int64.of_int
+  let[@inline] gen self = A.fetch_and_add self 1 |> Int64.of_int
+end
+
+module Trace_id_generator = struct
+  type t = int A.t
+
+  let create () = A.make 0
+  let[@inline] gen self = A.fetch_and_add self 1 |> Int64.of_int
 end

src/subscriber/trace_subscriber.mli

@@ -29,11 +29,20 @@ val collector : t -> Trace_core.collector
 
 (** A counter-based span generator.
     @since NEXT_RELEASE *)
-module Span_generator : sig
+module Span_id_generator : sig
   type t
 
   val create : unit -> t
-  val mk_span : t -> Span_sub.span_id
+  val gen : t -> Span_sub.span_id
+end
+
+(** A counter-based generator.
+    @since NEXT_RELEASE *)
+module Trace_id_generator : sig
+  type t
+
+  val create : unit -> t
+  val gen : t -> Span_sub.trace_id
 end
 
 (**/**)