add ambient_span_provider to get cur_span/with_cur_span

src/core/ambient_span_provider.ml

@@ -0,0 +1,20 @@
+(** Access/set the current span from some ambient context.
+    @since NEXT_RELEASE *)
+
+open Types
+
+module Callbacks = struct
+  type 'st t = {
+    with_current_span_set_to: 'a. 'st -> span -> (span -> 'a) -> 'a;
+        (** [with_current_span_set_to span f] sets the span as current span,
+            enters [f span], and restores the previous current span if any *)
+    get_current_span: 'st -> span option;
+        (** Access the current span from some ambient scope. This is only
+            supported for collectors that provide a [current_span_wrap] field.
+        *)
+  }
+end
+
+type t =
+  | ASP_none
+  | ASP_some : 'st * 'st Callbacks.t -> t

src/core/collector.ml

@@ -82,8 +82,7 @@ end
     This is only relevant to implementors of tracing backends; to instrument
     your code you only need to look at the {!Trace} module.
 
-    The definition changed since 0.11 to a record of callbacks + a state
+    The definition changed since 0.11 to a record of callbacks + a state *)
-*)
 type t =
   | C_none  (** No collector. *)
   | C_some : 'st * 'st Callbacks.t -> t

src/core/trace_core.ml

@@ -15,6 +15,10 @@ let collector : collector A.t = A.make Collector.C_none
 let default_level_ = A.make Level.Trace
 let current_level_ = A.make Level.Trace
 
+(** Global provider of span context *)
+let ambient_span_provider : Ambient_span_provider.t A.t =
+  A.make Ambient_span_provider.ASP_none
+
 (* ## implementation ## *)
 
 let data_empty_build_ () = []
@@ -27,6 +31,16 @@ let[@inline] get_current_level () = A.get current_level_
 let[@inline] check_level_ ~level st (cbs : _ Collector.Callbacks.t) : bool =
   Level.leq level (A.get current_level_) && cbs.enabled st level
 
+let[@inline] current_span () =
+  match A.get ambient_span_provider with
+  | ASP_none -> None
+  | ASP_some (st, cbs) -> cbs.get_current_span st
+
+let[@inline] with_current_span_set_to sp f =
+  match A.get ambient_span_provider with
+  | ASP_none -> f sp
+  | ASP_some (st, cbs) -> cbs.with_current_span_set_to st sp f
+
 let parent_of_span_opt_opt = function
   | None -> P_unknown
   | Some None -> P_none
@@ -46,7 +60,10 @@ let with_span_collector_ st (cbs : _ Collector.Callbacks.t) ?__FUNCTION__
     enter_span_st st cbs ?__FUNCTION__ ~__FILE__ ~__LINE__ ~level ?parent
       ?params ?data name
   in
-  match f sp with
+  match
+    (* set [sp] as current span before calling [f sp] *)
+    with_current_span_set_to sp f
+  with
   | res ->
     cbs.exit_span st sp;
     res

src/core/trace_core.mli

@@ -180,6 +180,22 @@ val counter_float :
       {!set_default_level}.
     @param data metadata for this metric (since 0.4) *)
 
+val current_span : unit -> span option
+(** Access the current span from some ambient scope, {b if supported}. This is
+    only supported if a {!Ambient_span_provider} has been set up.
+    @since NEXT_RELEASE *)
+
+val with_current_span_set_to : span -> (span -> 'a) -> 'a
+(** [with_current_span_set_to span f] sets the span as current span, enters
+    [f span], and restores the previous current span (if any).
+
+    This is only supported if a {!Ambient_span_provider} has been set up,
+    otherwise it is a no-op.
+
+    Automatically called by {!with_span}.
+
+    @since NEXT_RELEASE *)
+
 (** {2 Collector} *)
 
 type collector = Collector.t