From d92a9be08adef79dd515432f1b30f21900759b7f Mon Sep 17 00:00:00 2001
From: Matt Bray <mattjbray@gmail.com>
Date: Fri, 25 Mar 2022 15:56:01 +0000
Subject: [PATCH] feat: allow passing ~scope back to Trace.with_

---
 src/lwt/opentelemetry_lwt.ml | 14 +++++++++++++-
 src/opentelemetry.ml         | 15 +++++++++++++--
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/src/lwt/opentelemetry_lwt.ml b/src/lwt/opentelemetry_lwt.ml
index 0bc4367d..07f17e4f 100644
--- a/src/lwt/opentelemetry_lwt.ml
+++ b/src/lwt/opentelemetry_lwt.ml
@@ -15,8 +15,20 @@ module Trace = struct
   (** Sync span guard *)
   let with_
       ?trace_state ?service_name ?(attrs=[])
-      ?kind ?(trace_id=Trace_id.create()) ?parent ?links
+      ?kind ?trace_id ?parent ?scope ?links
       name (f:Trace.scope -> '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
diff --git a/src/opentelemetry.ml b/src/opentelemetry.ml
index ae606ab9..a7b53b09 100644
--- a/src/opentelemetry.ml
+++ b/src/opentelemetry.ml
@@ -463,9 +463,20 @@ module Trace = struct
   (** Sync span guard *)
   let with_
       ?trace_state ?service_name ?(attrs = [])
-      ?kind ?(trace_id=Trace_id.create()) ?parent ?links
+      ?kind ?trace_id ?parent ?scope ?links
       name (f: scope -> 'a) : 'a =
-
+    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