breaking: require subscribers to provide mk_span/mk_trace_id

we want control over this!
This commit is contained in:
Simon Cruanes 2025-12-04 12:31:14 -05:00
parent 15edb582d0
commit aeb2aff3b7
No known key found for this signature in database
GPG key ID: EBFFF6F283F3A2B4
9 changed files with 112 additions and 21 deletions

View file

@ -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 ->

View file

@ -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 *)

View file

@ -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 })

View file

@ -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.(

View file

@ -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:_ = ()

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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