From 99afa2ea3f4154ef1dfd34d138f93feaae2cefe9 Mon Sep 17 00:00:00 2001
From: Gabriel Buica <danutgabriel.buica@cloud.com>
Date: Thu, 23 Jan 2025 08:53:03 +0000
Subject: [PATCH] CP-52744: Update `trace_context` when starting a context span

Refresh the trace_context with the correct traceparent when creating a
span with `start_tracing_helper` in `context.ml`.

This ensures the tracing of the context has the correct parent.

Signed-off-by: Gabriel Buica <danutgabriel.buica@cloud.com>
---
 ocaml/libs/tracing/tracing.ml  |  6 ++++++
 ocaml/libs/tracing/tracing.mli |  2 ++
 ocaml/xapi/context.ml          | 15 ++++++++++++++-
 3 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/ocaml/libs/tracing/tracing.ml b/ocaml/libs/tracing/tracing.ml
index 98a792598bc..30ffdb4fea5 100644
--- a/ocaml/libs/tracing/tracing.ml
+++ b/ocaml/libs/tracing/tracing.ml
@@ -418,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)
diff --git a/ocaml/libs/tracing/tracing.mli b/ocaml/libs/tracing/tracing.mli
index 7239bf83675..b3f158c6833 100644
--- a/ocaml/libs/tracing/tracing.mli
+++ b/ocaml/libs/tracing/tracing.mli
@@ -146,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)
diff --git a/ocaml/xapi/context.ml b/ocaml/xapi/context.ml
index 5f357e110af..f9ba9877e68 100644
--- a/ocaml/xapi/context.ml
+++ b/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