Merge pull request #34 from c-cube/simon/string-trace-id-for-async-2025-04-09

breaking: feat(trace): pass a `string` trace_id in manual spans

.github/workflows/main.yml

@@ -48,7 +48,7 @@ jobs:
       - run: opam exec -- dune runtest -p trace-tef,trace-fuchsia
 
       # with depopts
-      - run: opam install hmap unix
+      - run: opam install hmap
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
 
       - run: opam install mtime

src/core/collector.ml

@@ -8,9 +8,13 @@
 open Types
 
 let dummy_span : span = Int64.min_int
+let dummy_trace_id : trace_id = "<dummy>"
 
 let dummy_explicit_span : explicit_span =
-  { span = dummy_span; meta = Meta_map.empty }
+  { span = dummy_span; trace_id = dummy_trace_id; meta = Meta_map.empty }
+
+let dummy_explicit_span_ctx : explicit_span_ctx =
+  { span = dummy_span; trace_id = dummy_trace_id }
 
 (** Signature for a collector.
 
@@ -26,9 +30,6 @@ module type S = sig
     (span -> 'a) ->
     'a
   (** Run the function in a new span.
-
-      This replaces the previous [enter_span] and [exit_span] which were too flexible
-      to be efficient to implement in async contexts.
      @since 0.3 *)
 
   val enter_span :
@@ -49,7 +50,7 @@ module type S = sig
       @since 0.6 *)
 
   val enter_manual_span :
-    parent:explicit_span option ->
+    parent:explicit_span_ctx option ->
     flavor:[ `Sync | `Async ] option ->
     __FUNCTION__:string option ->
     __FILE__:string ->
@@ -61,6 +62,9 @@ module type S = sig
       and this function can store as much metadata as it wants in the hmap
       in the {!explicit_span}'s [meta] field.
 
+      {b NOTE} the [parent] argument is now an {!explicit_span_ctx} and not
+      an {!explicit_span} since NEXT_RELEASE.
+
       This means that the collector doesn't need to implement contextual
       storage mapping {!span} to scopes, metadata, etc. on its side;
       everything can be transmitted in the {!explicit_span}.

src/core/trace_core.ml

@@ -17,6 +17,9 @@ let current_level_ = A.make Level.Trace
 
 (* ## implementation ## *)
 
+let[@inline] ctx_of_span (sp : explicit_span) : explicit_span_ctx =
+  { span = sp.span; trace_id = sp.trace_id }
+
 let data_empty_build_ () = []
 
 let[@inline] enabled () =
@@ -59,27 +62,19 @@ let[@inline] exit_span sp : unit =
   | None -> ()
   | Some (module C) -> C.exit_span sp
 
-let enter_explicit_span_collector_ (module C : Collector.S) ~parent ~flavor
+let enter_manual_span_collector_ (module C : Collector.S) ~parent ~flavor
     ?__FUNCTION__ ~__FILE__ ~__LINE__ ?(data = data_empty_build_) name :
     explicit_span =
   let data = data () in
   C.enter_manual_span ~parent ~flavor ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data
     name
 
-let[@inline] enter_manual_sub_span ~parent ?flavor ?level ?__FUNCTION__
-    ~__FILE__ ~__LINE__ ?data name : explicit_span =
-  match A.get collector with
-  | Some coll when check_level ?level () ->
-    enter_explicit_span_collector_ coll ~parent:(Some parent) ~flavor
-      ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
-  | _ -> Collector.dummy_explicit_span
-
-let[@inline] enter_manual_toplevel_span ?flavor ?level ?__FUNCTION__ ~__FILE__
+let[@inline] enter_manual_span ~parent ?flavor ?level ?__FUNCTION__ ~__FILE__
     ~__LINE__ ?data name : explicit_span =
   match A.get collector with
   | Some coll when check_level ?level () ->
-    enter_explicit_span_collector_ coll ~parent:None ~flavor ?__FUNCTION__
-      ~__FILE__ ~__LINE__ ?data name
+    enter_manual_span_collector_ coll ~parent ~flavor ?__FUNCTION__ ~__FILE__
+      ~__LINE__ ?data name
   | _ -> Collector.dummy_explicit_span
 
 let[@inline] exit_manual_span espan : unit =

src/core/trace_core.mli

@@ -31,6 +31,10 @@ val set_default_level : Level.t -> unit
     default value is [Level.Trace].
     @since 0.7 *)
 
+val ctx_of_span : explicit_span -> explicit_span_ctx
+(** Turn a span into a span context.
+    @since NEXT_RELEASE *)
+
 val with_span :
   ?level:Level.t ->
   ?__FUNCTION__:string ->
@@ -80,8 +84,8 @@ val add_data_to_span : span -> (string * user_data) list -> unit
     Behavior is not specified if the span has been exited.
     @since 0.4 *)
 
