From 3cbb4e5ea63612c9a9bc379bb65f8b22db8a1ada Mon Sep 17 00:00:00 2001
From: Corentin Leruth <corentin.leruth@gmail.com>
Date: Fri, 11 Oct 2024 10:39:25 +0200
Subject: [PATCH] web: fix wrong otel parent scope

---
 possibly_otel.mli     |  2 +-
 possibly_otel.real.ml |  6 +++++-
 possibly_otel.stub.ml |  2 +-
 web.ml                | 31 ++++++++++++++++---------------
 4 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/possibly_otel.mli b/possibly_otel.mli
index 4bca6b2..d880365 100644
--- a/possibly_otel.mli
+++ b/possibly_otel.mli
@@ -2,7 +2,7 @@ module Otrace := Trace_core
 
 module Traceparent : sig
   val name : string
-  val get_ambient : unit -> string option
+  val get_ambient : ?explicit_span:Trace_core.explicit_span -> unit -> string option
 end
 
 val enter_manual_span :
diff --git a/possibly_otel.real.ml b/possibly_otel.real.ml
index a54bd6b..f56d67c 100644
--- a/possibly_otel.real.ml
+++ b/possibly_otel.real.ml
@@ -5,8 +5,12 @@ let (let*) o f = Option.map f o
 module Traceparent = struct
   let name = Trace_context.Traceparent.name
 
-  let get_ambient () =
+  let get_ambient ?explicit_span () =
     let* Scope.{ trace_id; span_id; _ } = Scope.get_ambient_scope () in
+    let span_id = match explicit_span with
+      | Some {Trace_core.span; _} -> Opentelemetry_trace.Internal.otel_of_otrace span
+      | None -> span_id 
+    in
     Trace_context.Traceparent.to_value ~trace_id ~parent_id:span_id ()
 end
 
diff --git a/possibly_otel.stub.ml b/possibly_otel.stub.ml
index 6722e21..3801463 100644
--- a/possibly_otel.stub.ml
+++ b/possibly_otel.stub.ml
@@ -1,7 +1,7 @@
 module Traceparent = struct
   let name = "traceparent"
 
-  let get_ambient () = None
+  let get_ambient ?explicit_span () = None [@@warning "-27"]
 end
 
 
diff --git a/web.ml b/web.ml
index d617e53..5222130 100644
--- a/web.ml
+++ b/web.ml
@@ -257,17 +257,7 @@ module Http (IO : IO_TYPE) (Curl_IO : CURL with type 'a t = 'a IO.t) : HTTP with
     let open Curl in
     let action_name = string_of_http_action action in
 
-    let headers = match Possibly_otel.Traceparent.get_ambient () with
-    | None -> headers
-    | Some value -> Some Possibly_otel.Traceparent.(add_if_absent ~name ~value headers)
-    in
-
-    let set_body_and_headers h ct body =
-      set_httpheader h (("Content-Type: "^ct) :: Option.default [] headers);
-      set_postfields h body;
-      set_postfieldsize h (String.length body)
-    in
-    let setup h =
+    let setup ~headers set_body_and_headers h =
       begin match body with
       | Some (`Form args) -> set_body_and_headers h "application/x-www-form-urlencoded" (make_url_args args)
       | Some (`Raw (ct,body)) -> set_body_and_headers h ct body
@@ -311,17 +301,28 @@ module Http (IO : IO_TYPE) (Curl_IO : CURL with type 'a t = 'a IO.t) : HTTP with
         "url.full", `String url;
       ]
     in
-    let sid = Possibly_otel.enter_manual_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data:describe action_name in
+    let explicit_span = Possibly_otel.enter_manual_span ~__FUNCTION__ ~__FILE__ ~__LINE__ ~data:describe action_name in
+
+    let headers = match Possibly_otel.Traceparent.get_ambient ~explicit_span () with
+    | None -> headers
+    | Some value -> Some (add_if_absent ~name:(Possibly_otel.Traceparent.name) ~value headers)
+    in
+
+    let set_body_and_headers h ct body =
+      set_httpheader h (("Content-Type: "^ct) :: Option.default [] headers);
+      set_postfields h body;
+      set_postfieldsize h (String.length body)
+    in
 
     let t = new Action.timer in
     let result = Some (fun h code ->
       if verbose then verbose_curl_result nr_http action t h code;
-      Trace_core.add_data_to_manual_span sid ["http.response.status_code", `Int (Curl.get_httpcode h)];
-      Trace_core.exit_manual_span sid;
+      Trace_core.add_data_to_manual_span explicit_span ["http.response.status_code", `Int (Curl.get_httpcode h)];
+      Trace_core.exit_manual_span explicit_span;
       return ()
     ) in
 
-    http_gets ~setup ?timer ?result ?max_size url
+    http_gets ~setup:(setup ~headers set_body_and_headers) ?timer ?result ?max_size url
 
   let http_request ?ua ?timeout ?verbose ?setup ?timer ?max_size ?http_1_0 ?headers ?body (action:http_action) url =
     http_request' ?ua ?timeout ?verbose ?setup ?timer ?max_size ?http_1_0 ?headers ?body action url >>= fun res ->