Add explicit spans, for async tracing.

These spans require the user to pass the surrounding span, if any, when
entering a new span. They use the information inside (which is collector
specific) to track what asynchronous task is currently being executed.

Wrappers around `trace`, specific to an async library (e.g. Eio, lwt,
Async, etc.) can then smooth this over by providing a `with_async_span`
construct that uses some implicit contextual storage to carry the
`surrounding` scope around.

src/core/collector.ml

@@ -9,6 +9,9 @@ open Types
 
 let dummy_span : span = Int64.min_int
 
+let dummy_explicit_span : explicit_span =
+  { span = dummy_span; meta = Meta_map.empty }
+
 (** Signature for a collector.
 
     This is only relevant to implementors of tracing backends; to instrument
@@ -27,6 +30,27 @@ module type S = sig
   (** Exit given span. It can't be exited again. Spans must follow
       a strict stack discipline on each thread. *)
 
+  val enter_explicit_span :
+    surrounding:explicit_span option ->
+    ?__FUNCTION__:string ->
+    __FILE__:string ->
+    __LINE__:int ->
+    data:(string * user_data) list ->
+    string ->
+    explicit_span
+  (** Enter an explicit span. Surrounding scope is provided by [surrounding],
+      and this function can store as much metadata as it wants in the hmap
+      in the {!explicit_span}'s [meta] field.
+
+      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}.
+      @since NEXT_RELEASE *)
+
+  val exit_explicit_span : explicit_span -> unit
+  (** Exit an explicit span.
+      @since NEXT_RELEASE *)
+
   val message : ?span:span -> data:(string * user_data) list -> string -> unit
   (** Emit a message with associated metadata. *)
 

src/core/meta_map.ml

@@ -0,0 +1,83 @@
+module type KEY_IMPL = sig
+  type t
+
+  exception Store of t
+
+  val id : int
+end
+
+module Key = struct
+  type 'a t = (module KEY_IMPL with type t = 'a)
+
+  let _n = ref 0
+
+  let create (type k) () =
+    incr _n;
+    let id = !_n in
+    let module K = struct
+      type t = k
+
+      let id = id
+
+      exception Store of k
+    end in
+    (module K : KEY_IMPL with type t = k)
+
+  let id (type k) (module K : KEY_IMPL with type t = k) = K.id
+
+  let equal : type a b. a t -> b t -> bool =
+   fun (module K1) (module K2) -> K1.id = K2.id
+end
+
+type pair = Pair : 'a Key.t * 'a -> pair
+type exn_pair = E_pair : 'a Key.t * exn -> exn_pair
+
+let pair_of_e_pair (E_pair (k, e)) =
+  let module K = (val k) in
+  match e with
+  | K.Store v -> Pair (k, v)
+  | _ -> assert false
+
+module M = Map.Make (struct
+  type t = int
+
+  let compare (i : int) j = Stdlib.compare i j
+end)
+
+type t = exn_pair M.t
+
+let empty = M.empty
+let mem k t = M.mem (Key.id k) t
+
+let find_exn (type a) (k : a Key.t) t : a =
+  let module K = (val k) in
+  let (E_pair (_, e)) = M.find K.id t in
+  match e with
+  | K.Store v -> v
+  | _ -> assert false
+
+let find k t = try Some (find_exn k t) with Not_found -> None
+
+let add_e_pair_ p t =
+  let (E_pair ((module K), _)) = p in
+  M.add K.id p t
+
+let add_pair_ p t =
+  let (Pair (((module K) as k), v)) = p in
+  let p = E_pair (k, K.Store v) in
+  M.add K.id p t
+
+let add (type a) (k : a Key.t) v t =
+  let module K = (val k) in
+  add_e_pair_ (E_pair (k, K.Store v)) t
+
+let remove (type a) (k : a Key.t) t =
+  let module K = (val k) in
+  M.remove K.id t
+
+let cardinal t = M.cardinal t
+let length = cardinal
+let iter f t = M.iter (fun _ p -> f (pair_of_e_pair p)) t
+let to_list t = M.fold (fun _ p l -> pair_of_e_pair p :: l) t []
+let add_list t l = List.fold_right add_pair_ l t
+let of_list l = add_list empty l

