From 76703461eab929232e227768179e6cc11b070ba5 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 7 May 2025 11:08:49 -0400
Subject: [PATCH] feat(trace.subscriber): add `Span_tbl`, and a depopt on
 picos_aux

---
 .github/workflows/main.yml          |  3 +++
 dune-project                        |  1 +
 src/subscriber/callbacks.ml         | 11 ++++++++-
 src/subscriber/dune                 |  7 ++++++
 src/subscriber/span_tbl.ml          |  1 +
 src/subscriber/span_tbl.mli         | 21 ++++++++++++++++
 src/subscriber/tbl_.basic.ml        | 13 ++++++++++
 src/subscriber/tbl_.mli             |  7 ++++++
 src/subscriber/tbl_.picos.ml        | 18 ++++++++++++++
 src/subscriber/tbl_.thread.ml       | 38 +++++++++++++++++++++++++++++
 src/subscriber/trace_subscriber.ml  |  1 +
 src/subscriber/trace_subscriber.mli |  4 +++
 trace.opam                          |  1 +
 13 files changed, 125 insertions(+), 1 deletion(-)
 create mode 100644 src/subscriber/span_tbl.ml
 create mode 100644 src/subscriber/span_tbl.mli
 create mode 100644 src/subscriber/tbl_.basic.ml
 create mode 100644 src/subscriber/tbl_.mli
 create mode 100644 src/subscriber/tbl_.picos.ml
 create mode 100644 src/subscriber/tbl_.thread.ml

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 009de5a..3f894c3 100644
--- a/.github/workflows/main.yml
+++ b/.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
 
diff --git a/dune-project b/dune-project
index 3e8fcbf..0ef4777 100644
--- a/dune-project
+++ b/dune-project
@@ -28,6 +28,7 @@
  (depopts
   hmap
   unix
+  (picos_aux (>= 0.6))
   (mtime
    (>= 2.0)))
  (tags
diff --git a/src/subscriber/callbacks.ml b/src/subscriber/callbacks.ml
index 95926ea..3dd1b70 100644
--- a/src/subscriber/callbacks.ml
+++ b/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:_ = ()
diff --git a/src/subscriber/dune b/src/subscriber/dune
index 478132a..580214d 100644
--- a/src/subscriber/dune
+++ b/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
diff --git a/src/subscriber/span_tbl.ml b/src/subscriber/span_tbl.ml
new file mode 100644
index 0000000..e5113cc
--- /dev/null
+++ b/src/subscriber/span_tbl.ml
@@ -0,0 +1 @@
+include Tbl_
diff --git a/src/subscriber/span_tbl.mli b/src/subscriber/span_tbl.mli
new file mode 100644
index 0000000..32be058
--- /dev/null
+++ b/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
diff --git a/src/subscriber/tbl_.basic.ml b/src/subscriber/tbl_.basic.ml
new file mode 100644
index 0000000..e182d8a
--- /dev/null
+++ b/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 []
diff --git a/src/subscriber/tbl_.mli b/src/subscriber/tbl_.mli
new file mode 100644
index 0000000..78e2443
--- /dev/null
+++ b/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
diff --git a/src/subscriber/tbl_.picos.ml b/src/subscriber/tbl_.picos.ml
new file mode 100644
index 0000000..36dba69
--- /dev/null
+++ b/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
diff --git a/src/subscriber/tbl_.thread.ml b/src/subscriber/tbl_.thread.ml
new file mode 100644
index 0000000..54517f2
--- /dev/null
+++ b/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
diff --git a/src/subscriber/trace_subscriber.ml b/src/subscriber/trace_subscriber.ml
index d419587..eb35150 100644
--- a/src/subscriber/trace_subscriber.ml
+++ b/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
diff --git a/src/subscriber/trace_subscriber.mli b/src/subscriber/trace_subscriber.mli
index a3f2325..5b535bf 100644
--- a/src/subscriber/trace_subscriber.mli
+++ b/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
diff --git a/trace.opam b/trace.opam
index 83bc324..8813330 100644
--- a/trace.opam
+++ b/trace.opam
@@ -17,6 +17,7 @@ depends: [
 depopts: [
   "hmap"
   "unix"
+  "picos_aux" {>= "0.6"}
   "mtime" {>= "2.0"}
 ]
 build: [