Merge pull request #11 from c-cube/wip-add-data-dyn

add `add_data_to_span` and `add_data_to_explicit_span`

src/core/collector.ml

@@ -53,6 +53,15 @@ module type S = sig
   (** Exit an explicit span.
       @since 0.3 *)
 
+  val add_data : (string * user_data) list -> unit
+  (** @since Adds data to the current, implicit span.
+      NEXT_RELEASE *)
+
+  val add_data_to_manual_span :
+    explicit_span -> (string * user_data) list -> unit
+  (** Adds data to the given span.
+      @since 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,20 @@ let[@inline] exit_manual_span espan : unit =
   | None -> ()
   | Some (module C) -> C.exit_manual_span espan
 
+let[@inline] add_data data : unit =
+  if data <> [] then (
+    match A.get collector with
+    | None -> ()
+    | Some (module C) -> C.add_data data
+  )
+
+let[@inline] add_data_to_manual_span esp data : unit =
+  if data <> [] then (
+    match A.get collector with
+    | None -> ()
+    | Some (module C) -> C.add_data_to_manual_span esp 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,11 @@ val with_span :
     see {!enter_manual_span}.
 *)
 
+val add_data : (string * user_data) list -> unit
+(** Add structured data to the current, implicit span (see {!with_span}).
+    Behavior is not specified if there is no current span.
+    @since NEXT_RELEASE *)
+
 val enter_manual_sub_span :
   parent:explicit_span ->
   ?flavor:[ `Sync | `Async ] ->
@@ -75,6 +80,11 @@ val exit_manual_span : explicit_span -> unit
     {!enter_manual_toplevel_span}.
     @since 0.3 *)
 
+val add_data_to_manual_span : explicit_span -> (string * user_data) list -> unit
+(** [add_data_explicit esp data] adds [data] to the span [esp].
+      The behavior is not specified is the span has been exited already.
+    @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/core/types.ml

@@ -16,6 +16,6 @@ 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]. *)
-  meta: Meta_map.t;  (** Metadata for this span (and its context) *)
+  mutable meta: Meta_map.t;  (** Metadata for this span (and its context) *)
 }
 (** Explicit span, with collector-specific metadata *)

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 of {
+      tid: int;
+      data: (string * user_data) list;
+    }
   | E_enter_manual_span of {
       tid: int;
       name: string;
@@ -66,6 +72,7 @@ type event =
       name: string;
       time_us: float;
       flavor: [ `Sync | `Async ] option;
+      data: (string * user_data) list;
       id: int;
     }
   | E_counter of {
@@ -86,11 +93,18 @@ module Span_tbl = Hashtbl.Make (struct
   let hash : t -> int = Hashtbl.hash
 end)
 
+module Int_tbl = Hashtbl.Make (struct
+  type t = int
+
+  let equal : t -> t -> bool = ( = )
+  let hash : t -> int = Hashtbl.hash
+end)
+
 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 *)
@@ -99,6 +113,9 @@ let key_async_id : int Meta_map.Key.t = Meta_map.Key.create ()
 let key_async_data : (string * [ `Sync | `Async ] option) Meta_map.Key.t =
   Meta_map.Key.create ()
 
