diff --git a/src/lwt/opentelemetry_lwt.ml b/src/lwt/opentelemetry_lwt.ml index 8e8524cf..650b8016 100644 --- a/src/lwt/opentelemetry_lwt.ml +++ b/src/lwt/opentelemetry_lwt.ml @@ -11,13 +11,44 @@ module GC_metrics = GC_metrics module Metrics_callbacks = Metrics_callbacks module Trace_context = Trace_context +module Scope = struct + include Scope + + (**/**) + + let _global_scope_key : t Lwt.key = Lwt.new_key () + + (**/**) + + let get_surrounding ?scope () : t option = + let surrounding = Lwt.get _global_scope_key in + match scope, surrounding with + | Some _, _ -> scope + | None, Some _ -> surrounding + | None, None -> None + + let[@inline] with_scope (sc : t) (f : unit -> 'a) : 'a = + Lwt.with_value _global_scope_key (Some sc) f +end + +open struct + let get_surrounding_scope = Scope.get_surrounding +end + module Trace = struct open Proto.Trace include Trace (** Sync span guard *) - let with_ ?trace_state ?service_name ?(attrs = []) ?kind ?trace_id ?parent - ?scope ?links name (f : Scope.t -> 'a Lwt.t) : 'a Lwt.t = + let with_ ?(force_new_trace_id = false) ?trace_state ?service_name + ?(attrs = []) ?kind ?trace_id ?parent ?scope ?links name + (f : Scope.t -> 'a Lwt.t) : 'a Lwt.t = + let scope = + if force_new_trace_id then + None + else + get_surrounding_scope ?scope () + in let trace_id = match trace_id, scope with | Some trace_id, _ -> trace_id @@ -33,6 +64,8 @@ module Trace = struct let start_time = Timestamp_ns.now_unix_ns () in let span_id = Span_id.create () in let scope = { trace_id; span_id; events = []; attrs } in + (* set global scope in this thread *) + Scope.with_scope scope @@ fun () -> let finally ok = let status = match ok with