From 9ecd73cb26e63050627286abbe81a8ef326f15d6 Mon Sep 17 00:00:00 2001
From: Elliott Cable <elliott.cable@ahrefs.com>
Date: Wed, 30 Aug 2023 02:28:34 +0000
Subject: [PATCH] trace-collector: Add static names for well-known strings

---
 src/trace/opentelemetry_trace.ml  | 62 +++++++++++++++++++------------
 src/trace/opentelemetry_trace.mli | 35 ++++++++++++-----
 2 files changed, 64 insertions(+), 33 deletions(-)

diff --git a/src/trace/opentelemetry_trace.ml b/src/trace/opentelemetry_trace.ml
index 6873fdd7..bc9bb87b 100644
--- a/src/trace/opentelemetry_trace.ml
+++ b/src/trace/opentelemetry_trace.ml
@@ -2,6 +2,45 @@ module Otel = Opentelemetry
 module Otrace = Trace (* ocaml-trace *)
 module TLS = Ambient_context_tls.Thread_local
 
+module Well_known = struct
+  let spankind_key = "otrace.spankind"
+
+  let internal = `String "INTERNAL"
+
+  let server = `String "SERVER"
+
+  let client = `String "CLIENT"
+
+  let producer = `String "PRODUCER"
+
+  let consumer = `String "CONSUMER"
+
+  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 =
+    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
+end
+
+open Well_known
+
 module Internal = struct
   type span_begin = {
     id: Otel.Span_id.t;
@@ -45,29 +84,6 @@ module Internal = struct
     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 =
-    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
diff --git a/src/trace/opentelemetry_trace.mli b/src/trace/opentelemetry_trace.mli
index 02b2d28d..e7835195 100644
--- a/src/trace/opentelemetry_trace.mli
+++ b/src/trace/opentelemetry_trace.mli
@@ -13,20 +13,20 @@ module TLS := Ambient_context_tls.Thread_local
     and implicit scope (in {!Internal.M.with_span}, via {!Ambient_context}) are
     supported; see the detailed notes on {!Internal.M.enter_manual_span}.
 
-    {1 Well-known identifiers}
+    {1:wellknown Well-known identifiers}
 
     Because [ocaml-trace]'s API is a subset of OpenTelemetry functionality, this
     interface allows for a few 'well-known' identifiers to be used in
     [Trace]-instrumented libraries that wish to further support OpenTelemetry
     usage:
 
-    - If a key of ["otrace.spankind"] is included in the {!Trace.user_data}
-      passed to [with_span] et al., it will be used as the
+    - If a key of exactly ["otrace.spankind"] is included in the
+      {!Trace.user_data} passed to [with_span] et al., it will be used as the
       {!Opentelemetry.Span.kind} of the emitted span. (See
       {!Internal.spankind_of_string} for the list of supported values.)
 
       {[ocaml
-      let describe () = [ "otrace.spankind", `String "CLIENT" ] in
+      let describe () = [ Opentelemetry_trace.(spankind_key, client) ] in
       Trace.with_span ~__FILE__ ~__LINE__ ~data:describe "my-span" @@ fun _ ->
         (* ... *)
       ]}
@@ -41,6 +41,27 @@ val setup_with_otel_backend : Opentelemetry.Collector.backend -> unit
 val collector : unit -> Trace.collector
 (** Make a Trace collector that uses the OTEL backend to send spans and logs *)
 
+(** Static references for well-known identifiers; see {!label-wellknown}. *)
+module Well_known : sig
+  val spankind_key : string
+
+  val internal : Otrace.user_data
+
+  val server : Otrace.user_data
+
+  val client : Otrace.user_data
+
+  val producer : Otrace.user_data
+
+  val consumer : Otrace.user_data
+
+  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
+end
+
 (** Internal implementation details; do not consider these stable. *)
 module Internal : sig
   module M : sig
@@ -150,12 +171,6 @@ module Internal : sig
 
   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 ->