trace-collector: Expose Internal module

src/trace/opentelemetry_trace.ml

@@ -2,7 +2,8 @@ module Otel = Opentelemetry
 module Otrace = Trace (* ocaml-trace *)
 module TLS = Ambient_context_tls.Thread_local
 
-type span_begin = {
+module Internal = struct
+  type span_begin = {
     id: Otel.Span_id.t;
     start_time: int64;
     name: string;
@@ -14,17 +15,17 @@ type span_begin = {
     scope: Otel.Scope.t;
     parent_id: Otel.Span_id.t option;
     parent_scope: Otel.Scope.t option;
-}
+  }
 
-(** Table indexed by ocaml-trace spans *)
+  (** Table indexed by ocaml-trace spans *)
-module Active_span_tbl = Hashtbl.Make (struct
+  module Active_span_tbl = Hashtbl.Make (struct
     include Int64
 
     let hash : t -> int = Hashtbl.hash
-end)
+  end)
 
-(** Per-thread set of active spans. *)
+  (** Per-thread set of active spans. *)
-module Active_spans = struct
+  module Active_spans = struct
     type t = { tbl: span_begin Active_span_tbl.t } [@@unboxed]
 
     let create () : t = { tbl = Active_span_tbl.create 32 }
@@ -32,20 +33,20 @@ module Active_spans = struct
     let tls : t TLS.t = TLS.create ()
 
     let[@inline] get () : t = TLS.get_or_create tls ~create
-end
+  end
 
-let otrace_of_otel (id : Otel.Span_id.t) : int64 =
+  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 otel_of_otrace (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 spankind_of_string =
     let open Otel.Span in
     function
     | "INTERNAL" -> Span_kind_internal
@@ -55,7 +56,7 @@ let spankind_of_string =
     | "CONSUMER" -> Span_kind_consumer
     | _ -> Span_kind_unspecified
 
-let otel_attrs_of_otrace_data (data : Otel.Span.key_value list) =
+  let otel_attrs_of_otrace_data data =
     let kind : Otel.Span.kind ref = ref Otel.Span.Span_kind_unspecified in
     let data =
       List.filter_map
@@ -68,7 +69,8 @@ let otel_attrs_of_otrace_data (data : Otel.Span.key_value list) =
     in
     !kind, data
 
-let enter_span' ?explicit_parent ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name =
+  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
@@ -113,7 +115,7 @@ let enter_span' ?explicit_parent ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name =
 
     otrace_id, sb
 
-let exit_span' otrace_id
+  let exit_span' otrace_id
       {
         id = otel_id;
         start_time;
@@ -159,8 +161,7 @@ let exit_span' otrace_id
       ~end_time ~attrs name
     |> fst
 
-let collector () : Trace.collector =
+  module M = struct
-  let module M = struct
     let with_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name cb =
       let otrace_id, sb =
         enter_span' ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data name
@@ -226,8 +227,10 @@ let collector () : Trace.collector =
     let counter_float name cur_val : unit =
       let m = Otel.Metrics.(gauge ~name [ float cur_val ]) in
       Otel.Metrics.emit [ m ]
-  end in
+  end
-  (module M)
+end
+
+let collector () : Trace.collector = (module Internal.M)
 
 let setup () = Trace.setup_collector @@ collector ()
 

src/trace/opentelemetry_trace.mli

@@ -1,8 +1,101 @@
-val collector : unit -> Trace.collector
+module Otel := Opentelemetry
-(** Make a Trace collector that uses the OTEL backend to send spans and logs *)
+module Otrace := Trace
+module TLS := Ambient_context_tls.Thread_local
 
 val setup : unit -> unit
 (** Install the OTEL backend as a Trace collector *)
 
 val setup_with_otel_backend : Opentelemetry.Collector.backend -> unit
 (** Same as {!setup},  but also install the given backend as OTEL backend *)
+
+val collector : unit -> Trace.collector
+(** Make a Trace collector that uses the OTEL backend to send spans and logs *)
+
+(** Internal implementation details; do not consider these stable. *)
+module Internal : sig
+  module M : sig
+    val with_span :
+      __FUNCTION__:string option ->
+      __FILE__:string ->
+      __LINE__:int ->
+      data:(string * Otrace.user_data) list ->
+      string (* span name *) ->
+      (Otrace.span -> 'a) ->
+      'a
+
+    val enter_manual_span :
+      parent:Otrace.explicit_span option ->
+      flavor:'a ->
+      __FUNCTION__:string option ->
+      __FILE__:string ->
+      __LINE__:int ->
+      data:(string * Otrace.user_data) list ->
+      string (* span name *) ->
+      Otrace.explicit_span
+
+    val exit_manual_span : Otrace.explicit_span -> unit
+
+    val message :
+      ?span:Otrace.span ->
+      data:(string * Otrace.user_data) list ->
+      string ->
+      unit
+
+    val shutdown : unit -> unit
+
+    val name_process : string -> unit
+
+    val name_thread : string -> unit
+
+    val counter_int : string -> int -> unit
+
+    val counter_float : string -> float -> unit
+  end
+
+  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;
+  }
+
+  module Active_span_tbl : Hashtbl.S with type key = Otrace.span
+
+  module Active_spans : sig
+    type t = private { tbl: span_begin Active_span_tbl.t } [@@unboxed]
+
+    val create : unit -> t
+
+    val tls : t TLS.t
+
+    val get : unit -> t
+  end
+
+  val otrace_of_otel : Otel.Span_id.t -> Otrace.span
+
+  val otel_of_otrace : Otrace.span -> Otel.Span_id.t
+
+  val spankind_of_string : string -> Otel.Span.kind
+
+  val otel_attrs_of_otrace_data :
+    (string * Otrace.user_data) list ->
+    Otel.Span.kind * Otel.Span.key_value list
+
+  val enter_span' :
+    ?explicit_parent:Otrace.span ->
+    __FUNCTION__:string option ->
+    __FILE__:string ->
+    __LINE__:int ->
+    data:(string * Otrace.user_data) list ->
+    string ->
+    Otrace.span * span_begin
+
+  val exit_span' : Otrace.span -> span_begin -> Otel.Span.t
+end