CP-52744: Thread TraceContext as JSON inside debug_info

dune-project

@@ -103,11 +103,13 @@
    dune
    (alcotest :with-test)
    (fmt :with-test)
+   ppx_deriving_yojson
    re
    uri
    (uuid :with-test)
    (xapi-log (= :version))
    (xapi-stdext-threads (= :version))
+   yojson
   )
   (synopsis "Allows to instrument code to generate tracing information")
   (description "This library provides modules to allow gathering runtime traces.")

ocaml/libs/tracing/dune

@@ -1,7 +1,9 @@
 (library
  (name tracing)
  (modules tracing)
- (libraries re uri xapi-log xapi-stdext-threads threads.posix)
+ (libraries re uri yojson xapi-log xapi-stdext-threads threads.posix)
+ (preprocess
+  (pps ppx_deriving_yojson))
  (public_name xapi-tracing))
 
 (library

ocaml/libs/tracing/tracing.ml

@@ -211,11 +211,12 @@ end
 
 (* The context of a trace that can be propagated across service boundaries. *)
 module TraceContext = struct
-  type traceparent = string
+  type traceparent = string [@@deriving yojson]
 
-  type baggage = (string * string) list
+  type baggage = (string * string) list [@@deriving yojson]
 
   type t = {traceparent: traceparent option; baggage: baggage option}
+  [@@deriving yojson]
 
   let empty = {traceparent= None; baggage= None}
 
@@ -226,6 +227,10 @@ module TraceContext = struct
   let traceparent_of ctx = ctx.traceparent
 
   let baggage_of ctx = ctx.baggage
+
+  let to_json_string t = Yojson.Safe.to_string (to_yojson t)
+
+  let of_json_string s = of_yojson (Yojson.Safe.from_string s)
 end
 
 module SpanContext = struct
@@ -297,6 +302,8 @@ module Span = struct
 
   let get_context t = t.context
 
+  let get_trace_context t = t.context |> SpanContext.context_of_span_context
+
   let start ?(attributes = Attributes.empty)
       ?(trace_context : TraceContext.t option) ~name ~parent ~span_kind () =
     let trace_id, extra_context =
@@ -312,9 +319,9 @@ module Span = struct
     in
     let context =
       (* If trace_context is provided to the call, override any inherited trace context. *)
-      Option.fold ~none:context
-        ~some:(Fun.flip SpanContext.with_trace_context context)
-        trace_context
+      trace_context
+      |> Option.fold ~none:context
+           ~some:(Fun.flip SpanContext.with_trace_context context)
     in
     (* Using gettimeofday over Mtime as it is better for sharing timestamps between the systems *)
     let begin_time = Unix.gettimeofday () in
@@ -411,6 +418,12 @@ module Span = struct
         {span with status= {status_code; _description}}
     | _ ->
         span
+
+  let with_trace_context span trace_context =
+    let span_context =
+      span |> get_context |> SpanContext.with_trace_context trace_context
+    in
+    {span with context= span_context}
 end
 
 module TraceMap = Map.Make (Trace_id)

ocaml/libs/tracing/tracing.mli

@@ -94,6 +94,10 @@ module TraceContext : sig
   val traceparent_of : t -> traceparent option
 
   val baggage_of : t -> baggage option
+
+  val to_json_string : t -> string
+
+  val of_json_string : string -> (t, string) result
 end
 
 module SpanContext : sig
@@ -119,6 +123,8 @@ module Span : sig
 
   val get_context : t -> SpanContext.t
 
+  val get_trace_context : t -> TraceContext.t
+
   val add_link : t -> SpanContext.t -> (string * string) list -> t
 
   val add_event : t -> string -> (string * string) list -> t
@@ -140,6 +146,8 @@ module Span : sig
   val get_end_time : t -> float option
 
   val get_attributes : t -> (string * string) list
+
+  val with_trace_context : t -> TraceContext.t -> t
 end
 
 module TraceMap : module type of Map.Make (Trace_id)

ocaml/xapi-idl/lib/debug_info.ml

@@ -22,10 +22,34 @@ let make ~log ~tracing = {log; tracing}
 let of_string s =
   let open Tracing in
   match String.split_on_char separator s with
-  | [log; traceparent] ->
-      let spancontext = SpanContext.of_traceparent traceparent in
+  | [log; trace_context] ->
+      (* Process the tracing data:
+         1. We expect a JSON representing the trace_context.
+         2. If the JSON is valid but not representing a trace_context,
+         we ignore the tracing data.
+         3. If we get an exception from parsing the JSON string,
+         it means a traceparent string was received.*)
+      let trace_context =
+        try
+          let trace_context =
+            Tracing.TraceContext.of_json_string trace_context
+          in
+          match trace_context with
+          | Ok trace_context ->
+              Some trace_context
+          | Error _ ->
+              None
+        with _ ->
+          Some
+            (TraceContext.empty
+            |> TraceContext.with_traceparent (Some trace_context)
+            )
+      in
+      let spancontext =
+        Option.(join (map Tracing.SpanContext.of_trace_context trace_context))
+      in
       let tracing =
-        Option.map (fun tp -> Tracer.span_of_span_context tp log) spancontext
+        Option.map (Fun.flip Tracer.span_of_span_context log) spancontext
       in
       {log; tracing}
   | _ ->
@@ -37,11 +61,18 @@ let filter_separator = Astring.String.filter (( <> ) separator)
 let to_string t =
   Option.fold ~none:t.log
     ~some:(fun span ->
-      let traceparent =
-        Tracing.Span.get_context span |> Tracing.SpanContext.to_traceparent
+      let trace_context =
+        let traceparent =
+          span |> Tracing.Span.get_context |> Tracing.SpanContext.to_traceparent
+        in
+        span
+        |> Tracing.Span.get_context
+        |> Tracing.SpanContext.context_of_span_context
+        |> Tracing.TraceContext.with_traceparent (Some traceparent)
+        |> Tracing.TraceContext.to_json_string
       in
       Printf.sprintf "%s%c%s" (filter_separator t.log) separator
-        (filter_separator traceparent)
+        (filter_separator trace_context)
     )
     t.tracing
 
@@ -68,7 +99,21 @@ let with_dbg ?(with_thread = false) ~module_name ~name ~dbg f =
 
 let traceparent_of_dbg dbg =
   match String.split_on_char separator dbg with
-  | [_; traceparent] ->
-      Some traceparent
+  | [_; trace_context] -> (
+    (* Process the tracing data:
+         1. We expect a JSON representing the trace_context.
+         2. If the JSON is valid but not representing a trace_context,
+         we ignore the tracing data.
+         3. If we get an exception from parsing the JSON string,
+         it means a traceparent string was received.*)
+    try
+      let trace_context = Tracing.TraceContext.of_json_string trace_context in
+      match trace_context with
+      | Ok trace_context ->
+          Tracing.TraceContext.traceparent_of trace_context
+      | Error _ ->
+          None
+    with _ -> Some trace_context
+  )
   | _ ->
       None

ocaml/xapi/context.ml

@@ -337,7 +337,20 @@ let start_tracing_helper ?(span_attributes = []) parent_fn task_name =
       ~parent ()
   with
   | Ok x ->
-      x
+      Option.map
+        (fun span ->
+          let traceparent =
+            span |> Span.get_context |> SpanContext.to_traceparent
+          in
+          let trace_context =
+            span
+            |> Span.get_context
+            |> SpanContext.context_of_span_context
+            |> TraceContext.with_traceparent (Some traceparent)
+          in
+          Span.with_trace_context span trace_context
+        )
+        x
   | Error e ->
       R.warn "Failed to start tracing: %s" (Printexc.to_string e) ;
       None

xapi-tracing.opam

@@ -13,11 +13,13 @@ depends: [
   "dune" {>= "3.15"}
   "alcotest" {with-test}
   "fmt" {with-test}
+  "ppx_deriving_yojson"
   "re"
   "uri"
   "uuid" {with-test}
   "xapi-log" {= version}
   "xapi-stdext-threads" {= version}
+  "yojson"
   "odoc" {with-doc}
 ]
 build: [