From 927d1e3e648ecf83598e2d73ad80299d11fa38f0 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Wed, 9 Aug 2023 12:52:31 -0400 Subject: [PATCH] adapt to trace 0.3; fix missing "parent" argument at scope creation --- src/trace/dune | 2 +- src/trace/opentelemetry_trace.ml | 116 +++++++++++++++++++----------- src/trace/opentelemetry_trace.mli | 2 +- 3 files changed, 77 insertions(+), 43 deletions(-) diff --git a/src/trace/dune b/src/trace/dune index 8391f103..6b0294ef 100644 --- a/src/trace/dune +++ b/src/trace/dune @@ -4,4 +4,4 @@ (public_name opentelemetry.trace) (synopsis "Use opentelemetry as a collector for trace") (optional) - (libraries trace opentelemetry)) + (libraries trace.core opentelemetry)) diff --git a/src/trace/opentelemetry_trace.ml b/src/trace/opentelemetry_trace.ml index 2670429c..fa01124b 100644 --- a/src/trace/opentelemetry_trace.ml +++ b/src/trace/opentelemetry_trace.ml @@ -1,7 +1,8 @@ module Otel = Opentelemetry module TLS = Otel.Thread_local +module Meta_map = Trace_core.Meta_map -type span = Trace.span +type span = Trace_core.span (** Table indexed by Trace spans *) module Span_tbl = Hashtbl.Make (struct @@ -14,9 +15,10 @@ end) module Active_spans = struct type span_begin = { span_id: Otel.Span_id.t; + parent_span: Otel.Span_id.t option; start_time: int64; name: string; - data: (string * Trace.user_data) list; + data: (string * Trace_core.user_data) list; __FILE__: string; __LINE__: int; new_scope: Otel.Scope.t; @@ -45,19 +47,28 @@ let span_of_i64 (id : int64) : Otel.Span_id.t = Bytes.set_int64_le bs 0 id; Otel.Span_id.of_bytes bs -let collector () : Trace.collector = +let k_begin_span : Active_spans.span_begin Meta_map.Key.t = + Meta_map.Key.create () + +let collector () : Trace_core.collector = let module M = struct - let enter_span ?__FUNCTION__:_ ~__FILE__ ~__LINE__ ~data name : span = + let enter_span_ ~__FILE__ ~__LINE__ + ?(parent : Trace_core.explicit_span option) ~data name : + span * Active_spans.span_begin = let span_id = Otel.Span_id.create () in let span = conv_span_to_i64 span_id in let start_time = Otel.Timestamp_ns.now_unix_ns () in - let old_scope = Otel.Scope.get_surrounding () in - let trace_id = - match old_scope with - | None -> Otel.Trace_id.create () - | Some sc -> sc.trace_id + let parent_span, trace_id, old_scope = + match parent with + | Some p -> + let bsp = Meta_map.find_exn k_begin_span p.meta in + Some bsp.span_id, bsp.new_scope.trace_id, None + | None -> + (match Otel.Scope.get_surrounding () with + | None -> None, Otel.Trace_id.create (), None + | Some sc -> Some sc.span_id, sc.trace_id, Some sc) in let new_scope = @@ -65,55 +76,78 @@ let collector () : Trace.collector = in TLS.set Otel.Scope._global_scope new_scope; - let active_spans = Active_spans.get () in - Span_tbl.add active_spans.tbl span + ( span, { span_id; start_time; + parent_span; __FILE__; __LINE__; old_scope; new_scope; name; data; - }; + } ) + let enter_span ?__FUNCTION__:_ ~__FILE__ ~__LINE__ ~data name : span = + let span, bsp = enter_span_ ~__FILE__ ~__LINE__ ~data name in + let active_spans = Active_spans.get () in + Span_tbl.add active_spans.tbl span bsp; span + let enter_manual_span ~parent ~flavor:_ ~__FUNCTION__:_ ~__FILE__ ~__LINE__ + ~data name : Trace_core.explicit_span = + let span, bsp = enter_span_ ~__FILE__ ~__LINE__ ?parent ~data name in + { Trace_core.span; meta = Meta_map.(empty |> add k_begin_span bsp) } + + let exit_span_ (bsp : Active_spans.span_begin) : unit = + let { + Active_spans.span_id; + start_time; + parent_span; + name; + __FILE__; + __LINE__; + new_scope; + old_scope; + data; + } = + bsp + in + let end_time = Otel.Timestamp_ns.now_unix_ns () in + + (* restore previous scope *) + (match old_scope with + | None -> TLS.remove Otel.Scope._global_scope + | Some sc -> TLS.set Otel.Scope._global_scope sc); + + let o_span : Otel.Span.t = + let attrs = + [ "file", `String __FILE__; "line", `Int __LINE__ ] @ data + in + Otel.Span.create ~trace_id:new_scope.trace_id ~id:span_id ~start_time + ?parent:parent_span ~end_time ~attrs name + |> fst + in + + Otel.Trace.emit [ o_span ]; + + () + let exit_span (span : span) : unit = let active_spans = Active_spans.get () in match Span_tbl.find_opt active_spans.tbl span with | None -> () (* TODO: log warning *) - | Some - { - span_id; - start_time; - name; - __FILE__; - __LINE__; - new_scope; - old_scope; - data; - } -> - let end_time = Otel.Timestamp_ns.now_unix_ns () in + | Some bsp -> exit_span_ bsp - (* restore previous scope *) - (match old_scope with - | None -> TLS.remove Otel.Scope._global_scope - | Some sc -> TLS.set Otel.Scope._global_scope sc); + let exit_manual_span (es : Trace_core.explicit_span) : unit = + match Meta_map.find k_begin_span es.meta with + | None -> () (* TODO: log warning *) + | Some bsp -> exit_span_ bsp - let o_span : Otel.Span.t = - let attrs = - [ "file", `String __FILE__; "line", `Int __LINE__ ] @ data - in - Otel.Span.create ~trace_id:new_scope.trace_id ~id:span_id ~start_time - ~end_time ~attrs name - |> fst - in - - Otel.Trace.emit [ o_span ]; - - () + let with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name f = + let sp = enter_span ?__FUNCTION__ ~__FILE__ ~__LINE__ ~data name in + Fun.protect ~finally:(fun () -> exit_span sp) (fun () -> f sp) let message ?span ~data:_ msg : unit = (* gather information from context *) @@ -145,7 +179,7 @@ let collector () : Trace.collector = end in (module M) -let setup () = Trace.setup_collector @@ collector () +let setup () = Trace_core.setup_collector @@ collector () let setup_with_otel_backend b : unit = Otel.Collector.set_backend b; diff --git a/src/trace/opentelemetry_trace.mli b/src/trace/opentelemetry_trace.mli index 8e99befb..cf5db46a 100644 --- a/src/trace/opentelemetry_trace.mli +++ b/src/trace/opentelemetry_trace.mli @@ -1,4 +1,4 @@ -val collector : unit -> Trace.collector +val collector : unit -> Trace_core.collector (** Make a Trace collector that uses the OTEL backend to send spans and logs *) val setup : unit -> unit