add optional metadata to messages and spans

src/collector.ml

@@ -11,12 +11,22 @@ let dummy_span : span = Int64.min_int
 
 module type S = sig
   val enter_span :
-    ?__FUNCTION__:string -> __FILE__:string -> __LINE__:int -> string -> span
+    ?__FUNCTION__:string ->
+    __FILE__:string ->
+    __LINE__:int ->
+    data:(string * user_data) list ->
+    string ->
+    span
 
   val exit_span : span -> unit
 
   val message :
-    ?__FUNCTION__:string -> __FILE__:string -> __LINE__:int -> string -> unit
+    ?__FUNCTION__:string ->
+    __FILE__:string ->
+    __LINE__:int ->
+    data:(string * user_data) list ->
+    string ->
+    unit
 
   val shutdown : unit -> unit
   (** Shutdown collector, possibly waiting for it to finish sending data. *)

src/tef/trace_tef.ml

@@ -42,6 +42,7 @@ type event =
       tid: int;
       msg: string;
       time_us: float;
+      data: (string * user_data) list;
     }
   | E_define_span of {
       (*
@@ -53,6 +54,7 @@ type event =
       name: string;
       time_us: float;
       id: span;
+      data: (string * user_data) list;
     }
   | E_exit_span of {
       id: span;
@@ -74,6 +76,7 @@ type span_info = {
   tid: int;
   name: string;
   start_us: float;
+  data: (string * user_data) list;
 }
 
 module Writer = struct
@@ -132,6 +135,12 @@ module Writer = struct
     String.iter encode_char s;
     char oc '"'
 
+  let pp_user_data_ out : user_data -> unit = function
+    | `None -> Printf.fprintf out "null"
+    | `Int i -> Printf.fprintf out "%d" i
+    | `Bool b -> Printf.fprintf out "%b" b
+    | `String s -> str_val out s
+
   (* emit args, if not empty. [ppv] is used to print values. *)
   let emit_args_o_ ppv oc args : unit =
     if args <> [] then (
@@ -150,14 +159,18 @@ module Writer = struct
     emit_sep_ self;
     Printf.fprintf self.oc
       {json|{"pid": %d,"cat":"","tid": %d,"dur": %.2f,"ts": %.2f,"name":%a,"ph":"X"%a}|json}
-      self.pid tid dur ts str_val name (emit_args_o_ str_val) args;
+      self.pid tid dur ts str_val name
+      (emit_args_o_ pp_user_data_)
+      args;
     ()
 
   let emit_instant_event ~tid ~name ~ts ~args (self : t) : unit =
     emit_sep_ self;
     Printf.fprintf self.oc
       {json|{"pid": %d,"cat":"","tid": %d,"ts": %.2f,"name":%a,"ph":"I"%a}|json}
-      self.pid tid ts str_val name (emit_args_o_ str_val) args;
+      self.pid tid ts str_val name
+      (emit_args_o_ pp_user_data_)
+      args;
     ()
 end
 
@@ -169,24 +182,29 @@ let bg_thread ~out (events : event B_queue.t) : unit =
   (* how to deal with an event *)
   let handle_ev (ev : event) : unit =
     match ev with
-    | E_message { (* __FUNCTION__; __FILE__; __LINE__; *) tid; msg; time_us } ->
+    | E_message
-      Writer.emit_instant_event ~tid ~name:msg ~ts:time_us ~args:[] writer
+        { (* __FUNCTION__; __FILE__; __LINE__; *) tid; msg; time_us; data } ->
+      Writer.emit_instant_event ~tid ~name:msg ~ts:time_us ~args:data writer
     | E_define_span
-        { (* __FUNCTION__; __FILE__; __LINE__; *) tid; name; id; time_us } ->
+        { (* __FUNCTION__; __FILE__; __LINE__; *) tid; name; id; time_us; data }
+      ->
       (* save the span so we find it at exit *)
       Span_tbl.add spans id
         {
           (* __FUNCTION__; __FILE__; __LINE__; *) 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 *) ()
-      | Some { (* __FUNCTION__; __FILE__; __LINE__; *) tid; name; start_us } ->
+      | Some
+          { (* __FUNCTION__; __FILE__; __LINE__; *) tid; name; start_us; data }
+        ->
         Span_tbl.remove spans id;
         Writer.emit_duration_event ~tid ~name ~start:start_us ~end_:stop_us
-          ~args:[] writer)
+          ~args:data writer)
   in
 
   try
@@ -232,7 +250,7 @@ let collector ~out () : collector =
       else
         Thread.id (Thread.self ())
 
-    let enter_span ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ name : span =
+    let enter_span ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : span =
       let span = Int64.of_int (A.fetch_and_add span_id_gen_ 1) in
       let tid = get_tid_ () in
       let time_us = now_us () in
@@ -243,6 +261,7 @@ let collector ~out () : collector =
              name;
              time_us;
              id = span;
+             data;
            });
       span
 
