From dd432c4586c89f7971f3648489ff29c1115c00a9 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 17 Jan 2026 10:42:59 -0500
Subject: [PATCH] add trace.debug to find what spans were not closed on exit

---
 src/debug/dune            |   6 ++
 src/debug/track_spans.ml  | 120 ++++++++++++++++++++++++++++++++++++++
 src/debug/track_spans.mli |  24 ++++++++
 3 files changed, 150 insertions(+)
 create mode 100644 src/debug/dune
 create mode 100644 src/debug/track_spans.ml
 create mode 100644 src/debug/track_spans.mli

diff --git a/src/debug/dune b/src/debug/dune
new file mode 100644
index 0000000..b803f78
--- /dev/null
+++ b/src/debug/dune
@@ -0,0 +1,6 @@
+(library
+ (name trace_debug)
+ (public_name trace.debug)
+ (synopsis "helper to debug unclosed spans")
+ (optional) ; threads
+ (libraries threads trace.core trace.util))
diff --git a/src/debug/track_spans.ml b/src/debug/track_spans.ml
new file mode 100644
index 0000000..5e389ab
--- /dev/null
+++ b/src/debug/track_spans.ml
@@ -0,0 +1,120 @@
+open Trace_core
+
+let spf = Printf.sprintf
+let ( let@ ) = ( @@ )
+
+type unclosed_spans = {
+  num: int;
+  by_name: (string * int) list;
+}
+
+type out =
+  [ `Out of out_channel
+  | `Call of unclosed_spans -> unit
+  ]
+
+module type TRACKED_SPAN = sig
+  include Hashtbl.HashedType
+
+  val of_span : Trace_core.span -> t option
+  val name : t -> string
+end
+
+module Make_cbs (X : sig
+  module T_span : TRACKED_SPAN
+
+  type st
+
+  val cbs : st Collector.Callbacks.t
+  val out : out
+end) =
+struct
+  module Tbl = Hashtbl.Make (X.T_span)
+
+  let mutex = Mutex.create ()
+  let tbl_open_spans : unit Tbl.t = Tbl.create 32
+
+  let with_mutex mut f =
+    Mutex.lock mut;
+    Fun.protect f ~finally:(fun () -> Mutex.unlock mut)
+
+  let enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~params ~data ~parent name
+      : span =
+    let span =
+      X.cbs.enter_span st ~__FUNCTION__ ~__FILE__ ~__LINE__ ~params ~data
+        ~parent name
+    in
+
+    (match X.T_span.of_span span with
+    | None -> ()
+    | Some t_span ->
+      let@ () = with_mutex mutex in
+      Tbl.add tbl_open_spans t_span ());
+    span
+
+  let exit_span st span =
+    (match X.T_span.of_span span with
+    | None -> ()
+    | Some t_span ->
+      let@ () = with_mutex mutex in
+      Tbl.remove tbl_open_spans t_span);
+    X.cbs.exit_span st span
+
+  let emit (us : unclosed_spans) =
+    assert (us.by_name <> []);
+    match X.out with
+    | `Call f -> f us
+    | `Out out ->
+      Printf.fprintf out "trace: warning: %d spans were not closed" us.num;
+      List.iter
+        (fun (name, n) ->
+          Printf.fprintf out "  span %S was not closed (%d occurrences)" name n)
+        us.by_name;
+      flush out
+
+  let print_non_closed_spans_warning () =
+    let module Str_map = Map.Make (String) in
+    let@ () = with_mutex mutex in
+
+    let num = Tbl.length tbl_open_spans in
+    if num > 0 then (
+      let names_with_count =
+        Tbl.fold
+          (fun span () m ->
+            let name = X.T_span.name span in
+            Str_map.add name
+              (1 + try Str_map.find name m with Not_found -> 0)
+              m)
+          tbl_open_spans Str_map.empty
+      in
+      let unclosed_spans =
+        { num; by_name = Str_map.to_list names_with_count }
+      in
+      emit unclosed_spans
+    )
+
+  let shutdown st : unit =
+    print_non_closed_spans_warning ();
+    X.cbs.shutdown st
+
+  let new_callbacks : _ Collector.Callbacks.t =
+    { X.cbs with enter_span; exit_span; shutdown }
+end
+
+let track_ (type state) ~(out : out) (module T : TRACKED_SPAN) (st : state)
+    (cbs : state Collector.Callbacks.t) : Collector.t =
+  let module CBS = Make_cbs (struct
+    module T_span = T
+
+    type st = state
+
+    let cbs = cbs
+    let out = out
+  end) in
+  C_some (st, CBS.new_callbacks)
+
+let track ?(on_lingering_spans = `Out stderr) (module T : TRACKED_SPAN)
+    (c : Collector.t) : Collector.t =
+  match c with
+  | C_none -> C_none
+  | C_some (st, cbs) -> track_ ~out:on_lingering_spans (module T) st cbs
diff --git a/src/debug/track_spans.mli b/src/debug/track_spans.mli
new file mode 100644
index 0000000..2218a28
--- /dev/null
+++ b/src/debug/track_spans.mli
@@ -0,0 +1,24 @@
+open Trace_core
+
+module type TRACKED_SPAN = sig
+  include Hashtbl.HashedType
+
+  val of_span : Trace_core.span -> t option
+
+  val name : t -> string
+  (** Just the name of the span, nothing else *)
+end
+
+type unclosed_spans = {
+  num: int;
+  by_name: (string * int) list;
+}
+
+val track :
+  ?on_lingering_spans:[ `Out of out_channel | `Call of unclosed_spans -> unit ] ->
+  (module TRACKED_SPAN) ->
+  Collector.t ->
+  Collector.t
+(** Modify the enter/exit span functions to track the set of spans that are
+    open, and warn at the end if some are not closed.
+    @param on_lingering_spans what to do with the non-closed spans *)