src/core/meta_map.mli

@@ -0,0 +1,37 @@
+(** Associative containers with Heterogeneous Values *)
+
+(** Keys with a type witness. *)
+module Key : sig
+  type 'a t
+  (** A key of type ['a t] is used to access the portion of the
+      map or table that associates keys of type ['a] to values. *)
+
+  val create : unit -> 'a t
+  (** Make a new key. This is generative, so calling [create ()] twice with the
+      same return type will produce incompatible keys that cannot see each
+      other's bindings. *)
+
+  val equal : 'a t -> 'a t -> bool
+  (** Compare two keys that have compatible types. *)
+end
+
+type pair = Pair : 'a Key.t * 'a -> pair
+
+type t
+(** Immutable map from {!Key.t} to values *)
+
+val empty : t
+val mem : _ Key.t -> t -> bool
+val add : 'a Key.t -> 'a -> t -> t
+val remove : _ Key.t -> t -> t
+val length : t -> int
+val cardinal : t -> int
+val find : 'a Key.t -> t -> 'a option
+
+val find_exn : 'a Key.t -> t -> 'a
+(** @raise Not_found if the key is not in the table. *)
+
+val iter : (pair -> unit) -> t -> unit
+val add_list : t -> pair list -> t
+val of_list : pair list -> t
+val to_list : t -> pair list

src/core/trace_core.ml

@@ -1,6 +1,7 @@
 include Types
 module A = Atomic_
 module Collector = Collector
+module Meta_map = Meta_map
 
 type collector = (module Collector.S)
 
@@ -50,6 +51,26 @@ let[@inline] with_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name f =
     with_span_collector_ collector ?__FUNCTION__ ~__FILE__ ~__LINE__ ?data name
       f
 
+let enter_explicit_span_collector_ (module C : Collector.S) ~surrounding
+    ?__FUNCTION__ ~__FILE__ ~__LINE__ ?(data = fun () -> []) name :
+    explicit_span =
+  let data = data () in
+  C.enter_explicit_span ~surrounding ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data
+    name
+
+let[@inline] enter_explicit_span ~surrounding ?__FUNCTION__ ~__FILE__ ~__LINE__
+    ?data name : explicit_span =
+  match A.get collector with
+  | None -> Collector.dummy_explicit_span
+  | Some coll ->
+    enter_explicit_span_collector_ coll ~surrounding ?__FUNCTION__ ~__FILE__
+      ~__LINE__ ?data name
+
+let[@inline] exit_explicit_span espan : unit =
+  match A.get collector with
+  | None -> ()
+  | Some (module C) -> C.exit_explicit_span espan
+
 let message_collector_ (module C : Collector.S) ?span ?(data = fun () -> []) msg
     : unit =
   let data = data () in

src/core/trace_core.mli

@@ -2,6 +2,7 @@
 
 include module type of Types
 module Collector = Collector
+module Meta_map = Meta_map
 
 (** {2 Tracing} *)
 
@@ -47,7 +48,33 @@ val with_span :
     [sp] might be a dummy span if no collector is installed.
     When [f sp] returns or raises, the span [sp] is exited.
 
-    This is the recommended way to instrument most code. *)
+    This is the recommended way to instrument most code.
+
+    {b NOTE} an important restriction is that this is only supposed to
+    work for synchronous, direct style code. Monadic concurrency, Effect-based
+    fibers, etc. might not play well with this style of spans on some
+    or all backends. If you use cooperative concurrency,
+    see {!enter_explicit_span}.
+*)
+
+val enter_explicit_span :
+  surrounding:explicit_span option ->
+  ?__FUNCTION__:string ->
+  __FILE__:string ->
+  __LINE__:int ->
+  ?data:(unit -> (string * user_data) list) ->
+  string ->
+  explicit_span
+(** Like {!enter_span} but the caller is responsible for
+      providing the [surrounding] context, and carry the resulting
+      {!explicit_span} to  the matching {!exit_explicit_span}. 
+      @since NEXT_RELEASE *)
+
+val exit_explicit_span : explicit_span -> unit
+(** Exit an explicit span. This can be on another thread, in a
+    fiber or lightweight thread, etc. and will be supported by backends
+    nonetheless.
+    @since NEXT_RELEASE *)
 
 val message :
   ?span:span -> ?data:(unit -> (string * user_data) list) -> string -> unit

src/core/types.ml

@@ -11,3 +11,11 @@ type user_data =
   ]
 (** User defined data, generally passed as key/value pairs to
     whatever collector is installed (if any). *)
+
+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) *)
+}
+(** Explicit span, with collector-specific metadata *)

