Merge pull request #53 from imandra-ai/wip-add-span-ctx

add span context, rewrite w3c trace context

src/core/opentelemetry.ml

@@ -1,5 +1,9 @@
 (** Opentelemetry types and instrumentation *)
 
+open struct
+  let spf = Printf.sprintf
+end
+
 module Lock = Lock
 (** Global lock. *)
 
@@ -10,13 +14,6 @@ module AList = AList
 (** Atomic list, for internal usage
     @since NEXT_RELEASE *)
 
-open struct
-  let[@inline] result_bind x f =
-    match x with
-    | Error e -> Error e
-    | Ok x -> f x
-end
-
 (** {2 Wire format} *)
 
 module Proto = Opentelemetry_proto
@@ -216,41 +213,58 @@ module Collector = struct
       f ()
 end
 
-module Util_ = struct
+(**/**)
-  let bytes_to_hex (b : bytes) : string =
-    let i_to_hex (i : int) =
-      if i < 10 then
-        Char.chr (i + Char.code '0')
-      else
-        Char.chr (i - 10 + Char.code 'a')
-    in
 
-    let res = Bytes.create (2 * Bytes.length b) in
+module Util_ = struct
+  let int_to_hex (i : int) =
+    if i < 10 then
+      Char.chr (i + Char.code '0')
+    else
+      Char.chr (i - 10 + Char.code 'a')
+
+  let bytes_to_hex_into b res off : unit =
     for i = 0 to Bytes.length b - 1 do
       let n = Char.code (Bytes.get b i) in
-      Bytes.set res (2 * i) (i_to_hex ((n land 0xf0) lsr 4));
+      Bytes.set res ((2 * i) + off) (int_to_hex ((n land 0xf0) lsr 4));
-      Bytes.set res ((2 * i) + 1) (i_to_hex (n land 0x0f))
+      Bytes.set res ((2 * i) + 1 + off) (int_to_hex (n land 0x0f))
-    done;
+    done
+
+  let bytes_to_hex (b : bytes) : string =
+    let res = Bytes.create (2 * Bytes.length b) in
+    bytes_to_hex_into b res 0;
     Bytes.unsafe_to_string res
 
-  let bytes_of_hex (s : string) : bytes =
+  let int_of_hex = function
-    let n_of_c = function
+    | '0' .. '9' as c -> Char.code c - Char.code '0'
-      | '0' .. '9' as c -> Char.code c - Char.code '0'
+    | 'a' .. 'f' as c -> 10 + Char.code c - Char.code 'a'
-      | 'a' .. 'f' as c -> 10 + Char.code c - Char.code 'a'
+    | c -> raise (Invalid_argument (spf "invalid hex char: %C" c))
-      | _ -> raise (Invalid_argument "invalid hex char")
+
-    in
+  let bytes_of_hex_substring (s : string) off len =
-    if String.length s mod 2 <> 0 then
+    if len mod 2 <> 0 then
       raise (Invalid_argument "hex sequence must be of even length");
-    let res = Bytes.make (String.length s / 2) '\x00' in
+    let res = Bytes.make (len / 2) '\x00' in
-    for i = 0 to (String.length s / 2) - 1 do
+    for i = 0 to (len / 2) - 1 do
-      let n1 = n_of_c (String.get s (2 * i)) in
+      let n1 = int_of_hex (String.get s (off + (2 * i))) in
-      let n2 = n_of_c (String.get s ((2 * i) + 1)) in
+      let n2 = int_of_hex (String.get s (off + (2 * i) + 1)) in
       let n = (n1 lsl 4) lor n2 in
       Bytes.set res i (Char.chr n)
     done;
     res
+
+  let bytes_of_hex (s : string) : bytes =
+    bytes_of_hex_substring s 0 (String.length s)
+
+  let bytes_non_zero (self : bytes) : bool =
+    try
+      for i = 0 to Bytes.length self - 1 do
+        if Char.code (Bytes.unsafe_get self i) <> 0 then raise_notrace Exit
+      done;
+      false
+    with Exit -> true
 end
 
+(**/**)
+
 (** {2 Identifiers} *)
 
 (** Trace ID.
@@ -261,20 +275,30 @@ module Trace_id : sig
 
   val create : unit -> t
 
+  val dummy : t
+
   val pp : Format.formatter -> t -> unit
 
+  val is_valid : t -> bool
+
   val to_bytes : t -> bytes
 
   val of_bytes : bytes -> t
 
   val to_hex : t -> string
 
+  val to_hex_into : t -> bytes -> int -> unit
+
   val of_hex : string -> t
+
+  val of_hex_substring : string -> int -> t
 end = struct
   type t = bytes
 
   let to_bytes self = self
 
+  let dummy : t = Bytes.make 16 '\x00'
+
   let create () : t =
     let b = Collector.rand_bytes_16 () in
     assert (Bytes.length b = 16);
@@ -286,11 +310,18 @@ end = struct
     if Bytes.length b = 16 then
       b
     else
-      raise (Invalid_argument "trace IDs must be 16 bytes in length")
+      raise (Invalid_argument "trace ID must be 16 bytes in length")
 
-  let to_hex self = Util_.bytes_to_hex self
+  let is_valid = Util_.bytes_non_zero
 
-  let of_hex s = of_bytes (Util_.bytes_of_hex s)
+  let to_hex = Util_.bytes_to_hex
+
+  let to_hex_into = Util_.bytes_to_hex_into
+
+  let[@inline] of_hex s = of_bytes (Util_.bytes_of_hex s)
+
+  let[@inline] of_hex_substring s off =
+    of_bytes (Util_.bytes_of_hex_substring s off 32)
 
   let pp fmt t = Format.fprintf fmt "%s" (to_hex t)
 end
@@ -301,20 +332,30 @@ module Span_id : sig
 
   val create : unit -> t
 
+  val dummy : t
+
   val pp : Format.formatter -> t -> unit
 
+  val is_valid : t -> bool
+
   val to_bytes : t -> bytes
 
   val of_bytes : bytes -> t
 
   val to_hex : t -> string
 
+  val to_hex_into : t -> bytes -> int -> unit
+
   val of_hex : string -> t
+
+  val of_hex_substring : string -> int -> t
 end = struct
   type t = bytes
 
   let to_bytes self = self
 
+  let dummy : t = Bytes.make 8 '\x00'
+
   let create () : t =
     let b = Collector.rand_bytes_8 () in
     assert (Bytes.length b = 8);
@@ -322,19 +363,122 @@ end = struct
     Bytes.set b 0 (Char.unsafe_chr (Char.code (Bytes.get b 0) lor 1));
     b
 
+  let is_valid = Util_.bytes_non_zero
+
   let of_bytes b =
     if Bytes.length b = 8 then
       b
     else
       raise (Invalid_argument "span IDs must be 8 bytes in length")
 
-  let to_hex self = Util_.bytes_to_hex self
+  let to_hex = Util_.bytes_to_hex
 
-  let of_hex s = of_bytes (Util_.bytes_of_hex s)
+  let to_hex_into = Util_.bytes_to_hex_into
+
+  let[@inline] of_hex s = of_bytes (Util_.bytes_of_hex s)
+
+  let[@inline] of_hex_substring s off =
+    of_bytes (Util_.bytes_of_hex_substring s off 16)
 
   let pp fmt t = Format.fprintf fmt "%s" (to_hex t)
 end
 
+(** Span context. This bundles up a trace ID and parent ID.
+
+    https://opentelemetry.io/docs/specs/otel/trace/api/#spancontext
+    @since NEXT_RELEASE *)
+module Span_ctx : sig
+  type t
+
+  val make : trace_id:Trace_id.t -> parent_id:Span_id.t -> unit -> t
+
+  val dummy : t
+  (** Invalid span context, to be used as a placeholder *)
+
+  val is_valid : t -> bool
+
+  val trace_id : t -> Trace_id.t
+
+  val parent_id : t -> Span_id.t
+
+  val is_remote : t -> bool
+
+  val to_w3c_trace_context : t -> bytes
+
+  val of_w3c_trace_context : bytes -> (t, string) result
+
+  val of_w3c_trace_context_exn : bytes -> t
+  (** @raise Invalid_argument if parsing failed *)
+end = struct
+  (* TODO: trace flags *)
+  (* TODO: trace state *)
+
+  type t = {
+    trace_id: Trace_id.t;
+    parent_id: Span_id.t;
+    is_remote: bool;
+  }
+
+  let dummy =
+    { trace_id = Trace_id.dummy; parent_id = Span_id.dummy; is_remote = false }
+
+  let make ~trace_id ~parent_id () : t =
+    { trace_id; parent_id; is_remote = false }
+
+  let[@inline] is_valid self =
+    Trace_id.is_valid self.trace_id && Span_id.is_valid self.parent_id
+
+  let[@inline] is_remote self = self.is_remote
+
+  let[@inline] trace_id self = self.trace_id
+
+  let[@inline] parent_id self = self.parent_id
+
+  let to_w3c_trace_context (self : t) : bytes =
+    let bs = Bytes.create 55 in
+    Bytes.set bs 0 '0';
+    Bytes.set bs 1 '0';
+    Bytes.set bs 2 '-';
+    Trace_id.to_hex_into self.trace_id bs 3;
+    (* +32 *)
+    Bytes.set bs (3 + 32) '-';
+    Span_id.to_hex_into self.parent_id bs 36;
+    (* +16 *)
+    Bytes.set bs 52 '-';
+    Bytes.set bs 53 '0';
+    Bytes.set bs 54 '0';
+    bs
+
+  let of_w3c_trace_context bs : _ result =
+    try
+      if Bytes.length bs <> 55 then invalid_arg "trace context must be 55 bytes";
+      (match int_of_string_opt (Bytes.sub_string bs 0 2) with
+      | Some 0 -> ()
+      | Some n -> invalid_arg @@ spf "version is %d, expected 0" n
+      | None -> invalid_arg "expected 2-digit version");
+      if Bytes.get bs 2 <> '-' then invalid_arg "expected '-' before trace_id";
+      let trace_id =
+        try Trace_id.of_hex_substring (Bytes.unsafe_to_string bs) 3
+        with Invalid_argument msg -> invalid_arg (spf "in trace id: %s" msg)
+      in
+      if Bytes.get bs (3 + 32) <> '-' then
+        invalid_arg "expected '-' before parent_id";
+      let parent_id =
+        try Span_id.of_hex_substring (Bytes.unsafe_to_string bs) 36
+        with Invalid_argument msg -> invalid_arg (spf "in span id: %s" msg)
+      in
+      if Bytes.get bs 52 <> '-' then invalid_arg "expected '-' after parent_id";
+
+      (* ignore flags *)
+      Ok { trace_id; parent_id; is_remote = true }
+    with Invalid_argument msg -> Error msg
+
+  let of_w3c_trace_context_exn bs =
+    match of_w3c_trace_context bs with
+    | Ok t -> t
+    | Error msg -> invalid_arg @@ spf "invalid w3c trace context: %s" msg
+end
+
 (** {2 Attributes and conventions} *)
 
 module Conventions = struct
@@ -528,6 +672,10 @@ module Scope = struct
     mutable attrs: key_value list;
   }
 