@@ -250,11 +269,12 @@ let collector ~out () : collector =
       let time_us = now_us () in
       B_queue.push events (E_exit_span { id = span; time_us })
 
-    let message ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ msg : unit =
+    let message ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data msg : unit =
       let time_us = now_us () in
       let tid = get_tid_ () in
       B_queue.push events
-        (E_message { (* __FUNCTION__; __FILE__; __LINE__; *) tid; time_us; msg })
+        (E_message
+           { (* __FUNCTION__; __FILE__; __LINE__; *) tid; time_us; msg; data })
   end in
   (module M)
 

src/trace.ml

@@ -12,10 +12,16 @@ let[@inline] enabled () =
   | None -> false
   | Some _ -> true
 
-let[@inline] enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ name : span =
+let enter_span_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
+    ~__LINE__ ?(data = fun () -> []) name : span =
+  let data = data () in
+  C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
+
+let[@inline] enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name : span =
   match A.get collector with
   | None -> Collector.dummy_span
-  | Some (module C) -> C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ name
+  | Some coll ->
+    enter_span_collector_ coll ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
 
 let[@inline] exit_span span : unit =
   match A.get collector with
@@ -23,8 +29,9 @@ let[@inline] exit_span span : unit =
   | Some (module C) -> C.exit_span span
 
 let with_span_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
-    ~__LINE__ name f =
+    ~__LINE__ ?(data = fun () -> []) name f =
-  let sp = C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ name in
+  let data = data () in
+  let sp = C.enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data name in
   match f sp with
   | x ->
     C.exit_span sp;
@@ -34,26 +41,39 @@ let with_span_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
     C.exit_span sp;
     Printexc.raise_with_backtrace exn bt
 
-let[@inline] with_span ?__FUNCTION__ ~__FILE__ ~__LINE__ name f =
+let[@inline] with_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name f =
   match A.get collector with
   | None ->
     (* fast path: no collector, no span *)
     f Collector.dummy_span
   | Some collector ->
-    with_span_collector_ collector ?__FUNCTION__ ~__FILE__ ~__LINE__ name f
+    with_span_collector_ collector ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
+      f
 
-let[@inline] message ?__FUNCTION__ ~__FILE__ ~__LINE__ msg : unit =
+let message_collector_ (module C : Collector.S) ?__FUNCTION__ ~__FILE__
+    ~__LINE__ ?(data = fun () -> []) msg : unit =
+  let data = data () in
+  C.message ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data msg
+
+let[@inline] message ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data msg : unit =
   match A.get collector with
   | None -> ()
-  | Some (module C) -> C.message ?__FUNCTION__ ~__FILE__ ~__LINE__ msg
+  | Some coll ->
+    message_collector_ coll ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data msg
 
-let messagef ?__FUNCTION__ ~__FILE__ ~__LINE__ k =
+let messagef ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data k =
   match A.get collector with
   | None -> ()
   | Some (module C) ->
     k (fun fmt ->
         Format.kasprintf
-          (fun str -> C.message ?__FUNCTION__ ~__FILE__ ~__LINE__ str)
+          (fun str ->
+            let data =
+              match data with
+              | None -> []
+              | Some f -> f ()
+            in
+            C.message ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data str)
           fmt)
 
 let setup_collector c : unit =

src/trace.mli

@@ -12,7 +12,12 @@ val enabled : unit -> bool
     any span or message *)
 
 val enter_span :
-  ?__FUNCTION__:string -> __FILE__:string -> __LINE__:int -> string -> span
+  ?__FUNCTION__:string ->
+  __FILE__:string ->
+  __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
+  string ->
+  span
 
 val exit_span : span -> unit
 
@@ -20,12 +25,18 @@ val with_span :
   ?__FUNCTION__:string ->
   __FILE__:string ->
   __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
   string ->
   (span -> 'a) ->
   'a
 
 val message :
-  ?__FUNCTION__:string -> __FILE__:string -> __LINE__:int -> string -> unit
+  ?__FUNCTION__:string ->
+  __FILE__:string ->
+  __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
+  string ->
+  unit
 
 (* TODO: counter/plot/metric *)
 
@@ -33,6 +44,7 @@ val messagef :
   ?__FUNCTION__:string ->
   __FILE__:string ->
   __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
   ((('a, Format.formatter, unit, unit) format4 -> 'a) -> unit) ->
   unit
 

src/types.ml

@@ -1,2 +1,10 @@
 type span = int64
 (** A span identifier. *)
+
+type user_data =
+  [ `Int of int
+  | `String of string
+  | `Bool of bool
+  | `None
+  ]
+(** User defined data, generally passed as key/value pairs *)