src/tef/trace_tef.ml

@@ -51,6 +51,19 @@ type event =
       id: span;
       time_us: float;
     }
+  | E_enter_async_span of {
+      tid: int;
+      name: string;
+      time_us: float;
+      id: int;
+      data: (string * user_data) list;
+    }
+  | E_exit_async_span of {
+      tid: int;
+      name: string;
+      time_us: float;
+      id: int;
+    }
   | E_counter of {
       name: string;
       tid: int;
@@ -76,6 +89,11 @@ type span_info = {
   data: (string * user_data) list;
 }
 
+(** key used to carry a unique "id" for all spans in an async context *)
+let key_async_id : int Meta_map.Key.t = Meta_map.Key.create ()
+
+let key_async_name : string Meta_map.Key.t = Meta_map.Key.create ()
+
 module Writer = struct
   type t = {
     oc: out_channel;
@@ -164,6 +182,22 @@ module Writer = struct
       args;
     ()
 
+  let emit_async_begin ~tid ~name ~id ~ts ~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":"b"%a}|json}
+      self.pid id tid ts str_val name
+      (emit_args_o_ pp_user_data_)
+      args;
+    ()
+
+  let emit_async_end ~tid ~name ~id ~ts (self : t) : unit =
+    emit_sep_ self;
+    Printf.fprintf self.oc
+      {json|{"pid":%d,"cat":"trace","id":%d,"tid": %d,"ts": %.2f,"name":%a,"ph":"e"}|json}
+      self.pid id tid ts str_val name;
+    ()
+
   let emit_instant_event ~tid ~name ~ts ~args (self : t) : unit =
     emit_sep_ self;
     Printf.fprintf self.oc
@@ -222,6 +256,10 @@ let bg_thread ~out (events : event B_queue.t) : unit =
         Span_tbl.remove spans id;
         Writer.emit_duration_event ~tid ~name ~start:start_us ~end_:stop_us
           ~args:data writer)
+    | E_enter_async_span { tid; time_us; name; id; data } ->
+      Writer.emit_async_begin ~tid ~name ~id ~ts:time_us ~args:data writer
+    | E_exit_async_span { tid; time_us; name; id } ->
+      Writer.emit_async_end ~tid ~name ~id ~ts:time_us 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
@@ -302,6 +340,30 @@ let collector ~out () : collector =
       let time_us = now_us () in
       B_queue.push events (E_exit_span { id = span; time_us })
 
+    let enter_explicit_span ~(surrounding : explicit_span option)
+        ?__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_ ~data name : explicit_span =
+      (* get the id, or make a new one *)
+      let id =
+        match surrounding with
+        | Some m -> Meta_map.find_exn key_async_id m.meta
+        | None -> A.fetch_and_add span_id_gen_ 1
+      in
+      let time_us = now_us () in
+      B_queue.push events
+        (E_enter_async_span { id; time_us; tid = get_tid_ (); data; name });
+      {
+        span = 0L;
+        meta =
+          Meta_map.(empty |> add key_async_id id |> add key_async_name name);
+      }
+
+    let exit_explicit_span (es : explicit_span) : unit =
+      let id = Meta_map.find_exn key_async_id es.meta in
+      let name = Meta_map.find_exn key_async_name es.meta in
+      let time_us = now_us () in
+      let tid = get_tid_ () in
+      B_queue.push events (E_exit_async_span { tid; id; name; time_us })
+
     let message ?span:_ ~data msg : unit =
       let time_us = now_us () in
       let tid = get_tid_ () in