From d668f5c4727917b10bd15a89d41cc1f880e6ab64 Mon Sep 17 00:00:00 2001
From: Elliott Cable <elliott.cable@ahrefs.com>
Date: Tue, 1 Aug 2023 21:42:35 +0000
Subject: [PATCH] lwt: Share impl details with non-lwt with_

---
 src/lwt/opentelemetry_lwt.ml | 42 ++++++++----------------------------
 src/opentelemetry.ml         | 34 ++++++++++++++++++-----------
 2 files changed, 31 insertions(+), 45 deletions(-)

diff --git a/src/lwt/opentelemetry_lwt.ml b/src/lwt/opentelemetry_lwt.ml
index 8c50e284..887fe6f7 100644
--- a/src/lwt/opentelemetry_lwt.ml
+++ b/src/lwt/opentelemetry_lwt.ml
@@ -12,45 +12,21 @@ module Metrics_callbacks = Metrics_callbacks
 module Trace_context = Trace_context
 
 module Trace = struct
-  open Proto.Trace
   include Trace
 
   (** Sync span guard *)
-  let with_ ?trace_state ?service_name ?(attrs = []) ?kind ?trace_id ?parent
-      ?scope ?links name (f : Scope.t -> 'a Lwt.t) : 'a Lwt.t =
-    let trace_id =
-      match trace_id, scope with
-      | Some trace_id, _ -> trace_id
-      | None, Some scope -> scope.trace_id
-      | None, None -> Trace_id.create ()
-    in
-    let parent =
-      match parent, scope with
-      | Some span_id, _ -> Some span_id
-      | None, Some scope -> Some scope.span_id
-      | None, None -> None
-    in
-    let start_time = Timestamp_ns.now_unix_ns () in
-    let span_id = Span_id.create () in
-    let scope = { trace_id; span_id; events = []; attrs } in
-    let finally ok =
-      let status =
-        match ok with
-        | Ok () -> default_status ~code:Status_code_ok ()
-        | Error e -> default_status ~code:Status_code_error ~message:e ()
-      in
-      let span, _ =
-        Span.create ?kind ~trace_id ?parent ?links ~id:span_id ?trace_state
-          ~attrs:scope.attrs ~events:scope.events ~start_time
-          ~end_time:(Timestamp_ns.now_unix_ns ())
-          ~status name
-      in
-      emit ?service_name [ span ]
+  let with_ ?force_new_trace_id ?trace_state ?service_name ?attrs ?kind
+      ?trace_id ?parent ?scope ?links name (cb : Scope.t -> 'a Lwt.t) : 'a Lwt.t
+      =
+    let thunk, finally =
+      with_' ?force_new_trace_id ?trace_state ?service_name ?attrs ?kind
+        ?trace_id ?parent ?scope ?links name cb
     in
+
     try%lwt
-      let* x = f scope in
+      let* rv = thunk () in
       let () = finally (Ok ()) in
-      Lwt.return x
+      Lwt.return rv
     with e ->
       let () = finally (Error (Printexc.to_string e)) in
       Lwt.fail e
diff --git a/src/opentelemetry.ml b/src/opentelemetry.ml
index 60d374fa..6d27bf1d 100644
--- a/src/opentelemetry.ml
+++ b/src/opentelemetry.ml
@@ -716,17 +716,9 @@ module Trace = struct
 
   let add_attrs = Scope.add_attrs [@@deprecated "use Scope.add_attrs"]
 
-  (** Sync span guard.
-
-      @param force_new_trace_id if true (default false), the span will not use a
-      surrounding context, or [scope], or [trace_id], but will always
-      create a fresh new trace ID.
-
-      {b NOTE} be careful not to call this inside a Gc alarm, as it can
-      cause deadlocks. *)
-  let with_ ?(force_new_trace_id = false) ?trace_state ?service_name
+  let with_' ?(force_new_trace_id = false) ?trace_state ?service_name
       ?(attrs : (string * [< value ]) list = []) ?kind ?trace_id ?parent ?scope
-      ?links name (f : Scope.t -> 'a) : 'a =
+      ?links name cb =
     let scope =
       if force_new_trace_id then
         None
@@ -770,10 +762,28 @@ module Trace = struct
       in
       emit ?service_name [ span ]
     in
+    let thunk () = cb scope in
+    thunk, finally
+
+  (** Sync span guard.
+
+      @param force_new_trace_id if true (default false), the span will not use a
+      surrounding context, or [scope], or [trace_id], but will always
+      create a fresh new trace ID.
+
+      {b NOTE} be careful not to call this inside a Gc alarm, as it can
+      cause deadlocks. *)
+  let with_ ?force_new_trace_id ?trace_state ?service_name ?attrs ?kind
+      ?trace_id ?parent ?scope ?links name (cb : Scope.t -> 'a) : 'a =
+    let thunk, finally =
+      with_' ?force_new_trace_id ?trace_state ?service_name ?attrs ?kind
+        ?trace_id ?parent ?scope ?links name cb
+    in
+
     try
-      let x = f scope in
+      let rv = thunk () in
       finally (Ok ());
-      x
+      rv
     with e ->
       finally (Error (Printexc.to_string e));
       raise e