feat subscriber: avoid polyvariants entirely

src/subscriber/callbacks.ml

@@ -1,4 +1,5 @@
 open Trace_core
+open Types
 
 module type S = sig
   type st
@@ -64,7 +65,7 @@ module type S = sig
     parent:span option ->
     data:(string * user_data) list ->
     name:string ->
-    flavor:[ `Sync | `Async ] option ->
+    flavor:flavor option ->
     trace_id:int ->
     span ->
     unit
@@ -76,7 +77,7 @@ module type S = sig
     tid:int ->
     name:string ->
     data:(string * user_data) list ->
-    flavor:[ `Sync | `Async ] option ->
+    flavor:flavor option ->
     trace_id:int ->
     span ->
     unit

src/subscriber/trace_subscriber.ml

@@ -1,6 +1,7 @@
 open Trace_core
 module Callbacks = Callbacks
 module Subscriber = Subscriber
+include Types
 
 type t = Subscriber.t
 
@@ -27,7 +28,7 @@ open struct
 
   type manual_span_info = {
     name: string;
-    flavor: [ `Sync | `Async ] option;
+    flavor: flavor option;
     mutable data: (string * user_data) list;
   }
 
@@ -39,6 +40,26 @@ open struct
   let key_async_trace_id : int Meta_map.key = Meta_map.Key.create ()
 end
 
+let[@inline] conv_flavor = function
+  | `Async -> Async
+  | `Sync -> Sync
+
+let[@inline] conv_flavor_opt = function
+  | None -> None
+  | Some f -> Some (conv_flavor f)
+
+let[@inline] conv_user_data = function
+  | `Int i -> U_int i
+  | `Bool b -> U_bool b
+  | `Float f -> U_float f
+  | `String s -> U_string s
+  | `None -> U_none
+
+let rec conv_data = function
+  | [] -> []
+  | [ (k, v) ] -> [ k, conv_user_data v ]
+  | (k, v) :: tl -> (k, conv_user_data v) :: conv_data tl
+
 (** A collector that calls the callbacks of subscriber *)
 let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
   let open Private_ in
@@ -54,6 +75,7 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
       let span = Int64.of_int (new_span_ ()) in
       let tid = tid_ () in
       let time_ns = now_ns () in
+      let data = conv_data data in
       CB.on_enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~time_ns ~tid ~data
         ~name span;
       span
@@ -75,13 +97,18 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
         Printexc.raise_with_backtrace exn bt
 
     let add_data_to_span span data =
-      if data <> [] then CB.on_add_data st ~data span
+      if data <> [] then (
+        let data = conv_data data in
+        CB.on_add_data st ~data span
+      )
 
     let enter_manual_span ~(parent : explicit_span 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
+      let data = conv_data data in
+      let flavor = conv_flavor_opt flavor in
 
       (* get the common trace id, or make a new one *)
       let trace_id, parent =
@@ -117,6 +144,7 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
 
     let add_data_to_manual_span (es : explicit_span) data =
       if data <> [] then (
+        let data = conv_data data in
         match Meta_map.find key_manual_info es.meta with
         | None -> assert false
         | Some m -> m.data <- List.rev_append data m.data
@@ -125,11 +153,13 @@ let collector (Sub { st; callbacks = (module CB) } : Subscriber.t) : collector =
     let message ?span ~data msg : unit =
       let time_ns = now_ns () in
       let tid = tid_ () in
+      let data = conv_data data in
       CB.on_message st ~time_ns ~tid ~span ~data msg
 
     let counter_float ~data name f : unit =
       let time_ns = now_ns () in
       let tid = tid_ () in
+      let data = conv_data data in
       CB.on_counter st ~tid ~time_ns ~data ~name f
 
     let[@inline] counter_int ~data name i =

src/subscriber/trace_subscriber.mli

@@ -11,6 +11,10 @@
 module Callbacks = Callbacks
 module Subscriber = Subscriber
 
+include module type of struct
+  include Types
+end
+
 type t = Subscriber.t
 
 val collector : t -> Trace_core.collector

src/subscriber/types.ml

@@ -0,0 +1,10 @@
+type user_data =
+  | U_bool of bool
+  | U_float of float
+  | U_int of int
+  | U_none
+  | U_string of string
+
+type flavor =
+  | Sync
+  | Async

src/tef/event.ml

@@ -1,4 +1,5 @@
 open Trace_core
+module Sub = Trace_subscriber
 
 (** An event, specialized for TEF *)
 type t =
@@ -7,7 +8,7 @@ type t =
       tid: int;
       msg: string;
       time_us: float;
-      data: (string * user_data) list;
+      data: (string * Sub.user_data) list;
     }
   | E_define_span of {
       tid: int;
@@ -15,7 +16,7 @@ type t =
       time_us: float;
       id: span;
       fun_name: string option;
-      data: (string * user_data) list;
+      data: (string * Sub.user_data) list;
     }
   | E_exit_span of {
       id: span;
@@ -23,23 +24,23 @@ type t =
     }
   | E_add_data of {
       id: span;
-      data: (string * user_data) list;
+      data: (string * Sub.user_data) list;
     }
   | E_enter_manual_span of {
       tid: int;
       name: string;
       time_us: float;
       id: int;
-      flavor: [ `Sync | `Async ] option;
+      flavor: Sub.flavor option;
       fun_name: string option;
-      data: (string * user_data) list;
+      data: (string * Sub.user_data) list;
     }
   | E_exit_manual_span of {
       tid: int;
       name: string;
       time_us: float;
-      flavor: [ `Sync | `Async ] option;
-      data: (string * user_data) list;
+      flavor: Sub.flavor option;
+      data: (string * Sub.user_data) list;
       id: int;
     }
   | E_counter of {

src/tef/trace_tef.ml

@@ -29,7 +29,7 @@ type span_info = {
   tid: int;
   name: string;
   start_us: float;
-  mutable data: (string * user_data) list;
+  mutable data: (string * Sub.user_data) list;
 }
 
 (** Writer: knows how to write entries to a file in TEF format *)
@@ -110,12 +110,12 @@ module Writer = struct
     String.iter encode_char s;
     char buf '"'
 
-  let pp_user_data_ (out : Buffer.t) : [< user_data ] -> unit = function
-    | `None -> raw_string out "null"
-    | `Int i -> Printf.bprintf out "%d" i
-    | `Bool b -> Printf.bprintf out "%b" b
-    | `String s -> str_val out s
-    | `Float f -> Printf.bprintf out "%g" f
+  let pp_user_data_ (out : Buffer.t) : Sub.user_data -> unit = function
+    | U_none -> raw_string out "null"
+    | U_int i -> Printf.bprintf out "%d" i
+    | U_bool b -> Printf.bprintf out "%b" b
+    | U_string s -> str_val out s
+    | U_float f -> Printf.bprintf out "%g" f
 
   (* emit args, if not empty. [ppv] is used to print values. *)
   let emit_args_o_ ppv (out : Buffer.t) args : unit =
