setup ambient context in lwt/eio client setups

src/ambient-context/basic_map.ml

@@ -23,17 +23,11 @@ open struct
     let tid = Thread.id @@ Thread.self () in
 
     let ctx_ref =
-      try Int_map.find tid (Atomic.get self.m)
+      Opentelemetry_util.Util_atomic.update_cas self.m @@ fun m ->
+      try Int_map.find tid m, m
       with Not_found ->
         let r = ref Context.empty in
-        while
+        r, Int_map.add tid r m
-          let m = Atomic.get self.m in
-          let m' = Int_map.add tid r m in
-          not (Atomic.compare_and_set self.m m m')
-        do
-          ()
-        done;
-        r
     in
 
     let old_ctx = !ctx_ref in

src/ambient-context/dune

@@ -14,6 +14,7 @@
   threads
   (re_export opentelemetry.ambient-context.core)
   (re_export opentelemetry.atomic)
+  (re_export opentelemetry.util)
   (select
    default_.ml
    from

src/client-cohttp-eio/dune

@@ -8,6 +8,7 @@
   (re_export opentelemetry)
   (re_export opentelemetry-client)
   (re_export opentelemetry-client.sync)
+  opentelemetry.ambient-context.eio
   (re_export eio)
   (re_export eio.core)
   (re_export eio.unix)

src/client-cohttp-eio/opentelemetry_client_cohttp_eio.ml

@@ -180,6 +180,8 @@ let create_exporter ?(config = Config.make ()) ~sw ~env () =
 let create_backend = create_exporter
 
 let setup_ ~sw ?config env : unit =
+  Opentelemetry_ambient_context.set_current_storage
+    Opentelemetry_ambient_context_eio.storage;
   let exp = create_exporter ?config ~sw ~env () in
   Main_exporter.set exp
 

src/client-cohttp-lwt/opentelemetry_client_cohttp_lwt.ml

@@ -120,6 +120,7 @@ let create_exporter ?(config = Config.make ()) () =
 let create_backend = create_exporter
 
 let setup_ ?config () : unit =
+  Opentelemetry_client_lwt.Util_ambient_context.setup_ambient_context ();
   let exp = create_exporter ?config () in
   Main_exporter.set exp;
   ()

src/client-ocurl-lwt/opentelemetry_client_ocurl_lwt.ml

@@ -91,6 +91,7 @@ let create_exporter ?(config = Config.make ()) () =
 let create_backend = create_exporter
 
 let setup_ ?config () : Exporter.t =
+  Opentelemetry_client_lwt.Util_ambient_context.setup_ambient_context ();
   let exp = create_exporter ?config () in
   Main_exporter.set exp;
   exp

src/client/lwt/dune

@@ -16,6 +16,8 @@
   opentelemetry.emitter
   (re_export opentelemetry.core)
   (re_export opentelemetry)
+  (re_export opentelemetry.ambient-context)
+  opentelemetry.ambient-context.lwt
   (re_export opentelemetry-client)
   (re_export lwt)
   threads

src/client/lwt/util_ambient_context.ml

@@ -0,0 +1,4 @@
+(** Setup Lwt as the ambient context *)
+let setup_ambient_context () =
+  Opentelemetry_ambient_context.set_current_storage
+    Opentelemetry_ambient_context_lwt.storage

src/lwt/dune

@@ -2,8 +2,7 @@
  (name opentelemetry_lwt)
  (public_name opentelemetry-lwt)
  (synopsis "Lwt frontend for opentelemetry")
- (preprocess
-  (pps lwt_ppx))
  (libraries
   (re_export lwt)
+  opentelemetry.ambient-context.lwt
   (re_export opentelemetry)))

src/lwt/opentelemetry_lwt.ml

@@ -1,5 +1,10 @@
 include Opentelemetry
 
+(** Setup Lwt as the ambient context *)
+let setup_ambient_context () =
+  Opentelemetry_ambient_context.set_current_storage
+    Opentelemetry_ambient_context_lwt.storage
+
 module Main_exporter = struct
   include Main_exporter