From 2cc6ef8b3f03429baf68a7f816601e6aaf1f2c8c Mon Sep 17 00:00:00 2001
From: Gabriel Buica <danutgabriel.buica@cloud.com>
Date: Tue, 23 Jul 2024 15:19:47 +0100
Subject: [PATCH] CP-50444: Intrument `request_of_bio`

Intruments the request reading loop.

Adds new functionality to the tracing library to update a span with a
new parent.

This way we can retroactively set the parent span with the correnct one.
---
 ocaml/libs/http-lib/http_svr.ml | 20 ++++++++++++++++++++
 ocaml/libs/tracing/tracing.ml   | 20 ++++++++++++++++++++
 ocaml/libs/tracing/tracing.mli  |  4 ++++
 3 files changed, 44 insertions(+)

diff --git a/ocaml/libs/http-lib/http_svr.ml b/ocaml/libs/http-lib/http_svr.ml
index 0928642f8f7..e590c334aa3 100644
--- a/ocaml/libs/http-lib/http_svr.ml
+++ b/ocaml/libs/http-lib/http_svr.ml
@@ -453,9 +453,28 @@ let request_of_bio_exn ~proxy_seen ~read_timeout ~total_timeout ~max_length bio
     	already sent back a suitable error code and response to the client. *)
 let request_of_bio ?proxy_seen ~read_timeout ~total_timeout ~max_length ic =
   try
+    let tracer = Tracing.Tracer.get_tracer ~name:"http_tracer" in
+    let loop_span =
+      match Tracing.Tracer.start ~tracer ~name:__FUNCTION__ ~parent:None () with
+      | Ok span ->
+          span
+      | Error _ ->
+          None
+    in
     let r, proxy =
       request_of_bio_exn ~proxy_seen ~read_timeout ~total_timeout ~max_length ic
     in
+    let parent_span = traceparent_of_request r in
+    let loop_span =
+      Option.fold ~none:None
+        ~some:(fun span ->
+          Tracing.Tracer.update_span_with_parent span parent_span
+        )
+        loop_span
+    in
+    let _ : (Tracing.Span.t option, exn) result =
+      Tracing.Tracer.finish loop_span
+    in
     (Some r, proxy)
   with e ->
     D.warn "%s (%s)" (Printexc.to_string e) __LOC__ ;
@@ -583,6 +602,7 @@ let handle_connection ~header_read_timeout ~header_total_timeout
       request_of_bio ?proxy_seen ~read_timeout ~total_timeout
         ~max_length:max_header_length ic
     in
+
     (* 2. now we attempt to process the request *)
     let finished =
       Option.fold ~none:true
diff --git a/ocaml/libs/tracing/tracing.ml b/ocaml/libs/tracing/tracing.ml
index 28a5b2f687b..4d3c55ea8bf 100644
--- a/ocaml/libs/tracing/tracing.ml
+++ b/ocaml/libs/tracing/tracing.ml
@@ -151,6 +151,8 @@ end
 module SpanContext = struct
   type t = {trace_id: string; span_id: string} [@@deriving rpcty]
 
+  let context trace_id span_id = {trace_id; span_id}
+
   let to_traceparent t = Printf.sprintf "00-%s-%s-01" t.trace_id t.span_id
 
   let of_traceparent traceparent =
@@ -624,6 +626,24 @@ module Tracer = struct
       let span = Span.start ~attributes ~name ~parent ~span_kind () in
       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
+
   let finish ?error span =
     Ok
       (Option.map
diff --git a/ocaml/libs/tracing/tracing.mli b/ocaml/libs/tracing/tracing.mli
index cd6e4204602..4e1f14adf96 100644
--- a/ocaml/libs/tracing/tracing.mli
+++ b/ocaml/libs/tracing/tracing.mli
@@ -57,6 +57,8 @@ end
 module SpanContext : sig
   type t
 
+  val context : string -> string -> t
+
   val to_traceparent : t -> string
 
   val of_traceparent : string -> t option
@@ -125,6 +127,8 @@ module Tracer : sig
     -> unit
     -> (Span.t option, exn) result
 
+  val update_span_with_parent : Span.t -> Span.t option -> Span.t option
+
   val finish :
     ?error:exn * string -> Span.t option -> (Span.t option, exn) result