trace-collector: Support manual spans

This commit is contained in:
Elliott Cable 2023-08-29 00:11:48 +00:00
parent de35c80799
commit 1822c1acaa

View file

@ -1,77 +1,203 @@
module Otel = Opentelemetry
module Otrace = Trace (* ocaml-trace *)
module TLS = Ambient_context_tls.Thread_local
let conv_span_to_i64 (id : Otel.Span_id.t) : int64 =
type span_begin = {
id: Otel.Span_id.t;
start_time: int64;
name: string;
data: (string * Otrace.user_data) list;
__FILE__: string;
__LINE__: int;
__FUNCTION__: string option;
trace_id: Otel.Trace_id.t;
scope: Otel.Scope.t;
parent_id: Otel.Span_id.t option;
parent_scope: Otel.Scope.t option;
}
(** Table indexed by ocaml-trace spans *)
module Active_span_tbl = Hashtbl.Make (struct
include Int64
let hash : t -> int = Hashtbl.hash
end)
(** Per-thread set of active spans. *)
module Active_spans = struct
type t = { tbl: span_begin Active_span_tbl.t } [@@unboxed]
let create () : t = { tbl = Active_span_tbl.create 32 }
let tls : t TLS.t = TLS.create ()
let[@inline] get () : t = TLS.get_or_create tls ~create
end
let otrace_of_otel (id : Otel.Span_id.t) : int64 =
let bs = Otel.Span_id.to_bytes id in
(* lucky that it coincides! *)
assert (Bytes.length bs = 8);
Bytes.get_int64_le bs 0
let span_of_i64 (id : int64) : Otel.Span_id.t =
let otel_of_otrace (id : int64) : Otel.Span_id.t =
let bs = Bytes.create 8 in
Bytes.set_int64_le bs 0 id;
Otel.Span_id.of_bytes bs
let spankind_of_string =
let open Otel.Span in
function
| "INTERNAL" -> Span_kind_internal
| "SERVER" -> Span_kind_server
| "CLIENT" -> Span_kind_client
| "PRODUCER" -> Span_kind_producer
| "CONSUMER" -> Span_kind_consumer
| _ -> Span_kind_unspecified
let otel_attrs_of_otrace_data (data : Otel.Span.key_value list) =
let kind : Otel.Span.kind ref = ref Otel.Span.Span_kind_unspecified in
let data =
List.filter_map
(function
| name, `String v when name = "otrace.spankind" ->
kind := spankind_of_string v;
None
| x -> Some x)
data
in
!kind, data
let enter_span' ?explicit_parent ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name =
let open Otel in
let otel_id = Span_id.create () in
let otrace_id = otrace_of_otel otel_id in
let parent_scope = Scope.get_ambient_scope () in
let trace_id =
match parent_scope with
| Some sc -> sc.trace_id
| None -> Trace_id.create ()
in
let parent_id =
match explicit_parent, parent_scope with
| Some p, _ -> Some (otel_of_otrace p)
| None, Some parent -> Some parent.span_id
| None, None -> None
in
let new_scope =
{ Scope.span_id = otel_id; trace_id; events = []; attrs = [] }
in
let start_time = Timestamp_ns.now_unix_ns () in
let sb =
{
id = otel_id;
start_time;
name;
data;
__FILE__;
__LINE__;
__FUNCTION__;
trace_id;
scope = new_scope;
parent_id;
parent_scope;
}
in
let active_spans = Active_spans.get () in
Active_span_tbl.add active_spans.tbl otrace_id sb;
otrace_id, sb
let exit_span' otrace_id
{
id = otel_id;
start_time;
name;
data;
__FILE__;
__LINE__;
__FUNCTION__;
trace_id;
scope = _;
parent_id;
parent_scope = _;
} =
let open Otel in
let active_spans = Active_spans.get () in
Active_span_tbl.remove active_spans.tbl otrace_id;
let end_time = Timestamp_ns.now_unix_ns () in
let kind, attrs = otel_attrs_of_otrace_data data in
let attrs =
match __FUNCTION__ with
| None ->
[ "code.filepath", `String __FILE__; "code.lineno", `Int __LINE__ ]
@ attrs
| Some __FUNCTION__ ->
let last_dot = String.rindex __FUNCTION__ '.' in
let module_path = String.sub __FUNCTION__ 0 last_dot in
let function_name =
String.sub __FUNCTION__ (last_dot + 1)
(String.length __FUNCTION__ - last_dot - 1)
in
[
"code.filepath", `String __FILE__;
"code.lineno", `Int __LINE__;
"code.function", `String function_name;
"code.namespace", `String module_path;
]
@ attrs
in
Span.create ~kind ~trace_id ?parent:parent_id ~id:otel_id ~start_time
~end_time ~attrs name
|> fst
let collector () : Trace.collector =
let module M = struct
let with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name cb =
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_ambient_scope () in
let trace_id =
match old_scope with
| None -> Otel.Trace_id.create ()
| Some sc -> sc.trace_id
let otrace_id, sb =
enter_span' ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
in
let new_scope =
{ Otel.Scope.span_id; trace_id; events = []; attrs = [] }
in
Otel.Scope.with_ambient_scope sb.scope @@ fun () ->
let rv = cb otrace_id in
Otel.Scope.with_ambient_scope new_scope @@ fun () ->
let rv = cb span in
let end_time = Otel.Timestamp_ns.now_unix_ns () in
let o_span : Otel.Span.t =
let attrs =
match __FUNCTION__ with
| None ->
[ "code.filepath", `String __FILE__; "code.lineno", `Int __LINE__ ]
| Some __FUNCTION__ ->
let last_dot = String.rindex __FUNCTION__ '.' in
let module_path = String.sub __FUNCTION__ 0 last_dot in
let function_name =
String.sub __FUNCTION__ (last_dot + 1)
(String.length __FUNCTION__ - last_dot - 1)
in
[
"code.filepath", `String __FILE__;
"code.lineno", `Int __LINE__;
"code.function", `String function_name;
"code.namespace", `String module_path;
]
in
Otel.Span.create ~trace_id:new_scope.trace_id ~id:span_id ~start_time
~end_time ~attrs:(attrs @ data) name
|> fst
in
Otel.Trace.emit [ o_span ];
let otel_span = exit_span' otrace_id sb in
Otel.Trace.emit [ otel_span ];
rv
let enter_manual_span ~parent:_ ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_
~data:_ _name : Trace.explicit_span =
(* TODO: print debugging warning if OTEL_OCAML_DEBUG is enabled *)
Trace.Collector.dummy_explicit_span
let enter_manual_span ~(parent : Otrace.explicit_span option) ~flavor:_
~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name : Otrace.explicit_span =
let otrace_id, sb =
match parent with
| None -> enter_span' ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
| Some { span; _ } ->
enter_span' ~explicit_parent:span ~__FUNCTION__ ~__FILE__ ~__LINE__
~data name
in
let exit_manual_span es =
(* TODO: print debugging warning if OTEL_OCAML_DEBUG is enabled *)
assert (es == Trace.Collector.dummy_explicit_span);
()
let active_spans = Active_spans.get () in
Active_span_tbl.add active_spans.tbl otrace_id sb;
Otrace.{ span = otrace_id; meta = Meta_map.empty }
let exit_manual_span Otrace.{ span = otrace_id; _ } =
let active_spans = Active_spans.get () in
match Active_span_tbl.find_opt active_spans.tbl otrace_id with
| None ->
(* FIXME: some kind of error/debug logging *)
()
| Some sb ->
let otel_span = exit_span' otrace_id sb in
Otel.Trace.emit [ otel_span ]
let message ?span ~data:_ msg : unit =
(* gather information from context *)
@ -80,7 +206,7 @@ let collector () : Trace.collector =
let span_id =
match span with
| Some id -> Some (span_of_i64 id)
| Some id -> Some (otel_of_otrace id)
| None -> Option.map (fun sc -> sc.Otel.Scope.span_id) old_scope
in