CP-51772: Propagate trace context through spans

Adds TraceContext to the SpanContext data structure and attempts to
ensure its inheritance through each part of the code base.

The API exposed by the tracing library is a bit problematic, it ought to
be simplified and adapted to various use cases.

Signed-off-by: Colin James <[email protected]>

ocaml/libs/http-lib/http_svr.ml

@@ -442,17 +442,24 @@ let read_request_exn ~proxy_seen ~read_timeout ~total_timeout ~max_length fd =
     	already sent back a suitable error code and response to the client. *)
 let read_request ?proxy_seen ~read_timeout ~total_timeout ~max_length fd =
   try
+    (* TODO: Restore functionality of tracing this function. We rely on the request
+       to contain information we want spans to inherit. However, it is the reading of the
+       request that we intend to trace. *)
+    let r, proxy =
+      read_request_exn ~proxy_seen ~read_timeout ~total_timeout ~max_length fd
+    in
+    let trace_context = Tracing_propagator.Propagator.Http.extract_from r in
     let tracer = Tracing.Tracer.get_tracer ~name:"http_tracer" in
     let loop_span =
-      match Tracing.Tracer.start ~tracer ~name:__FUNCTION__ ~parent:None () with
+      match
+        Tracing.Tracer.start ~tracer ~trace_context ~name:__FUNCTION__
+          ~parent:None ()
+      with
       | Ok span ->
           span
       | Error _ ->
           None
     in
-    let r, proxy =
-      read_request_exn ~proxy_seen ~read_timeout ~total_timeout ~max_length fd
-    in
     let parent_span = Helper.traceparent_of r in
     let loop_span =
       Option.fold ~none:None

ocaml/libs/tracing/tracing.ml

@@ -95,7 +95,7 @@ let validate_attribute (key, value) =
   && W3CBaggage.Key.is_valid_key key
 
 module SpanKind = struct
-  type t = Server | Consumer | Client | Producer | Internal [@@deriving rpcty]
+  type t = Server | Consumer | Client | Producer | Internal
 
   let to_string = function
     | Server ->
@@ -127,7 +127,7 @@ let endpoint_to_string = function
 let ok_none = Ok None
 
 module Status = struct
-  type status_code = Unset | Ok | Error [@@deriving rpcty]
+  type status_code = Unset | Ok | Error
 
   type t = {status_code: status_code; _description: string option}
 end
@@ -229,9 +229,14 @@ module TraceContext = struct
 end
 
 module SpanContext = struct
-  type t = {trace_id: Trace_id.t; span_id: Span_id.t} [@@deriving rpcty]
+  type t = {
+      trace_id: Trace_id.t
+    ; span_id: Span_id.t
+    ; trace_context: TraceContext.t
+  }
 
-  let context trace_id span_id = {trace_id; span_id}
+  let context trace_id span_id =
+    {trace_id; span_id; trace_context= TraceContext.empty}
 
   let to_traceparent t =
     let tid = Trace_id.to_string t.trace_id in
@@ -246,13 +251,23 @@ module SpanContext = struct
           {
             trace_id= Trace_id.of_string trace_id
           ; span_id= Span_id.of_string span_id
+          ; trace_context= TraceContext.empty
           }
     | _ ->
         None
 
   let trace_id_of_span_context t = t.trace_id
 
   let span_id_of_span_context t = t.span_id
+
+  let context_of_span_context t = t.trace_context
+
+  let with_trace_context trace_context t = {t with trace_context}
+
+  let of_trace_context trace_context =
+    let traceparent = TraceContext.traceparent_of trace_context in
+    let span_context = Option.(join (map of_traceparent traceparent)) in
+    Option.map (with_trace_context trace_context) span_context
 end
 
 module SpanLink = struct
@@ -282,16 +297,25 @@ module Span = struct
 
   let get_context t = t.context
 
-  let start ?(attributes = Attributes.empty) ~name ~parent ~span_kind () =
-    let trace_id =
+  let start ?(attributes = Attributes.empty)
+      ?(trace_context : TraceContext.t option) ~name ~parent ~span_kind () =
+    let trace_id, extra_context =
       match parent with
       | None ->
-          Trace_id.make ()
+          (Trace_id.make (), TraceContext.empty)
       | Some span_parent ->
-          span_parent.context.trace_id
+          (span_parent.context.trace_id, span_parent.context.trace_context)
     in
     let span_id = Span_id.make () in
-    let context : SpanContext.t = {trace_id; span_id} in
+    let context : SpanContext.t =
+      {trace_id; span_id; trace_context= extra_context}
+    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
+    in
     (* Using gettimeofday over Mtime as it is better for sharing timestamps between the systems *)
     let begin_time = Unix.gettimeofday () in
     let end_time = None in
@@ -669,15 +693,18 @@ module Tracer = struct
     ; attributes= Attributes.empty
     }
 
-  let start ~tracer:t ?(attributes = []) ?(span_kind = SpanKind.Internal) ~name
-      ~parent () : (Span.t option, exn) result =
+  let start ~tracer:t ?(attributes = []) ?trace_context
+      ?(span_kind = SpanKind.Internal) ~name ~parent () :
+      (Span.t option, exn) result =
     let open TracerProvider in
     (* Do not start span if the TracerProvider is disabled*)
     if not t.enabled then
       ok_none
     else
       let attributes = Attributes.merge_into t.attributes attributes in
