add trace.thread-local-storage optional library

a basic ambient span provider.

dune-project

@@ -29,6 +29,7 @@
   dune)
  (depopts
   unix
+  (thread-local-storage (>= 0.2))
   (mtime
    (>= 2.0)))
  (tags

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 =

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 = ..

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))

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

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 *)

trace.opam

@@ -21,6 +21,7 @@ depends: [
 ]
 depopts: [
   "unix"
+  "thread-local-storage" {>= "0.2"}
   "mtime" {>= "2.0"}
 ]
 build: [