+  (** Turn the scope into a span context *)
+  let[@inline] to_span_ctx (self : t) : Span_ctx.t =
+    Span_ctx.make ~trace_id:self.trace_id ~parent_id:self.span_id ()
+
   (** Add an event to the scope. It will be aggregated into the span.
 
       Note that this takes a function that produces an event, and will only
@@ -582,6 +730,8 @@ module Span_link : sig
     ?dropped_attributes_count:int ->
     unit ->
     t
+
+  val of_span_ctx : ?attrs:key_value list -> Span_ctx.t -> t
 end = struct
   open Proto.Trace
 
@@ -597,6 +747,10 @@ end = struct
       ~trace_id:(Trace_id.to_bytes trace_id)
       ~span_id:(Span_id.to_bytes span_id) ?trace_state ~attributes
       ?dropped_attributes_count ()
+
+  let[@inline] of_span_ctx ?attrs ctx : t =
+    make ~trace_id:(Span_ctx.trace_id ctx) ~span_id:(Span_ctx.parent_id ctx)
+      ?attrs ()
 end
 
 (** Spans.
@@ -1091,53 +1245,13 @@ module Trace_context = struct
         [{flags}] are currently ignored.
     *)
     let of_value str : (Trace_id.t * Span_id.t, string) result =
-      let ( let* ) = result_bind in
+      match Span_ctx.of_w3c_trace_context (Bytes.unsafe_of_string str) with
-      let blit ~offset ~len ~or_ =
+      | Ok sp -> Ok (Span_ctx.trace_id sp, Span_ctx.parent_id sp)
-        let buf = Bytes.create len in
+      | Error _ as e -> e
-        let* str =
-          match Bytes.blit_string str offset buf 0 len with
-          | () -> Ok (Bytes.unsafe_to_string buf)
-          | exception Invalid_argument _ -> Error or_
-        in
-        Ok (str, offset + len)
-      in
-      let consume expected ~offset ~or_ =
-        let len = String.length expected in
-        let* str, offset = blit ~offset ~len ~or_ in
-        if str = expected then
-          Ok offset
-        else
-          Error or_
-      in
-      let offset = 0 in
-      let* offset = consume "00" ~offset ~or_:"Expected version 00" in
-      let* offset = consume "-" ~offset ~or_:"Expected delimiter" in
-      let* trace_id, offset =
-        blit ~offset ~len:32 ~or_:"Expected 32-digit trace-id"
-      in
-      let* trace_id =
-        match Trace_id.of_hex trace_id with
-        | trace_id -> Ok trace_id
-        | exception Invalid_argument _ -> Error "Expected hex-encoded trace-id"
-      in
-      let* offset = consume "-" ~offset ~or_:"Expected delimiter" in
-      let* parent_id, offset =
-        blit ~offset ~len:16 ~or_:"Expected 16-digit parent-id"
-      in
-      let* parent_id =
-        match Span_id.of_hex parent_id with
-        | parent_id -> Ok parent_id
-        | exception Invalid_argument _ -> Error "Expected hex-encoded parent-id"
-      in
-      let* offset = consume "-" ~offset ~or_:"Expected delimiter" in
-      let* _flags, _offset =
-        blit ~offset ~len:2 ~or_:"Expected 2-digit flags"
-      in
-      Ok (trace_id, parent_id)
 
     let to_value ~(trace_id : Trace_id.t) ~(parent_id : Span_id.t) () : string =
-      Printf.sprintf "00-%s-%s-00" (Trace_id.to_hex trace_id)
+      let span_ctx = Span_ctx.make ~trace_id ~parent_id () in
-        (Span_id.to_hex parent_id)
+      Bytes.unsafe_to_string @@ Span_ctx.to_w3c_trace_context span_ctx
   end
 end
 

tests/core/test_trace_context.expected

@@ -1,25 +1,31 @@
 Trace_context.Traceparent.of_value "xx":
-  Error "Expected version 00"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00":
-  Error "Expected delimiter"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-xxxx":
-  Error "Expected 32-digit trace-id"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx":
-  Error "Expected hex-encoded trace-id"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef":
-  Error "Expected delimiter"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef-xxxx":
-  Error "Expected 16-digit parent-id"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef-xxxxxxxxxxxxxxxx":
-  Error "Expected hex-encoded parent-id"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef-0123456789abcdef":
-  Error "Expected delimiter"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef-0123456789abcdef-":
-  Error "Expected 2-digit flags"
+  Error "trace context must be 55 bytes"
 Trace_context.Traceparent.of_value "00-0123456789abcdef0123456789abcdef-0123456789abcdef-00":
   Ok trace_id:"0123456789abcdef0123456789abcdef" parent_id:"0123456789abcdef"
 Trace_context.Traceparent.of_value "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01":
   Ok trace_id:"4bf92f3577b34da6a3ce929d0e0e4736" parent_id:"00f067aa0ba902b7"
+Trace_context.Traceparent.of_value "03-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01":
+  Error "version is 3, expected 0"
+Trace_context.Traceparent.of_value "00-ohnonohex7b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01":
+  Error "in trace id: invalid hex char: 'o'"
+Trace_context.Traceparent.of_value "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aazzzzzzb7-01":
+  Error "in span id: invalid hex char: 'z'"
 
 Trace_context.Traceparent.to_value trace_id:"4bf92f3577b34da6a3ce929d0e0e4736" parent_id:"00f067aa0ba902b7":
   "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-00"

tests/core/test_trace_context.ml

@@ -36,6 +36,12 @@ let () = test_of_value "00-0123456789abcdef0123456789abcdef-0123456789abcdef-00"
 
 let () = test_of_value "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
 
+let () = test_of_value "03-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+
+let () = test_of_value "00-ohnonohex7b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"
+
+let () = test_of_value "00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aazzzzzzb7-01"
+
 let () = print_endline ""
 
 let test_to_value trace_id parent_id =