CP-50444: Add specialized function for tracing http requests to `Http…

….Request`

Like the title says, this commit adds `Http.Request.with_tracing` to
`Http.Request`.

This should enable updating the traceparent field of a request while we
process it.

Signed-off-by: Gabriel Buica <danutgabriel.buica@cloud.com>

ocaml/libs/http-lib/http.ml

@@ -694,6 +694,29 @@ module Request = struct
     let headers, body = to_headers_and_body x in
     let frame_header = if x.frame then make_frame_header headers else "" in
     frame_header ^ headers ^ body
+
+  let traceparent_of req =
+    let open Tracing in
+    let ( let* ) = Option.bind in
+    let* traceparent = req.traceparent in
+    let* span_context = SpanContext.of_traceparent traceparent in
+    let span = Tracer.span_of_span_context span_context req.uri in
+    Some span
+
+  let with_tracing ?attributes ~name req f =
+    let open Tracing in
+    let parent = traceparent_of req in
+    with_child_trace ?attributes parent ~name (fun (span : Span.t option) ->
+        match span with
+        | Some span ->
+            let traceparent =
+              Some (span |> Span.get_context |> SpanContext.to_traceparent)
+            in
+            let req = {req with traceparent} in
+            f req
+        | None ->
+            f req
+    )
 end
 
 module Response = struct

ocaml/libs/http-lib/http.mli

@@ -128,6 +128,11 @@ module Request : sig
 
   val to_wire_string : t -> string
   (** [to_wire_string t] returns a string which could be sent to a server *)
+
+  val traceparent_of : t -> Tracing.Span.t option
+
+  val with_tracing :
+    ?attributes:(string * string) list -> name:string -> t -> (t -> 'a) -> 'a
 end
 
 (** Parsed form of the HTTP response *)

ocaml/libs/http-lib/http_svr.ml

@@ -101,18 +101,9 @@ let response_of_request req hdrs =
     ~headers:(connection :: cache :: hdrs)
     "200" "OK"
 
-let traceparent_of_request req =
-  let open Tracing in
-  let ( let* ) = Option.bind in
-  let* traceparent = req.Http.Request.traceparent in
-  let* span_context = SpanContext.of_traceparent traceparent in
-  let span = Tracer.span_of_span_context span_context req.uri in
-  Some span
-
 let response_fct req ?(hdrs = []) s (response_length : int64)
     (write_response_to_fd_fn : Unix.file_descr -> unit) =
-  let parent = traceparent_of_request req in
-  let@ _ = Tracing.with_child_trace parent ~name:__FUNCTION__ in
+  let@ req = Http.Request.with_tracing ~name:__FUNCTION__ req in
   let res =
     {
       (response_of_request req hdrs) with
@@ -464,7 +455,7 @@ let request_of_bio ?proxy_seen ~read_timeout ~total_timeout ~max_length ic =
     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 parent_span = Http.Request.traceparent_of r in
     let loop_span =
       Option.fold ~none:None
         ~some:(fun span ->
@@ -517,8 +508,8 @@ let request_of_bio ?proxy_seen ~read_timeout ~total_timeout ~max_length ic =
     (None, None)
 
 let handle_one (x : 'a Server.t) ss context req =
-  let parent = traceparent_of_request req in
-  let@ span = Tracing.with_child_trace parent ~name:__FUNCTION__ in
+  let@ req = Http.Request.with_tracing ~name:__FUNCTION__ req in
+  let span = Http.Request.traceparent_of req in
   let ic = Buf_io.of_fd ss in
   let finished = ref false in
   try
@@ -533,13 +524,6 @@ let handle_one (x : 'a Server.t) ss context req =
         (Radix_tree.longest_prefix req.Request.uri method_map)
     in
     let@ _ = Tracing.with_child_trace span ~name:"handler" in
-    let traceparent =
-      let open Tracing in
-      Option.map
-        (fun span -> Span.get_context span |> SpanContext.to_traceparent)
-        span
-    in
-    let req = {req with traceparent} in
     ( match te.TE.handler with
     | BufIO handlerfn ->
         handlerfn req ic context

ocaml/libs/http-lib/http_svr.mli

@@ -143,5 +143,3 @@ val https_client_of_req : Http.Request.t -> Ipaddr.t option
 val client_of_req_and_fd : Http.Request.t -> Unix.file_descr -> client option
 
 val string_of_client : client -> string
-
-val traceparent_of_request : Http.Request.t -> Tracing.Span.t option

ocaml/xapi/api_server.ml

@@ -5,8 +5,7 @@ let ( let@ ) f x = f x
 
 (* This bit is called directly by the fake_rpc callback *)
 let callback1 ?(json_rpc_version = Jsonrpc.V1) is_json req fd call =
-  let parent = Http_svr.traceparent_of_request req in
-  let@ span = Tracing.with_child_trace parent ~name:__FUNCTION__ in
+  let@ req = Http.Request.with_tracing ~name:__FUNCTION__ req in
   (* We now have the body string, the xml and the call name, and can also tell *)
   (* if we're a master or slave and whether the call came in on the unix domain socket or the tcp socket *)
   (* If we're a slave, and the call is from the unix domain socket or from the HIMN, and the call *isn't* *)
@@ -25,7 +24,7 @@ let callback1 ?(json_rpc_version = Jsonrpc.V1) is_json req fd call =
     forward req call is_json
   else
     let response =
-      let@ _ = Tracing.with_child_trace span ~name:"Server.dispatch_call" in
+      let@ req = Http.Request.with_tracing ~name:"Server.dispatch_call" req in
       Server.dispatch_call req fd call
     in
     let translated =
@@ -92,8 +91,8 @@ let create_thumbprint_header req response =
 
 (** HTML callback that dispatches an RPC and returns the response. *)
 let callback is_json req bio _ =
-  let parent = Http_svr.traceparent_of_request req in
-  let@ span = Tracing.with_child_trace parent ~name:__FUNCTION__ in
+  let@ req = Http.Request.with_tracing ~name:__FUNCTION__ req in
+  let span = Http.Request.traceparent_of req in
   let fd = Buf_io.fd_of bio in
   (* fd only used for writing *)
   let body =
@@ -147,8 +146,7 @@ let callback is_json req bio _ =
 
 (** HTML callback that dispatches an RPC and returns the response. *)
 let jsoncallback req bio _ =
-  let parent = Http_svr.traceparent_of_request req in
-  let@ _ = Tracing.with_child_trace parent ~name:__FUNCTION__ in
+  let@ req = Http.Request.with_tracing ~name:__FUNCTION__ req in
   let fd = Buf_io.fd_of bio in
   (* fd only used for writing *)
   let body =

ocaml/xapi/server_helpers.ml

@@ -18,6 +18,8 @@ open D
 
 exception Dispatcher_FieldNotFound of string
 
+let ( let@ ) f x = f x
+
 let my_assoc fld assoc_list =
   try List.assoc fld assoc_list
   with Not_found -> raise (Dispatcher_FieldNotFound fld)
@@ -120,6 +122,7 @@ let dispatch_exn_wrapper f =
 let do_dispatch ?session_id ?forward_op ?self:_ supports_async called_fn_name
     op_fn marshaller fd http_req label sync_ty generate_task_for =
   (* if the call has been forwarded to us, then they are responsible for completing the task, so we don't need to complete it *)
+  let@ http_req = Http.Request.with_tracing ~name:__FUNCTION__ http_req in
   let called_async = sync_ty <> `Sync in
   if called_async && not supports_async then
     API.response_of_fault