trace-collector: Support manual spans

src/trace/opentelemetry_trace.ml

@@ -1,45 +1,145 @@
 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 collector () : Trace.collector =
+let spankind_of_string =
-  let module M = struct
+  let open Otel.Span in
-    let with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name cb =
+  function
-      let span_id = Otel.Span_id.create () in
+  | "INTERNAL" -> Span_kind_internal
-      let span = conv_span_to_i64 span_id in
+  | "SERVER" -> Span_kind_server
+  | "CLIENT" -> Span_kind_client
+  | "PRODUCER" -> Span_kind_producer
+  | "CONSUMER" -> Span_kind_consumer
+  | _ -> Span_kind_unspecified
 
-      let start_time = Otel.Timestamp_ns.now_unix_ns () in
+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 old_scope = Otel.Scope.get_ambient_scope () in
+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 old_scope with
+    match parent_scope with
-        | None -> Otel.Trace_id.create ()
     | 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 =
-        { Otel.Scope.span_id; trace_id; events = []; attrs = [] }
+    { Scope.span_id = otel_id; trace_id; events = []; attrs = [] }
   in
 
-      Otel.Scope.with_ambient_scope new_scope @@ fun () ->
+  let start_time = Timestamp_ns.now_unix_ns () in
-      let rv = cb span in
 
-      let end_time = Otel.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 o_span : Otel.Span.t =
   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
@@ -53,25 +153,51 @@ let collector () : Trace.collector =
         "code.function", `String function_name;
         "code.namespace", `String module_path;
       ]
+      @ attrs
   in
-        Otel.Span.create ~trace_id:new_scope.trace_id ~id:span_id ~start_time
+  Span.create ~kind ~trace_id ?parent:parent_id ~id:otel_id ~start_time
-          ~end_time ~attrs:(attrs @ data) name
+    ~end_time ~attrs name
   |> fst
+
+let collector () : Trace.collector =
+  let module M = struct
+    let with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name cb =
+      let otrace_id, sb =
+        enter_span' ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
       in
 
-      Otel.Trace.emit [ o_span ];
+      Otel.Scope.with_ambient_scope sb.scope @@ fun () ->
+      let rv = cb otrace_id in
+
+      let otel_span = exit_span' otrace_id sb in
+      Otel.Trace.emit [ otel_span ];
 
       rv
 
-    let enter_manual_span ~parent:_ ~__FUNCTION__:_ ~__FILE__:_ ~__LINE__:_
+    let enter_manual_span ~(parent : Otrace.explicit_span option) ~flavor:_
-        ~data:_ _name : Trace.explicit_span =
+        ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name : Otrace.explicit_span =
-      (* TODO: print debugging warning if OTEL_OCAML_DEBUG is enabled *)
+      let otrace_id, sb =
-      Trace.Collector.dummy_explicit_span
+        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 =
+      let active_spans = Active_spans.get () in
-      (* TODO: print debugging warning if OTEL_OCAML_DEBUG is enabled *)
+      Active_span_tbl.add active_spans.tbl otrace_id sb;
-      assert (es == Trace.Collector.dummy_explicit_span);
+
+      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