fixup! CP-50444: Intrument request_of_bio

xapi-project · Aug 8, 2024 · 28207c0 · 28207c0
1 parent 0e5efd8
commit 28207c0
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 16 deletions.
diff --git a/ocaml/libs/tracing/tracing.ml b/ocaml/libs/tracing/tracing.ml
@@ -627,22 +627,28 @@ module Tracer = struct
       Spans.add_to_spans ~span ; Ok (Some span)
 
   let update_span_with_parent span (parent : Span.t option) =
-    match parent with
-    | None ->
-        Some span
-    | Some parent ->
-        let _ : Span.t option = Spans.remove_from_spans span in
-        let old_context = Span.get_context span in
-        let new_context : SpanContext.t =
-          SpanContext.context
-            (SpanContext.trace_id_of_span_context parent.context)
-            old_context.span_id
-        in
-        let updated_span = {span with parent= Some parent} in
-        let updated_span = {updated_span with context= new_context} in
-
-        let () = Spans.add_to_spans ~span:updated_span in
-        Some updated_span
+    if Atomic.get observe then
+      match parent with
+      | None ->
+          Some span
+      | Some parent ->
+          span
+          |> Spans.remove_from_spans
+          |> Option.map (fun existing_span ->
+                 let old_context = Span.get_context existing_span in
+                 let new_context : SpanContext.t =
+                   SpanContext.context
+                     (SpanContext.trace_id_of_span_context parent.context)
+                     old_context.span_id
+                 in
+                 let updated_span = {existing_span with parent= Some parent} in
+                 let updated_span = {updated_span with context= new_context} in
+
+                 let () = Spans.add_to_spans ~span:updated_span in
+                 updated_span
+             )
+    else
+      Some span
 
   let finish ?error span =
     Ok

diff --git a/ocaml/libs/tracing/tracing.mli b/ocaml/libs/tracing/tracing.mli
@@ -128,6 +128,14 @@ module Tracer : sig
     -> (Span.t option, exn) result
 
   val update_span_with_parent : Span.t -> Span.t option -> Span.t option
+  (**[update_span_with_parent s p] returns [Some span] where [span] is an 
+  updated verison of the span [s].
+  If [p] is [Some parent], [span] is a child of [parent], otherwise it is the 
+  original [s].
+  
+  If the span [s] is finished or is no longer considered an on-going span, 
+  returns [None].
+  *)
 
   val finish :
     ?error:exn * string -> Span.t option -> (Span.t option, exn) result