From e0a705e391f257f1bec16f5ea3635ea45699a741 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 26 Feb 2026 09:36:06 -0500
Subject: [PATCH] add trace.thread-local-storage optional library

a basic ambient span provider.
---
 dune-project                           |  1 +
 src/core/trace_core.ml                 |  3 +++
 src/core/trace_core.mli                |  8 +++++++
 src/tls/dune                           |  7 +++++++
 src/tls/trace_thread_local_storage.ml  | 29 ++++++++++++++++++++++++++
 src/tls/trace_thread_local_storage.mli | 14 +++++++++++++
 trace.opam                             |  1 +
 7 files changed, 63 insertions(+)
 create mode 100644 src/tls/dune
 create mode 100644 src/tls/trace_thread_local_storage.ml
 create mode 100644 src/tls/trace_thread_local_storage.mli

diff --git a/dune-project b/dune-project
index a9bc09b..1850f8f 100644
--- a/dune-project
+++ b/dune-project
@@ -29,6 +29,7 @@
   dune)
  (depopts
   unix
+  (thread-local-storage (>= 0.2))
   (mtime
    (>= 2.0)))
  (tags
diff --git a/src/core/trace_core.ml b/src/core/trace_core.ml
index 94b9714..446e51f 100644
--- a/src/core/trace_core.ml
+++ b/src/core/trace_core.ml
@@ -3,6 +3,7 @@ module A = Atomic_
 module Collector = Collector
 module Level = Level
 module Core_ext = Core_ext
+module Ambient_span_provider = Ambient_span_provider
 
 type collector = Collector.t
 
@@ -175,6 +176,8 @@ let with_setup_collector c f =
   setup_collector c;
   Fun.protect ~finally:shutdown f
 
+let set_ambient_context_provider p = A.set ambient_span_provider p
+
 type extension_event = Types.extension_event = ..
 
 let[@inline] extension_event ?(level = A.get default_level_) ev : unit =
diff --git a/src/core/trace_core.mli b/src/core/trace_core.mli
index 6c032fc..fd8215c 100644
--- a/src/core/trace_core.mli
+++ b/src/core/trace_core.mli
@@ -7,6 +7,7 @@
 include module type of Types
 module Collector = Collector
 module Level = Level
+module Ambient_span_provider = Ambient_span_provider
 
 (**/**)
 
@@ -224,6 +225,13 @@ val with_setup_collector : Collector.t -> (unit -> 'a) -> 'a
     [f()] is done.
     @since 0.11 *)
 
+(** {2 ambient span provider} *)
+
+val set_ambient_context_provider : Ambient_span_provider.t -> unit
+(** Install a provider for {!current_span} and {!with_current_span_set_to}. The
+    default provider does nothing (ie [current_span ()] is always [None]).
+    @since NEXT_RELEASE *)
+
 (** {2 Extensions} *)
 
 type extension_event = Types.extension_event = ..
diff --git a/src/tls/dune b/src/tls/dune
new file mode 100644
index 0000000..f885110
--- /dev/null
+++ b/src/tls/dune
@@ -0,0 +1,7 @@
+
+(library
+  (name trace_thread_local_storage)
+  (public_name trace.thread-local-storage)
+  (synopsis "Use thread-local-storage for ambient spans")
+  (optional) ; thread-local-storage
+  (libraries trace.core thread-local-storage))
diff --git a/src/tls/trace_thread_local_storage.ml b/src/tls/trace_thread_local_storage.ml
new file mode 100644
index 0000000..24787ea
--- /dev/null
+++ b/src/tls/trace_thread_local_storage.ml
@@ -0,0 +1,29 @@
+open Trace_core
+
+let k_span : span Thread_local_storage.t = Thread_local_storage.create ()
+
+open struct
+  let get_current_span () = Thread_local_storage.get_opt k_span
+
+  let with_current_span_set_to () span f =
+    let prev_span =
+      try Thread_local_storage.get_exn k_span
+      with Thread_local_storage.Not_set -> Collector.dummy_span
+    in
+    Thread_local_storage.set k_span span;
+
+    match f span with
+    | res ->
+      Thread_local_storage.set k_span prev_span;
+      res
+    | exception exn ->
+      let bt = Printexc.get_raw_backtrace () in
+      Thread_local_storage.set k_span prev_span;
+      Printexc.raise_with_backtrace exn bt
+
+  let callbacks : unit Ambient_span_provider.Callbacks.t =
+    { get_current_span; with_current_span_set_to }
+end
+
+let provider : Ambient_span_provider.t = ASP_some ((), callbacks)
+let setup () = Trace_core.set_ambient_context_provider provider
diff --git a/src/tls/trace_thread_local_storage.mli b/src/tls/trace_thread_local_storage.mli
new file mode 100644
index 0000000..e0bae7b
--- /dev/null
+++ b/src/tls/trace_thread_local_storage.mli
@@ -0,0 +1,14 @@
+(** use [thread-local-storage] to store ambient spans.
+
+    This doesn't work with cooperative concurrency (Eio, Lwt, etc) but is fine
+    in a threaded context. *)
+
+open Trace_core
+
+val k_span : span Thread_local_storage.t
+(** Key to access the current span *)
+
+val provider : Ambient_span_provider.t
+
+val setup : unit -> unit
+(** Install the provider *)
diff --git a/trace.opam b/trace.opam
index 81473ad..f7922ed 100644
--- a/trace.opam
+++ b/trace.opam
@@ -21,6 +21,7 @@ depends: [
 ]
 depopts: [
   "unix"
+  "thread-local-storage" {>= "0.2"}
   "mtime" {>= "2.0"}
 ]
 build: [