feat(trace.subscriber): add Span_tbl, and a depopt on picos_aux

.github/workflows/main.yml

@@ -51,6 +51,9 @@ jobs:
       - run: opam install hmap
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
 
+      - run: opam install picos_aux
+      - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
+
       - run: opam install mtime
       - run: opam exec -- dune build '@install' -p trace,trace-tef,trace-fuchsia
 

dune-project

@@ -28,6 +28,7 @@
  (depopts
   hmap
   unix
+  (picos_aux (>= 0.6))
   (mtime
    (>= 2.0)))
  (tags

src/subscriber/callbacks.ml

@@ -121,7 +121,16 @@ type 'st t = (module S with type st = 'st)
 
 (** Dummy callbacks. It can be useful to reuse some of these functions in a real
     subscriber that doesn't want to handle {b all} events, but only some of
-    them. *)
+    them.
+
+    To write a subscriber that only supports some callbacks, this can be handy:
+    {[
+      module My_callbacks = struct
+      type st = my_own_state
+      include Callbacks.Dummy
+      let on_counter (st:st) ~time_ns ~tid ~data ~name v : unit = ...
+      end
+    ]} *)
 module Dummy = struct
   let on_init _ ~time_ns:_ = ()
   let on_shutdown _ ~time_ns:_ = ()

src/subscriber/dune

@@ -1,6 +1,7 @@
 (library
  (name trace_subscriber)
  (public_name trace.subscriber)
+ (private_modules time_ thread_ tbl_)
  (libraries
   (re_export trace.core)
   (select
@@ -8,6 +9,12 @@
    from
    (threads -> thread_.real.ml)
    (-> thread_.dummy.ml))
+  (select
+   tbl_.ml
+   from
+   (picos_aux.htbl -> tbl_.picos.ml)
+   (threads -> tbl_.thread.ml)
+   (-> tbl_.basic.ml))
   (select
    time_.ml
    from

src/subscriber/span_tbl.ml

@@ -0,0 +1 @@
+include Tbl_

src/subscriber/span_tbl.mli

@@ -0,0 +1,21 @@
+(** A table that can be used to remember information about spans.
+
+    This is convenient when we want to rememner information from a span begin,
+    when dealing with the corresponding span end.
+
+    {b NOTE}: this is thread safe when threads are enabled. *)
+
+open Trace_core
+
+type 'v t
+
+val create : unit -> 'v t
+val add : 'v t -> span -> 'v -> unit
+
+val find_exn : 'v t -> span -> 'v
+(** @raise Not_found if information isn't found *)
+
+val remove : _ t -> span -> unit
+(** Remove the span if present *)
+
+val to_list : 'v t -> (span * 'v) list

src/subscriber/tbl_.basic.ml

@@ -0,0 +1,13 @@
+module T = Hashtbl.Make (struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end)
+
+type 'v t = 'v T.t
+
+let create () : _ t = T.create 32
+let find_exn = T.find
+let remove = T.remove
+let add = T.replace
+let to_list self : _ list = T.fold (fun k v l -> (k, v) :: l) self []

src/subscriber/tbl_.mli

@@ -0,0 +1,7 @@
+type 'v t
+
+val create : unit -> 'v t
+val add : 'v t -> int64 -> 'v -> unit
+val find_exn : 'v t -> int64 -> 'v
+val remove : _ t -> int64 -> unit
+val to_list : 'v t -> (int64 * 'v) list

src/subscriber/tbl_.picos.ml

@@ -0,0 +1,18 @@
+module H = Picos_aux_htbl
+
+module Key = struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end
+
+type 'v t = (int64, 'v) H.t
+
+let create () : _ t = H.create ~hashed_type:(module Key) ()
+let find_exn = H.find_exn
+let[@inline] remove self k = ignore (H.try_remove self k : bool)
+
+let[@inline] add self k v =
+  if not (H.try_add self k v) then ignore (H.try_set self k v)
+
+let[@inline] to_list self = H.to_seq self |> List.of_seq

src/subscriber/tbl_.thread.ml

@@ -0,0 +1,38 @@
+module T = Hashtbl.Make (struct
+  include Int64
+
+  let hash = Hashtbl.hash
+end)
+
+type 'v t = {
+  tbl: 'v T.t;
+  lock: Mutex.t;
+}
+
+let create () : _ t = { tbl = T.create 32; lock = Mutex.create () }
+
+let find_exn self k =
+  Mutex.lock self.lock;
+  try
+    let v = T.find self.tbl k in
+    Mutex.unlock self.lock;
+    v
+  with e ->
+    Mutex.unlock self.lock;
+    raise e
+
+let remove self k =
+  Mutex.lock self.lock;
+  T.remove self.tbl k;
+  Mutex.unlock self.lock
+
+let add self k v =
+  Mutex.lock self.lock;
+  T.replace self.tbl k v;
+  Mutex.unlock self.lock
+
+let to_list self : _ list =
+  Mutex.lock self.lock;
+  let l = T.fold (fun k v l -> (k, v) :: l) self.tbl [] in
+  Mutex.unlock self.lock;
+  l

src/subscriber/trace_subscriber.ml

@@ -1,6 +1,7 @@
 open Trace_core
 module Callbacks = Callbacks
 module Subscriber = Subscriber
+module Span_tbl = Span_tbl
 include Types
 
 type t = Subscriber.t

src/subscriber/trace_subscriber.mli

@@ -4,10 +4,14 @@
     trace event. It also defines a collector that needs to be installed for the
     subscriber(s) to be called.
 
+    Thanks to {!Subscriber.tee_l} it's possible to combine multiple subscribers
+    into a single collector.
+
     @since 0.8 *)
 
 module Callbacks = Callbacks
 module Subscriber = Subscriber
+module Span_tbl = Span_tbl
 
 include module type of struct
   include Types

trace.opam

@@ -17,6 +17,7 @@ depends: [
 depopts: [
   "hmap"
   "unix"
+  "picos_aux" {>= "0.6"}
   "mtime" {>= "2.0"}
 ]
 build: [