-val enter_manual_sub_span :
-  parent:explicit_span ->
+val enter_manual_span :
+  parent:explicit_span_ctx option ->
   ?flavor:[ `Sync | `Async ] ->
   ?level:Level.t ->
   ?__FUNCTION__:string ->
@@ -93,6 +97,12 @@ val enter_manual_sub_span :
 (** Like {!with_span} but the caller is responsible for
     obtaining the [parent] span from their {e own} caller, and carry the resulting
     {!explicit_span} to  the matching {!exit_manual_span}. 
+
+    {b NOTE} this replaces [enter_manual_sub_span] and [enter_manual_toplevel_span]
+    by just making [parent] an explicit option. It is breaking anyway because we now pass
+    an {!explicit_span_ctx} instead of a full {!explicit_span} (the reason being that we
+    might receive this explicit_span_ctx from another process or machine).
+
     @param flavor a description of the span that can be used by the {!Collector.S}
       to decide how to represent the span. Typically, [`Sync] spans
       start and stop on one thread, and are nested purely by their timestamp;
@@ -100,24 +110,7 @@ val enter_manual_sub_span :
       Lwt, Eio, Async, etc.) which impacts how the collector might represent them.
     @param level optional level for this span. since 0.7.
       Default is set via {!set_default_level}.
-    @since 0.3 *)
-
-val enter_manual_toplevel_span :
-  ?flavor:[ `Sync | `Async ] ->
-  ?level:Level.t ->
-  ?__FUNCTION__:string ->
-  __FILE__:string ->
-  __LINE__:int ->
-  ?data:(unit -> (string * user_data) list) ->
-  string ->
-  explicit_span
-(** Like {!with_span} but the caller is responsible for carrying this
-    [explicit_span] around until it's exited with {!exit_manual_span}.
-    The span can be used as a parent in {!enter_manual_sub_span}.
-    @param flavor see {!enter_manual_sub_span} for more details.
-    @param level optional level for this span. since 0.7.
-      Default is set via {!set_default_level}.
-    @since 0.3 *)
+    @since NEXT_RELEASE *)
 
 val exit_manual_span : explicit_span -> unit
 (** Exit an explicit span. This can be on another thread, in a

src/core/types.ml

@@ -3,6 +3,11 @@ type span = int64
 
     The meaning of the identifier depends on the collector. *)
 
+type trace_id = string
+(** A bytestring representing a (possibly distributed) trace made of async spans.
+    With opentelemetry this is 16 bytes.
+    @since NEXT_RELEASE *)
+
 type user_data =
   [ `Int of int
   | `String of string
@@ -13,16 +18,26 @@ type user_data =
 (** User defined data, generally passed as key/value pairs to
     whatever collector is installed (if any). *)
 
+type explicit_span_ctx = {
+  span: span;  (** The current span *)
+  trace_id: trace_id;  (** The trace this belongs to *)
+}
+(** A context, passed around for async traces.
+    @since NEXT_RELEASE *)
+
 type explicit_span = {
   span: span;
       (** Identifier for this span. Several explicit spans might share the same
       identifier since we can differentiate between them via [meta]. *)
+  trace_id: trace_id;  (** The trace this belongs to *)
   mutable meta: Meta_map.t;
       (** Metadata for this span (and its context). This can be used by collectors to
         carry collector-specific information from the beginning
         of the span, to the end of the span. *)
 }
-(** Explicit span, with collector-specific metadata *)
+(** Explicit span, with collector-specific metadata.
+    This is richer than {!explicit_span_ctx} but not intended to be passed around
+    (or sent across the wire), unlike {!explicit_span_ctx}. *)
 
 type extension_event = ..
 (** An extension event, used to add features that are backend specific

src/fuchsia/fcollector.ml

@@ -104,7 +104,6 @@ end = struct
 end
 
 type async_span_info = {
-  async_id: int;
   flavor: [ `Sync | `Async ] option;
   name: string;
   mutable data: (string * user_data) list;
@@ -140,6 +139,12 @@ type state = {
         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 () =
@@ -298,36 +303,33 @@ struct
     | 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 option) ~flavor ~__FUNCTION__:_
-      ~__FILE__:_ ~__LINE__:_ ~data name : explicit_span =
+  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 async_id =
+    let trace_id =
       match parent with
-      | Some m -> (Meta_map.find_exn key_async_data m.meta).async_id
-      | None -> A.fetch_and_add st.span_id_gen 1
+      | 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 ();
+      ~time_ns ~async_id:trace_id ();
     {
       span = 0L;
-      meta =
-        Meta_map.(
-          empty |> add key_async_data { async_id; name; flavor; data = [] });
+      trace_id;
+      meta = Meta_map.(empty |> add key_async_data { name; flavor; data = [] });
     }
 
   let exit_manual_span (es : explicit_span) : unit =
-    let { async_id; name; data; flavor = _ } =
-      Meta_map.find_exn key_async_data es.meta
-    in
+    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 ()
+      ~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

src/fuchsia/write/trace_fuchsia_write.ml

@@ -9,6 +9,9 @@ module Buf_pool = Buf_pool
 
 open struct
   let spf = Printf.sprintf
+
+  let[@inline] int64_of_trace_id_ (id : Trace_core.trace_id) : int64 =
+    Bytes.get_int64_le (Bytes.unsafe_of_string id) 0
 end
 
 open Util
@@ -469,7 +472,7 @@ module Event = struct
       + Arguments.size_word args + 1 (* async id *)
 
     let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
-        ~(async_id : int) ~args () : unit =
+        ~(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
@@ -494,7 +497,7 @@ module Event = struct
 
       Buf.add_string buf name;
       Arguments.encode buf args;
-      Buf.add_i64 buf (I64.of_int async_id);
+      Buf.add_i64 buf (int64_of_trace_id_ async_id);
       ()
   end
 
@@ -505,7 +508,7 @@ module Event = struct
       + Arguments.size_word args + 1 (* async id *)
 
     let encode (out : Output.t) ~name ~(t_ref : Thread_ref.t) ~time_ns
-        ~(async_id : int) ~args () : unit =
+        ~(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
@@ -530,7 +533,7 @@ module Event = struct
 
       Buf.add_string buf name;
       Arguments.encode buf args;
-      Buf.add_i64 buf (I64.of_int async_id);
+      Buf.add_i64 buf (int64_of_trace_id_ async_id);
       ()
   end
 end

src/subscriber/callbacks.ml

@@ -16,6 +16,9 @@
 
     (* … other custom callbacks … *)
   end ]}
+
+  {b NOTE}: the [trace_id] passed alongside manual spans is guaranteed to be at
+  least 64 bits.
 *)
 
 open Trace_core
@@ -88,7 +91,7 @@ module type S = sig
     data:(string * user_data) list ->
     name:string ->
     flavor:flavor option ->
-    trace_id:int ->
+    trace_id:trace_id ->
     span ->
     unit
   (** Enter a manual (possibly async) span *)
@@ -100,7 +103,7 @@ module type S = sig
     name:string ->
     data:(string * user_data) list ->
     flavor:flavor option ->
-    trace_id:int ->
+    trace_id:trace_id ->
     span ->
     unit
   (** Exit a manual span *)

src/subscriber/time_.mtime.ml

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

src/subscriber/trace_subscriber.ml

@@ -33,9 +33,6 @@ open struct
   (** Key used to carry some information between begin and end of
     manual spans, by way of the meta map *)
   let key_manual_info : manual_span_info Meta_map.key = Meta_map.Key.create ()
-
-  (** key used to carry a unique "id" for all spans in an async context *)
-  let key_async_trace_id : int Meta_map.key = Meta_map.Key.create ()
 end
 
 let[@inline] conv_flavor = function
@@ -64,6 +61,12 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
   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
@@ -100,8 +103,8 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
         CB.on_add_data st ~data span
       )
 
-    let enter_manual_span ~(parent : explicit_span option) ~flavor ~__FUNCTION__
-        ~__FILE__ ~__LINE__ ~data name : explicit_span =
+    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 tid = tid_ () in
       let time_ns = now_ns () in
@@ -111,8 +114,8 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
       (* get the common trace id, or make a new one *)
       let trace_id, parent =
         match parent with
-        | Some m -> Meta_map.find_exn key_async_trace_id m.meta, Some m.span
-        | None -> A.fetch_and_add trace_id_gen_ 1, None
+        | Some m -> m.trace_id, Some m.span
+        | None -> mk_trace_id (), None
       in
 
       CB.on_enter_manual_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~parent ~data
@@ -120,18 +123,13 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
       let meta =
         Meta_map.empty
         |> Meta_map.add key_manual_info { name; flavor; data = [] }
-        |> Meta_map.add key_async_trace_id trace_id
       in
-      { span; meta }
+      { span; trace_id; meta }
 
     let exit_manual_span (es : explicit_span) : unit =
       let time_ns = now_ns () in
       let tid = tid_ () in
-      let trace_id =
-        match Meta_map.find key_async_trace_id es.meta with
-        | None -> assert false
-        | Some id -> id
-      in
+      let trace_id = es.trace_id in
       let minfo =
         match Meta_map.find key_manual_info es.meta with
         | None -> assert false

src/tef/event.ml

@@ -30,7 +30,7 @@ type t =
       tid: int;
       name: string;
       time_us: float;
-      id: int;
+      id: trace_id;
       flavor: Sub.flavor option;
       fun_name: string option;
       data: (string * Sub.user_data) list;
@@ -41,7 +41,7 @@ type t =
       time_us: float;
       flavor: Sub.flavor option;
       data: (string * Sub.user_data) list;
-      id: int;
+      id: trace_id;
     }
   | E_counter of {
       name: string;

src/tef/trace_tef.ml

@@ -6,6 +6,9 @@ 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 =
+  Bytes.get_int64_le (Bytes.unsafe_of_string id) 0
+
 module Mock_ = struct
   let enabled = ref false
   let now = ref 0
@@ -142,12 +145,12 @@ module Writer = struct
       args;
     Buffer.output_buffer self.oc self.buf
 
-  let emit_manual_begin ~tid ~name ~id ~ts ~args ~(flavor : Sub.flavor option)
-      (self : t) : unit =
+  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":%d,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
-      self.pid id tid ts str_val name
+      {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')
@@ -155,12 +158,12 @@ module Writer = struct
       args;
     Buffer.output_buffer self.oc self.buf
 
-  let emit_manual_end ~tid ~name ~id ~ts ~(flavor : Sub.flavor option) ~args
-      (self : t) : unit =
+  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":%d,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"%a}|json}
-      self.pid id tid ts str_val name
+      {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')

test/t1.ml

@@ -7,7 +7,7 @@ let run () =
   for _i = 1 to 50 do
     Trace.with_span ~__FILE__ ~__LINE__ "outer.loop" @@ fun _sp ->
     let pseudo_async_sp =
-      Trace.enter_manual_toplevel_span ~__FILE__ ~__LINE__ "fake_sleep"
+      Trace.enter_manual_span ~parent:None ~__FILE__ ~__LINE__ "fake_sleep"
     in
 
     for _j = 2 to 5 do
@@ -22,7 +22,8 @@ let run () =
       if _j = 2 then (
         Trace.add_data_to_span _sp [ "j", `Int _j ];
         let _sp =
-          Trace.enter_manual_sub_span ~parent:pseudo_async_sp
+          Trace.enter_manual_span
+            ~parent:(Some (Trace.ctx_of_span pseudo_async_sp))
             ~flavor:
               (if _i mod 3 = 0 then
                  `Sync