+let key_data : (string * user_data) list ref Meta_map.Key.t =
+  Meta_map.Key.create ()
+
 (** Writer: knows how to write entries to a file in TEF format *)
 module Writer = struct
   type t = {
@@ -206,14 +223,16 @@ module Writer = struct
       args;
     ()
 
-  let emit_manual_end ~tid ~name ~id ~ts ~flavor (self : t) : unit =
+  let emit_manual_end ~tid ~name ~id ~ts ~flavor ~args (self : t) : unit =
     emit_sep_ self;
     Printf.fprintf self.oc
-      {json|{"pid":%d,"cat":"trace","id":%d,"tid": %d,"ts": %.2f,"name":%a,"ph":"%c"}|json}
+      {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
       (match flavor with
       | None | Some `Async -> 'e'
-      | Some `Sync -> 'E');
+      | Some `Sync -> 'E')
+      (emit_args_o_ pp_user_data_)
+      args;
 
     ()
 
@@ -261,6 +280,7 @@ let bg_thread ~out (events : event B_queue.t) : unit =
   Writer.with_ ~out @@ fun writer ->
   (* local state, to keep track of span information and implicit stack context *)
   let spans : span_info Span_tbl.t = Span_tbl.create 32 in
+  let ambient_span : span_info Int_tbl.t = Int_tbl.create 16 in
   let local_q = Queue.create () in
 
   (* add function name, if provided, to the metadata *)
@@ -277,22 +297,33 @@ let bg_thread ~out (events : event B_queue.t) : unit =
     | E_message { tid; msg; time_us; data } ->
       Writer.emit_instant_event ~tid ~name:msg ~ts:time_us ~args:data writer
     | E_define_span { tid; name; id; time_us; fun_name; data } ->
-      (* save the span so we find it at exit *)
       let data = add_fun_name_ fun_name data in
-      Span_tbl.add spans id { tid; name; start_us = time_us; data }
+      let info = { tid; name; start_us = time_us; data } in
+      (* make this span the "ambient" one for the given thread *)
+      Int_tbl.add ambient_span tid info;
+      (* save the span so we find it at exit *)
+      Span_tbl.add spans id info
     | 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;
+        Int_tbl.remove ambient_span tid;
         Writer.emit_duration_event ~tid ~name ~start:start_us ~end_:stop_us
           ~args:data writer)
+    | E_add_data { tid; data } ->
+      (match Int_tbl.find_opt ambient_span tid with
+      | None ->
+        !on_tracing_error
+          (Printf.sprintf "cannot find ambient span for thread %d" tid)
+      | Some info -> info.data <- List.rev_append data 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
         writer
-    | E_exit_manual_span { tid; time_us; name; id; flavor } ->
-      Writer.emit_manual_end ~tid ~name ~id ~ts:time_us ~flavor writer
+    | E_exit_manual_span { tid; time_us; name; id; flavor; data } ->
+      Writer.emit_manual_end ~tid ~name ~id ~ts:time_us ~flavor ~args:data
+        writer
     | E_counter { tid; name; time_us; n } ->
       Writer.emit_counter ~name ~tid ~ts:time_us writer n
     | E_name_process { name } -> Writer.emit_name_process ~name writer
@@ -379,6 +410,12 @@ let collector ~out () : collector =
 
       Fun.protect ~finally (fun () -> f span)
 
+    let add_data data =
+      if data <> [] then (
+        let tid = get_tid_ () in
+        B_queue.push events (E_add_data { tid; data })
+      )
+
     let enter_manual_span ~(parent : explicit_span option) ~flavor
         ~__FUNCTION__:fun_name ~__FILE__:_ ~__LINE__:_ ~data name :
         explicit_span =
@@ -402,10 +439,24 @@ let collector ~out () : collector =
     let exit_manual_span (es : explicit_span) : unit =
       let id = Meta_map.find_exn key_async_id es.meta in
       let name, flavor = Meta_map.find_exn key_async_data es.meta in
+      let data =
+        try !(Meta_map.find_exn key_data es.meta) with Not_found -> []
+      in
       let time_us = now_us () in
       let tid = get_tid_ () in
       B_queue.push events
-        (E_exit_manual_span { tid; id; name; time_us; flavor })
+        (E_exit_manual_span { tid; id; name; time_us; data; flavor })
+
+    let add_data_to_manual_span (es : explicit_span) data =
+      if data <> [] then (
+        let data_ref, add =
+          try Meta_map.find_exn key_data es.meta, false
+          with Not_found -> ref [], true
+        in
+        let new_data = List.rev_append data !data_ref in
+        data_ref := new_data;
+        if add then es.meta <- Meta_map.add key_data data_ref es.meta
+      )
 
     let message ?span:_ ~data msg : unit =
       let time_us = now_us () in
@@ -450,4 +501,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
 
 (**/**)

test/t1.ml

@@ -17,8 +17,10 @@ let run () =
       Trace.message "world";
       Trace.counter_int "n" !n;
 
+      Trace.add_data [ "i", `Int _i ];
+
       if _j = 2 then (
-        (* fake micro sleep *)
+        Trace.add_data [ "j", `Int _j ];
         let _sp =
           Trace.enter_manual_sub_span ~parent:pseudo_async_sp
             ~flavor:
@@ -28,11 +30,15 @@ let run () =
                 `Async)
             ~__FILE__ ~__LINE__ "sub-sleep"
         in
+
+        (* fake micro sleep *)
         Thread.delay 0.005;
         Trace.exit_manual_span _sp
-      ) else if _j = 3 then
+      ) else if _j = 3 then (
         (* pretend some task finished. Note that this is not well scoped wrt other spans. *)
+        Trace.add_data_to_manual_span pseudo_async_sp [ "slept", `Bool true ];
         Trace.exit_manual_span pseudo_async_sp
+      )
     done
   done