breaking: require subscribers to provide mk_span/mk_trace_id

we want control over this!

src/core/trace_core.mli

@@ -84,7 +84,7 @@ val add_data_to_span : span -> (string * user_data) list -> unit
 
 val enter_manual_span :
   parent:explicit_span_ctx option ->
-  ?flavor:[ `Sync | `Async ] ->
+  ?flavor:span_flavor ->
   ?level:Level.t ->
   ?__FUNCTION__:string ->
   __FILE__:string ->

src/core/types.ml

@@ -18,6 +18,13 @@ type user_data =
 (** User defined data, generally passed as key/value pairs to whatever collector
     is installed (if any). *)
 
+type span_flavor =
+  [ `Sync
+  | `Async
+  ]
+(** Some information about the span.
+    @since NEXT_RELEASE *)
+
 type explicit_span_ctx = {
   span: span;  (** The current span *)
   trace_id: trace_id;  (** The trace this belongs to *)

src/event/subscriber.ml

@@ -6,9 +6,18 @@ open Event
 type event_consumer = { on_event: Event.t -> unit } [@@unboxed]
 (** Callback for events. *)
 
+open struct
+  (* just use the same ones for everyone *)
+
+  let span_gen = Sub.Span_generator.create ()
+  let trace_id_gen = Sub.Trace_id_8B_generator.create ()
+end
+
 module Callbacks : Sub.Callbacks.S with type st = event_consumer = struct
   type st = event_consumer
 
+  let new_span (_self : st) = Sub.Span_generator.mk_span span_gen
+  let new_trace_id _self = Sub.Trace_id_8B_generator.mk_trace_id trace_id_gen
   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 })
 

src/fuchsia/subscriber.ml

@@ -21,6 +21,8 @@ type t = {
   spans: span_info Span_tbl.t;
   buf_chain: Buf_chain.t;
   exporter: Exporter.t;
+  span_gen: Sub.Span_generator.t;
+  trace_id_gen: Sub.Trace_id_8B_generator.t;
 }
 (** Subscriber state *)
 
@@ -67,11 +69,24 @@ let flush (self : t) : unit =
 
 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 () }
+  {
+    active = A.make true;
+    buf_chain;
+    exporter;
+    pid;
+    spans = Span_tbl.create ();
+    span_gen = Sub.Span_generator.create ();
+    trace_id_gen = Sub.Trace_id_8B_generator.create ();
+  }
 
 module Callbacks = struct
   type st = t
 
+  let new_span (self : st) = Sub.Span_generator.mk_span self.span_gen
+
+  let new_trace_id self =
+    Sub.Trace_id_8B_generator.mk_trace_id self.trace_id_gen
+
   let on_init (self : st) ~time_ns:_ =
     Writer.Metadata.Magic_record.encode self.buf_chain;
     Writer.Metadata.Initialization_record.(

src/subscriber/callbacks.ml

@@ -32,6 +32,14 @@ module type S = sig
   val on_init : st -> time_ns:int64 -> unit
   (** Called when the subscriber is initialized in a collector *)
 
+  val new_span : st -> span
+  (** How to generate a new span?
+      @since NEXT_RELEASE *)
+
+  val new_trace_id : st -> trace_id
+  (** How to generate a new trace ID?
+      @since NEXT_RELEASE *)
+
   val on_shutdown : st -> time_ns:int64 -> unit
   (** Called when the collector is shutdown *)
 
@@ -133,6 +141,8 @@ type 'st t = (module S with type st = 'st)
     ]} *)
 module Dummy = struct
   let on_init _ ~time_ns:_ = ()
+  let new_span _ = Collector.dummy_span
+  let new_trace_id _ = Collector.dummy_trace_id
   let on_shutdown _ ~time_ns:_ = ()
   let on_name_thread _ ~time_ns:_ ~tid:_ ~name:_ = ()
   let on_name_process _ ~time_ns:_ ~tid:_ ~name:_ = ()

src/subscriber/subscriber.ml

@@ -20,6 +20,14 @@ open struct
   module Tee_cb : Callbacks.S with type st = t array = struct
     type nonrec st = t array
 
+    let new_span st =
+      let (Sub { st = s; callbacks = (module CB) }) = Array.get st 0 in
+      CB.new_span s
+
+    let new_trace_id st =
+      let (Sub { st = s; callbacks = (module CB) }) = Array.get st 0 in
+      CB.new_trace_id s
+
     let on_init st ~time_ns =
       for i = 0 to Array.length st - 1 do
         let (Sub { st = s; callbacks = (module CB) }) = Array.get st i in
