From f88cd7651cbfafe74682326eac9e84fd72b87ad3 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Sat, 17 Jan 2026 20:54:07 -0500
Subject: [PATCH] trace-tef: ?debug option to track spans

---
 src/tef/dune          |  9 ++++++++-
 src/tef/trace_tef.ml  | 31 +++++++++++++++++++++----------
 src/tef/trace_tef.mli |  9 ++++++---
 3 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/src/tef/dune b/src/tef/dune
index 2f3b5a8..ab476ae 100644
--- a/src/tef/dune
+++ b/src/tef/dune
@@ -3,4 +3,11 @@
  (public_name trace-tef)
  (synopsis
   "Simple and lightweight tracing using TEF/Catapult format, in-process")
- (libraries trace.core trace.util mtime mtime.clock.os unix threads))
+ (libraries
+  trace.core
+  trace.util
+  trace.debug
+  mtime
+  mtime.clock.os
+  unix
+  threads))
diff --git a/src/tef/trace_tef.ml b/src/tef/trace_tef.ml
index d26b78a..ca3a6b9 100644
--- a/src/tef/trace_tef.ml
+++ b/src/tef/trace_tef.ml
@@ -76,27 +76,38 @@ open struct
       )
 end
 
-let setup ?(out = `Env) () =
+let setup ?(debug = false) ?(out = `Env) () =
   register_atexit ();
+
+  let setup_col c =
+    let c =
+      if debug then
+        Trace_debug.Track_spans.track c
+      else
+        c
+    in
+    Trace_core.setup_collector c
+  in
+
   match out with
-  | `Stderr -> Trace_core.setup_collector @@ collector ~out:`Stderr ()
-  | `Stdout -> Trace_core.setup_collector @@ collector ~out:`Stdout ()
-  | `File path -> Trace_core.setup_collector @@ collector ~out:(`File path) ()
+  | `Stderr -> setup_col @@ collector ~out:`Stderr ()
+  | `Stdout -> setup_col @@ collector ~out:`Stdout ()
+  | `File path -> setup_col @@ collector ~out:(`File path) ()
   | `Env ->
     (match Sys.getenv_opt "TRACE" with
     | Some ("1" | "true") ->
       let path = "trace.json" in
       let c = collector ~out:(`File path) () in
-      Trace_core.setup_collector c
-    | Some "stdout" -> Trace_core.setup_collector @@ collector ~out:`Stdout ()
-    | Some "stderr" -> Trace_core.setup_collector @@ collector ~out:`Stderr ()
+      setup_col c
+    | Some "stdout" -> setup_col @@ collector ~out:`Stdout ()
+    | Some "stderr" -> setup_col @@ collector ~out:`Stderr ()
     | Some path ->
       let c = collector ~out:(`File path) () in
-      Trace_core.setup_collector c
+      setup_col c
     | None -> ())
 
-let with_setup ?out () f =
-  setup ?out ();
+let with_setup ?debug ?out () f =
+  setup ?debug ?out ();
   Fun.protect ~finally:Trace_core.shutdown f
 
 module Private_ = struct
diff --git a/src/tef/trace_tef.mli b/src/tef/trace_tef.mli
index b66e3fa..b796d52 100644
--- a/src/tef/trace_tef.mli
+++ b/src/tef/trace_tef.mli
@@ -19,7 +19,7 @@ val collector : out:[< output ] -> unit -> Trace_core.collector
 (** Make a collector that writes into the given output. See {!setup} for more
     details. *)
 
-val setup : ?out:[ output | `Env ] -> unit -> unit
+val setup : ?debug:bool -> ?out:[ output | `Env ] -> unit -> unit
 (** [setup ()] installs the collector depending on [out].
 
     @param out
@@ -31,9 +31,12 @@ val setup : ?out:[ output | `Env ] -> unit -> unit
     - If it's set to "stdout", then logging happens on stdout (since 0.2)
     - If it's set to "stderr", then logging happens on stdout (since 0.2)
     - Otherwise, if it's set to a non empty string, the value is taken to be the
-      file path into which to write. *)
+      file path into which to write.
 
-val with_setup : ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
+    @param debug if true, use {!Trace_debug}. Default [false]. *)
+
+val with_setup :
+  ?debug:bool -> ?out:[ output | `Env ] -> unit -> (unit -> 'a) -> 'a
 (** [with_setup () f] (optionally) sets a collector up, calls [f()], and makes
     sure to shutdown before exiting. since 0.2 a () argument was added. *)