-      let span = Span.start ~attributes ~name ~parent ~span_kind () in
+      let span =
+        Span.start ~attributes ?trace_context ~name ~parent ~span_kind ()
+      in
       Spans.add_to_spans ~span ; Ok (Some span)
 
   let update_span_with_parent span (parent : Span.t option) =
@@ -691,9 +718,11 @@ module Tracer = struct
           |> Option.map (fun existing_span ->
                  let old_context = Span.get_context existing_span in
                  let new_context : SpanContext.t =
+                   let trace_context = span.Span.context.trace_context in
                    SpanContext.context
                      (SpanContext.trace_id_of_span_context parent.context)
                      old_context.span_id
+                   |> SpanContext.with_trace_context trace_context
                  in
                  let updated_span = {existing_span with parent= Some parent} in
                  let updated_span = {updated_span with context= new_context} in
@@ -730,10 +759,10 @@ end
 let enable_span_garbage_collector ?(timeout = 86400.) () =
   Spans.GC.initialise_thread ~timeout
 
-let with_tracing ?(attributes = []) ?(parent = None) ~name f =
+let with_tracing ?(attributes = []) ?(parent = None) ?trace_context ~name f =
   let tracer = Tracer.get_tracer ~name in
   if tracer.enabled then (
-    match Tracer.start ~tracer ~attributes ~name ~parent () with
+    match Tracer.start ~tracer ?trace_context ~attributes ~name ~parent () with
     | Ok span -> (
       try
         let result = f span in
@@ -751,12 +780,12 @@ let with_tracing ?(attributes = []) ?(parent = None) ~name f =
   ) else
     f None
 
-let with_child_trace ?attributes parent ~name f =
+let with_child_trace ?attributes ?trace_context parent ~name f =
   match parent with
   | None ->
       f None
   | Some _ as parent ->
-      with_tracing ?attributes ~parent ~name f
+      with_tracing ?attributes ?trace_context ~parent ~name f
 
 module EnvHelpers = struct
   let traceparent_key = "TRACEPARENT"
@@ -824,6 +853,9 @@ module Propagator = struct
       let trace_context = P.extract_from carrier in
       let* parent = TraceContext.traceparent_of trace_context in
       let* span_context = SpanContext.of_traceparent parent in
+      let span_context =
+        SpanContext.with_trace_context trace_context span_context
+      in
       let name = P.name_span carrier in
       Some (Tracer.span_of_span_context span_context name)
 
@@ -845,6 +877,7 @@ module Propagator = struct
         | _ ->
             f carrier
       in
-      with_child_trace ?attributes parent ~name continue_with_child
+      with_child_trace ?attributes ~trace_context parent ~name
+        continue_with_child
   end
 end

ocaml/libs/tracing/tracing.mli

@@ -103,9 +103,13 @@ module SpanContext : sig
 
   val of_traceparent : string -> t option
 
+  val of_trace_context : TraceContext.t -> t option
+
   val trace_id_of_span_context : t -> Trace_id.t
 
   val span_id_of_span_context : t -> Span_id.t
+
+  val context_of_span_context : t -> TraceContext.t
 end
 
 module Span : sig
@@ -164,6 +168,7 @@ module Tracer : sig
   val start :
        tracer:t
     -> ?attributes:(string * string) list
+    -> ?trace_context:TraceContext.t
     -> ?span_kind:SpanKind.t
     -> name:string
     -> parent:Span.t option
@@ -250,12 +255,14 @@ val enable_span_garbage_collector : ?timeout:float -> unit -> unit
 val with_tracing :
      ?attributes:(string * string) list
   -> ?parent:Span.t option
+  -> ?trace_context:TraceContext.t
   -> name:string
   -> (Span.t option -> 'a)
   -> 'a
 
 val with_child_trace :
      ?attributes:(string * string) list
+  -> ?trace_context:TraceContext.t
   -> Span.t option
   -> name:string
   -> (Span.t option -> 'a)

ocaml/libs/tracing/tracing_export.ml

@@ -82,6 +82,16 @@ module Content = struct
                  {timestamp; value}
              )
         in
+        let tags =
+          let span_context = Span.get_context s in
+          let trace_context =
+            SpanContext.context_of_span_context span_context
+          in
+          let baggage =
+            TraceContext.baggage_of trace_context |> Option.value ~default:[]
+          in
+          Span.get_attributes s @ baggage
+        in
         {
           id=
             s
@@ -117,7 +127,7 @@ module Content = struct
             |> Option.map SpanKind.to_string
         ; localEndpoint= {serviceName}
         ; annotations
-        ; tags= Span.get_attributes s
+        ; tags
         }
 
       let content_of (spans : Span.t list) =
@@ -270,7 +280,10 @@ module Destination = struct
             ; ("xs.tracing.finished_spans_table.count", string_of_int count)
             ]
           in
-          let@ _ = with_tracing ~parent ~attributes ~name in
+          let@ _ =
+            with_tracing ~trace_context:TraceContext.empty ~parent ~attributes
+              ~name
+          in
           all_spans
           |> Content.Json.ZipkinV2.content_of
           |> export
@@ -283,7 +296,8 @@ module Destination = struct
     let ((_span_list, span_count) as span_info) = Spans.since () in
     let attributes = [("export.traces.count", string_of_int span_count)] in
     let@ parent =
-      with_tracing ~parent:None ~attributes ~name:"Tracing.flush_spans"
+      with_tracing ~trace_context:TraceContext.empty ~parent:None ~attributes
+        ~name:"Tracing.flush_spans"
     in
     TracerProvider.get_tracer_providers ()
     |> List.filter TracerProvider.get_enabled

ocaml/xapi-cli-server/xapi_cli.ml

@@ -132,11 +132,17 @@ module TraceHelper = struct
     let open Tracing in
     let span_context = Option.map Span.get_context span in
     let traceparent = Option.map SpanContext.to_traceparent span_context in
-    let trace_context = TraceContext.(with_traceparent traceparent empty) in
+    let trace_context =
+      Option.map SpanContext.context_of_span_context span_context
+    in
+    let trace_context =
+      Option.value ~default:TraceContext.empty trace_context
+      |> TraceContext.with_traceparent traceparent
+    in
     Tracing_propagator.Propagator.Http.inject_into trace_context
 end
 
-let do_rpcs _req s username password minimal cmd session args tracing =
+let do_rpcs req s username password minimal cmd session args =
   let cmdname = get_cmdname cmd in
   let cspec =
     try Hashtbl.find cmdtable cmdname
@@ -151,8 +157,21 @@ let do_rpcs _req s username password minimal cmd session args tracing =
   let _ = check_required_keys cmd cspec.reqd in
   try
     let generic_rpc = get_rpc () in
+    let trace_context = Tracing_propagator.Propagator.Http.extract_from req in
+    let parent =
+      (* This is a "faux" span in the sense that it's not exported by the program. It exists
+         so that the derived child span can refer to its span-id as its parent during exportation
+         (along with inheriting the trace-id). *)
+      let open Tracing in
+      let ( let* ) = Option.bind in
+      let* traceparent = TraceContext.traceparent_of trace_context in
+      let* span_context = SpanContext.of_traceparent traceparent in
+      let span = Tracer.span_of_span_context span_context (get_cmdname cmd) in
+      Some span
+    in
     (* NB the request we've received is for the /cli. We need an XMLRPC request for the API *)
-    Tracing.with_tracing ~parent:tracing ~name:("xe " ^ cmdname) @@ fun span ->
+    Tracing.with_tracing ~trace_context ~parent ~name:("xe " ^ cmdname)
+    @@ fun span ->
     let req = Xmlrpc_client.xmlrpc ~version:"1.1" "/" in
     let req = TraceHelper.inject_span_into_req span req in
     let rpc = generic_rpc req s in
@@ -204,15 +223,6 @@ let uninteresting_cmd_postfixes = ["help"; "-get"; "-list"]
 
 let exec_command req cmd s session args =
   let params = get_params cmd in
-  let tracing =
-    let ( let* ) = Option.bind in
-    let open Tracing in
-    let context = Tracing_propagator.Propagator.Http.extract_from req in
-    let* traceparent = TraceContext.traceparent_of context in
-    let* span_context = SpanContext.of_traceparent traceparent in
-    let span = Tracer.span_of_span_context span_context (get_cmdname cmd) in
-    Some span
-  in
   let minimal =
     List.assoc_opt "minimal" params
     |> Option.fold ~none:false ~some:bool_of_string
@@ -271,7 +281,7 @@ let exec_command req cmd s session args =
             params
          )
       ) ;
-    do_rpcs req s u p minimal cmd session args tracing
+    do_rpcs req s u p minimal cmd session args
 
 let get_line str i =
   try

ocaml/xapi/context.ml

@@ -223,8 +223,7 @@ let parent_of_origin (origin : origin) span_name =
   | Http (req, _) ->
       let context = Tracing_propagator.Propagator.Http.extract_from req in
       let open Tracing in
-      let* traceparent = TraceContext.traceparent_of context in
-      let* span_context = SpanContext.of_traceparent traceparent in
+      let* span_context = SpanContext.of_trace_context context in
       let span = Tracer.span_of_span_context span_context span_name in
       Some span
   | _ ->

ocaml/xapi/helpers.ml

@@ -392,7 +392,13 @@ module TraceHelper = struct
     let open Tracing in
     let span_context = Option.map Span.get_context span in
     let traceparent = Option.map SpanContext.to_traceparent span_context in
-    let trace_context = TraceContext.(with_traceparent traceparent empty) in
+    let trace_context =
+      Option.map SpanContext.context_of_span_context span_context
+    in
+    let trace_context =
+      Option.value ~default:TraceContext.empty trace_context
+      |> TraceContext.with_traceparent traceparent
+    in
     Tracing_propagator.Propagator.Http.inject_into trace_context
 end