add add_data_to_span and add_data_to_explicit_span

src/core/collector.ml

@@ -53,6 +53,10 @@ module type S = sig
   (** Exit an explicit span.
       @since 0.3 *)
 
+  val add_data_to_span : span -> (string * user_data) list -> unit
+  (** @since Adds data to the given span.
+      NEXT_RELEASE *)
+
   val message : ?span:span -> data:(string * user_data) list -> string -> unit
   (** Emit a message with associated metadata. *)
 

src/core/trace_core.ml

@@ -55,6 +55,16 @@ let[@inline] exit_manual_span espan : unit =
   | None -> ()
   | Some (module C) -> C.exit_manual_span espan
 
+let[@inline] add_data_to_span (span : span) data : unit =
+  if data <> [] then (
+    match A.get collector with
+    | None -> ()
+    | Some (module C) -> C.add_data_to_span span data
+  )
+
+let[@inline] add_data_to_explicit_span esp data : unit =
+  add_data_to_span esp.span data
+
 let message_collector_ (module C : Collector.S) ?span ?(data = fun () -> []) msg
     : unit =
   let data = data () in

src/core/trace_core.mli

@@ -34,6 +34,10 @@ val with_span :
     see {!enter_manual_span}.
 *)
 
+val add_data_to_span : span -> (string * user_data) list -> unit
+(** Add structured data to the given span.
+    @since NEXT_RELEASE *)
+
 val enter_manual_sub_span :
   parent:explicit_span ->
   ?flavor:[ `Sync | `Async ] ->
@@ -75,6 +79,12 @@ val exit_manual_span : explicit_span -> unit
     {!enter_manual_toplevel_span}.
     @since 0.3 *)
 
+val add_data_to_explicit_span :
+  explicit_span -> (string * user_data) list -> unit
+(** [add_data_explicit esp data] is [add_data_to_span esp.span data], ie. it adds
+    the pairs [k:v] from [data] to the span inside [esp].
+    @since NEXT_RELEASE *)
+
 val message :
   ?span:span -> ?data:(unit -> (string * user_data) list) -> string -> unit
 (** [message msg] logs a message [msg] (if a collector is installed).

src/tef/trace_tef.ml

@@ -32,6 +32,8 @@ let protect ~finally f =
     finally ();
     Printexc.raise_with_backtrace exn bt
 
+let on_tracing_error = ref (fun s -> Printf.eprintf "trace-tef error: %s\n%!" s)
+
 type event =
   | E_tick
   | E_message of {
@@ -52,6 +54,10 @@ type event =
       id: span;
       time_us: float;
     }
+  | E_add_data_to_span of {
+      id: span;
+      data: (string * user_data) list;
+    }
   | E_enter_manual_span of {
       tid: int;
       name: string;
@@ -90,7 +96,7 @@ type span_info = {
   tid: int;
   name: string;
   start_us: float;
-  data: (string * user_data) list;
+  mutable data: (string * user_data) list;
 }
 
 (** key used to carry a unique "id" for all spans in an async context *)
@@ -282,11 +288,15 @@ let bg_thread ~out (events : event B_queue.t) : unit =
       Span_tbl.add spans id { tid; name; start_us = time_us; data }
     | E_exit_span { id; time_us = stop_us } ->
       (match Span_tbl.find_opt spans id with
-      | None -> (* bug! TODO: emit warning *) ()
+      | 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_to_span { id; data } ->
+      (match Span_tbl.find_opt spans id with
+      | None -> !on_tracing_error (Printf.sprintf "cannot find span %Ld" id)
+      | Some span_info -> span_info.data <- List.rev_append data span_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
@@ -379,6 +389,10 @@ let collector ~out () : collector =
 
       Fun.protect ~finally (fun () -> f span)
 
+    let add_data_to_span (span : span) data =
+      if data <> [] then
+        B_queue.push events (E_add_data_to_span { id = span; data })
+
     let enter_manual_span ~(parent : explicit_span option) ~flavor
         ~__FUNCTION__:fun_name ~__FILE__:_ ~__LINE__:_ ~data name :
         explicit_span =
@@ -450,4 +464,5 @@ let with_setup ?out () f =
 
 module Internal_ = struct
   let mock_all_ () = Mock_.enabled := true
+  let on_tracing_error = on_tracing_error
 end

src/tef/trace_tef.mli

@@ -41,7 +41,9 @@ val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
 
 module Internal_ : sig
   val mock_all_ : unit -> unit
-  (* use fake, deterministic timestamps, TID, PID *)
+  (** use fake, deterministic timestamps, TID, PID *)
+
+  val on_tracing_error : (string -> unit) ref
 end
 
 (**/**)