@@ -142,26 +142,28 @@ module Writer = struct
       args;
     Buffer.output_buffer self.oc self.buf
 
-  let emit_manual_begin ~tid ~name ~id ~ts ~args ~flavor (self : t) : unit =
+  let emit_manual_begin ~tid ~name ~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
       (match flavor with
-      | None | Some `Async -> 'b'
-      | Some `Sync -> 'B')
+      | None | Some Async -> 'b'
+      | Some Sync -> 'B')
       (emit_args_o_ pp_user_data_)
       args;
     Buffer.output_buffer self.oc self.buf
 
-  let emit_manual_end ~tid ~name ~id ~ts ~flavor ~args (self : t) : unit =
+  let emit_manual_end ~tid ~name ~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
       (match flavor with
-      | None | Some `Async -> 'e'
-      | Some `Sync -> 'E')
+      | None | Some Async -> 'e'
+      | Some Sync -> 'E')
       (emit_args_o_ pp_user_data_)
       args;
     Buffer.output_buffer self.oc self.buf
@@ -181,7 +183,7 @@ module Writer = struct
       {json|{"pid":%d,"tid": %d,"name":"thread_name","ph":"M"%a}|json} self.pid
       tid
       (emit_args_o_ pp_user_data_)
-      [ "name", `String name ];
+      [ "name", U_string name ];
     Buffer.output_buffer self.oc self.buf
 
   let emit_name_process ~name (self : t) : unit =
@@ -189,7 +191,7 @@ module Writer = struct
     Printf.bprintf self.buf
       {json|{"pid":%d,"name":"process_name","ph":"M"%a}|json} self.pid
       (emit_args_o_ pp_user_data_)
-      [ "name", `String name ];
+      [ "name", U_string name ];
     Buffer.output_buffer self.oc self.buf
 
   let emit_counter ~name ~tid ~ts (self : t) f : unit =
@@ -198,7 +200,7 @@ module Writer = struct
       {json|{"pid":%d,"tid":%d,"ts":%.2f,"name":"c","ph":"C"%a}|json} self.pid
       tid ts
       (emit_args_o_ pp_user_data_)
-      [ name, `Float f ];
+      [ name, U_float f ];
     Buffer.output_buffer self.oc self.buf
 end
 
@@ -215,7 +217,7 @@ let bg_thread ~mode ~out (events : Event.t B_queue.t) : unit =
   let add_fun_name_ fun_name data : _ list =
     match fun_name with
     | None -> data
-    | Some f -> ("function", `String f) :: data
+    | Some f -> ("function", Sub.U_string f) :: data
   in
 
   (* how to deal with an event *)