@@ -100,8 +108,11 @@ open struct
   end
 end
 
-(** Tee multiple subscribers, ie return a subscriber that forwards to all the
-    subscribers in [subs]. *)
+(** Tee multiple subscribers, ie return a subscriber that forwards to every
+    subscriber in [subs].
+
+    To generate a new span or trace ID, the first subscriber of the list is
+    used. *)
 let tee_l (subs : t list) : t =
   match subs with
   | [] -> dummy

src/subscriber/trace_subscriber.ml

@@ -63,21 +63,8 @@ let rec conv_data = function
 let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
   let open Private_ in
   let module M = struct
-    let trace_id_gen_ = A.make 0
-
-    let[@inline] mk_trace_id () : trace_id =
-      let n = A.fetch_and_add trace_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
-
-    (** generator for span ids *)
-    let new_span_ : unit -> int =
-      let span_id_gen_ = A.make 0 in
-      fun [@inline] () -> A.fetch_and_add span_id_gen_ 1
-
     let enter_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name : span =
-      let span = Int64.of_int (new_span_ ()) in
+      let span = CB.new_span st in
       let tid = tid_ () in
       let time_ns = now_ns () in
       let data = conv_data data in
@@ -109,7 +96,7 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
 
     let enter_manual_span ~(parent : explicit_span_ctx option) ~flavor
         ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name : explicit_span =
-      let span = Int64.of_int (new_span_ ()) in
+      let span = CB.new_span st in
       let tid = tid_ () in
       let time_ns = now_ns () in
       let data = conv_data data in
@@ -119,7 +106,7 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
       let trace_id, parent =
         match parent with
         | Some m -> m.trace_id, Some m.span
-        | None -> mk_trace_id (), None
+        | None -> CB.new_trace_id st, None
       in
 
       CB.on_enter_manual_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~parent ~data
@@ -190,3 +177,22 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
       CB.on_init st ~time_ns
   end in
   (module M)
+
+module Span_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
+end
+
+module Trace_id_8B_generator = struct
+  type t = int A.t
+
+  let create () = A.make 0
+
+  let[@inline] mk_trace_id (self : t) : trace_id =
+    let n = A.fetch_and_add self 1 in
+    let b = Bytes.create 8 in
+    Bytes.set_int64_le b 0 (Int64.of_int n);
+    Bytes.unsafe_to_string b
+end

src/subscriber/trace_subscriber.mli

@@ -32,6 +32,24 @@ val collector : t -> Trace_core.collector
 
     It uses [mtime] (if available) to obtain timestamps. *)
 
+(** A counter-based span generator.
+    @since NEXT_RELEASE *)
+module Span_generator : sig
+  type t
+
+  val create : unit -> t
+  val mk_span : t -> Trace_core.span
+end
+
+(** A counter-based trace ID generator, producing 8-byte trace IDs.
+    @since NEXT_RELEASE *)
+module Trace_id_8B_generator : sig
+  type t
+
+  val create : unit -> t
+  val mk_trace_id : t -> Trace_core.trace_id
+end
+
 (**/**)
 
 module Private_ : sig

src/tef/subscriber.ml

@@ -42,6 +42,8 @@ type t = {
   spans: span_info Span_tbl.t;
   buf_pool: Buf_pool.t;
   exporter: Exporter.t;
+  span_gen: Sub.Span_generator.t;
+  trace_id_gen: Sub.Trace_id_8B_generator.t;
 }
 (** Subscriber state *)
 
@@ -76,11 +78,24 @@ 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 () }
+  {
+    active = A.make true;
+    exporter;
+    buf_pool;
+    pid;
+    spans = Span_tbl.create ();
+    span_gen = Sub.Span_generator.create ();
+    trace_id_gen = Sub.Trace_id_8B_generator.create ();
+  }
 
 module Callbacks = struct
   type st = t
 
+  let new_span (self : st) = Sub.Span_generator.mk_span self.span_gen
+
+  let new_trace_id self =
+    Sub.Trace_id_8B_generator.mk_trace_id self.trace_id_gen
+
   let on_init _ ~time_ns:_ = ()
   let on_shutdown (self : st) ~time_ns:_ = close self