(new lwt) Use Lwt's SAS instead of Thread_local

This stores the implicit 'global' parent-scope off the callstack in Lwt's "sequence-associated storage" instead of our own Thread_local, as this makes more sense in the context of green threads.

src/lwt/opentelemetry_lwt.ml

@@ -11,13 +11,44 @@ module GC_metrics = GC_metrics
 module Metrics_callbacks = Metrics_callbacks
 module Trace_context = Trace_context
 
+module Scope = struct
+  include Scope
+
+  (**/**)
+
+  let _global_scope_key : t Lwt.key = Lwt.new_key ()
+
+  (**/**)
+
+  let get_surrounding ?scope () : t option =
+    let surrounding = Lwt.get _global_scope_key in
+    match scope, surrounding with
+    | Some _, _ -> scope
+    | None, Some _ -> surrounding
+    | None, None -> None
+
+  let[@inline] with_scope (sc : t) (f : unit -> 'a) : 'a =
+    Lwt.with_value _global_scope_key (Some sc) f
+end
+
+open struct
+  let get_surrounding_scope = Scope.get_surrounding
+end
+
 module Trace = struct
   open Proto.Trace
   include Trace
 
   (** Sync span guard *)
-  let with_ ?trace_state ?service_name ?(attrs = []) ?kind ?trace_id ?parent
+  let with_ ?(force_new_trace_id = false) ?trace_state ?service_name
-      ?scope ?links name (f : Scope.t -> 'a Lwt.t) : 'a Lwt.t =
+      ?(attrs = []) ?kind ?trace_id ?parent ?scope ?links name
+      (f : Scope.t -> 'a Lwt.t) : 'a Lwt.t =
+    let scope =
+      if force_new_trace_id then
+        None
+      else
+        get_surrounding_scope ?scope ()
+    in
     let trace_id =
       match trace_id, scope with
       | Some trace_id, _ -> trace_id
@@ -33,6 +64,8 @@ module Trace = struct
     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
+    (* set global scope in this thread *)
+    Scope.with_scope scope @@ fun () ->
     let finally ok =
       let status =
         match ok with