From 9584a7426f6ad8c16e1cba2ed3f0f4a30d2ab848 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 5 Sep 2024 16:11:09 -0400
Subject: [PATCH 1/5] wip: inline ambient-context into opentelemetry

---
 src/ambient-context/dune                      |  11 ++
 src/ambient-context/hmap_key_.new.ml          |   1 +
 src/ambient-context/hmap_key_.rcontext.ml     |   1 +
 .../opentelemetry_ambient_context.ml          | 124 ++++++++++++++++++
 .../opentelemetry_ambient_context.mli         |  54 ++++++++
 src/ambient-context/types/dune                |   4 +
 .../opentelemetry_ambient_context_types.ml    |  19 +++
 .../opentelemetry_ambient_context_types.mli   |  30 +++++
 8 files changed, 244 insertions(+)
 create mode 100644 src/ambient-context/dune
 create mode 100644 src/ambient-context/hmap_key_.new.ml
 create mode 100644 src/ambient-context/hmap_key_.rcontext.ml
 create mode 100644 src/ambient-context/opentelemetry_ambient_context.ml
 create mode 100644 src/ambient-context/opentelemetry_ambient_context.mli
 create mode 100644 src/ambient-context/types/dune
 create mode 100644 src/ambient-context/types/opentelemetry_ambient_context_types.ml
 create mode 100644 src/ambient-context/types/opentelemetry_ambient_context_types.mli

diff --git a/src/ambient-context/dune b/src/ambient-context/dune
new file mode 100644
index 00000000..f0f87c7d
--- /dev/null
+++ b/src/ambient-context/dune
@@ -0,0 +1,11 @@
+(library
+ (name opentelemetry_ambient_context)
+ (public_name opentelemetry.ambient-context)
+ (synopsis
+  "Abstraction over thread-local storage and fiber-local storage mechanisms")
+ (private_modules hmap_key_)
+ (libraries thread-local-storage threads atomic
+  opentelemetry.ambient-context.types
+  (select hmap_key_.ml from
+          (rcontext hmap -> hmap_key_.rcontext.ml)
+          (-> hmap_key_.new.ml))))
diff --git a/src/ambient-context/hmap_key_.new.ml b/src/ambient-context/hmap_key_.new.ml
new file mode 100644
index 00000000..1925b70e
--- /dev/null
+++ b/src/ambient-context/hmap_key_.new.ml
@@ -0,0 +1 @@
+let key : Hmap.t Thread_local_storage.t = Thread_local_storage.create ()
diff --git a/src/ambient-context/hmap_key_.rcontext.ml b/src/ambient-context/hmap_key_.rcontext.ml
new file mode 100644
index 00000000..99f136e6
--- /dev/null
+++ b/src/ambient-context/hmap_key_.rcontext.ml
@@ -0,0 +1 @@
+let key : Hmap.t Thread_local_storage.t = Rcontext.Ambient_hmap.k_hmap
diff --git a/src/ambient-context/opentelemetry_ambient_context.ml b/src/ambient-context/opentelemetry_ambient_context.ml
new file mode 100644
index 00000000..7c622eb7
--- /dev/null
+++ b/src/ambient-context/opentelemetry_ambient_context.ml
@@ -0,0 +1,124 @@
+module TLS = Thread_local_storage
+include Opentelemetry_ambient_context_types
+
+type 'a key = int * 'a Hmap.key
+
+let debug =
+  match Sys.getenv_opt "OCAML_AMBIENT_CONTEXT_DEBUG" with
+  | Some ("1" | "true") -> true
+  | _ -> false
+
+let _debug_id_ = Atomic.make 0
+
+let[@inline] generate_debug_id () = Atomic.fetch_and_add _debug_id_ 1
+
+let compare_key : int -> int -> int = Stdlib.compare
+
+module Storage_tls_hmap = struct
+  let[@inline] ( let* ) o f =
+    match o with
+    | None -> None
+    | Some x -> f x
+
+  let key : Hmap.t TLS.t = Hmap_key_.key
+
+  let name = "Storage_tls"
+
+  let[@inline] get_map () = TLS.get_opt key
+
+  let[@inline] with_map m cb =
+    let old = TLS.get_opt key |> Option.value ~default:Hmap.empty in
+    TLS.set key m;
+    Fun.protect ~finally:(fun () -> TLS.set key old) cb
+
+  let create_key = Hmap.Key.create
+
+  let get k =
+    let* context = get_map () in
+    Hmap.find k context
+
+  let with_binding k v cb =
+    let new_context =
+      match get_map () with
+      | None -> Hmap.singleton k v
+      | Some old_context -> Hmap.add k v old_context
+    in
+    with_map new_context @@ fun _context -> cb ()
+
+  let without_binding k cb =
+    match get_map () with
+    | None -> cb ()
+    | Some old_context ->
+      let new_context = Hmap.rem k old_context in
+      with_map new_context @@ fun _context -> cb ()
+end
+
+let default_storage : storage = (module Storage_tls_hmap)
+
+let k_current_storage : storage TLS.t = TLS.create ()
+
+let get_current_storage () =
+  match TLS.get_exn k_current_storage with
+  | v -> v
+  | exception TLS.Not_set ->
+    let v = default_storage in
+    TLS.set k_current_storage v;
+    v
+
+let create_key () =
+  let (module Store : STORAGE) = get_current_storage () in
+  if not debug then
+    0, Store.create_key ()
+  else (
+    let id = generate_debug_id () in
+    Printf.printf "%s: create_key %i\n%!" Store.name id;
+    id, Store.create_key ()
+  )
+
+let get (id, k) =
+  let (module Store : STORAGE) = get_current_storage () in
+  if not debug then
+    Store.get k
+  else (
+    let rv = Store.get k in
+    (match rv with
+    | Some _ -> Printf.printf "%s: get %i -> Some\n%!" Store.name id
+    | None -> Printf.printf "%s: get %i -> None\n%!" Store.name id);
+    rv
+  )
+
+let with_binding : 'a key -> 'a -> (unit -> 'r) -> 'r =
+ fun (id, k) v cb ->
+  let (module Store : STORAGE) = get_current_storage () in
+  if not debug then
+    Store.with_binding k v cb
+  else (
+    Printf.printf "%s: with_binding %i enter\n%!" Store.name id;
+    let rv = Store.with_binding k v cb in
+    Printf.printf "%s: with_binding %i exit\n%!" Store.name id;
+    rv
+  )
+
+let without_binding (id, k) cb =
+  let (module Store : STORAGE) = get_current_storage () in
+  if not debug then
+    Store.without_binding k cb
+  else (
+    Printf.printf "%s: without_binding %i enter\n%!" Store.name id;
+    let rv = Store.without_binding k cb in
+    Printf.printf "%s: without_binding %i exit\n%!" Store.name id;
+    rv
+  )
+
+let set_storage_provider store_new =
+  let store_before = get_current_storage () in
+  if store_new == store_before then
+    ()
+  else
+    TLS.set k_current_storage store_new;
+  if debug then (
+    let (module Store_before : STORAGE) = store_before in
+    let (module Store_new : STORAGE) = store_new in
+    Printf.printf "set_storage_provider %s (previously %s)\n%!" Store_new.name
+      Store_before.name
+  )
diff --git a/src/ambient-context/opentelemetry_ambient_context.mli b/src/ambient-context/opentelemetry_ambient_context.mli
new file mode 100644
index 00000000..4cb1cc51
--- /dev/null
+++ b/src/ambient-context/opentelemetry_ambient_context.mli
@@ -0,0 +1,54 @@
+(** Ambient context.
+
+    The ambient context, like the Matrix, is everywhere around you.
+
+    It is responsible for keeping track of that context in a manner that's consistent with
+    the program's choice of control flow paradigm:
+
+    - for synchronous/threaded/direct style code, {b TLS} ("thread local storage") keeps
+      track of a global variable per thread. Each thread has its own copy of the variable
+      and updates it independently of other threads.
+
+    - for Lwt, any ['a Lwt.t] created inside the [with_binding k v (fun _ -> …)] will
+      inherit the [k := v] assignment.
+
+    - for Eio, fibers created inside [with_binding k v (fun () -> …)] will inherit the
+      [k := v] assignment. This is consistent with the structured concurrency approach of
+      Eio.
+
+    The only data stored by this storage is a {!Hmap.t}, ie a heterogeneous map. Various
+    users (libraries, user code, etc.) can create their own {!key} to store what they are
+    interested in, without affecting other parts of the storage. *)
+
+module Types := Opentelemetry_ambient_context_types
+
+module type STORAGE = Types.STORAGE
+
+type storage = (module STORAGE)
+
+val default_storage : storage
+
+val get_current_storage : unit -> storage
+
+val set_storage_provider : storage -> unit
+
+type 'a key
+(** A key that can be mapped to values of type ['a] in the ambient context. *)
+
+val compare_key : int -> int -> int
+(** Total order on keys *)
+
+val create_key : unit -> 'a key
+(** Create a new fresh key, distinct from any previously created key. *)
+
+val get : 'a key -> 'a option
+(** Get the current value for a given key, or [None] if no value was associated with the
+    key in the ambient context. *)
+
+val with_binding : 'a key -> 'a -> (unit -> 'r) -> 'r
+(** [with_binding k v cb] calls [cb()] in a context in which [k] is bound to [v]. This
+    does not affect storage outside of [cb()]. *)
+
+val without_binding : 'a key -> (unit -> 'b) -> 'b
+(** [without_binding k cb] calls [cb()] in a context where [k] has no binding (possibly
+    shadowing the current ambient binding of [k] if it exists). *)
diff --git a/src/ambient-context/types/dune b/src/ambient-context/types/dune
new file mode 100644
index 00000000..b9e4146c
--- /dev/null
+++ b/src/ambient-context/types/dune
@@ -0,0 +1,4 @@
+(library
+ (name opentelemetry_ambient_context_types)
+ (public_name opentelemetry.ambient-context.types)
+ (libraries hmap thread-local-storage))
diff --git a/src/ambient-context/types/opentelemetry_ambient_context_types.ml b/src/ambient-context/types/opentelemetry_ambient_context_types.ml
new file mode 100644
index 00000000..829f7789
--- /dev/null
+++ b/src/ambient-context/types/opentelemetry_ambient_context_types.ml
@@ -0,0 +1,19 @@
+type 'a key = 'a Hmap.key
+
+module type STORAGE = sig
+  val name : string
+
+  val get_map : unit -> Hmap.t option
+
+  val with_map : Hmap.t -> (unit -> 'b) -> 'b
+
+  val create_key : unit -> 'a key
+
+  val get : 'a key -> 'a option
+
+  val with_binding : 'a key -> 'a -> (unit -> 'b) -> 'b
+
+  val without_binding : 'a key -> (unit -> 'b) -> 'b
+end
+
+type storage = (module STORAGE)
diff --git a/src/ambient-context/types/opentelemetry_ambient_context_types.mli b/src/ambient-context/types/opentelemetry_ambient_context_types.mli
new file mode 100644
index 00000000..cded6589
--- /dev/null
+++ b/src/ambient-context/types/opentelemetry_ambient_context_types.mli
@@ -0,0 +1,30 @@
+(** Storage implementation.
+
+    There is a singleton storage for a given program, responsible for providing ambient
+    context to the rest of the program. *)
+
+type 'a key = 'a Hmap.key
+
+module type STORAGE = sig
+  val name : string
+  (** Name of the storage implementation. *)
+
+  val get_map : unit -> Hmap.t option
+  (** Get the hmap from the current ambient context, or [None] if there is no ambient
+      context. *)
+
+  val with_map : Hmap.t -> (unit -> 'b) -> 'b
+  (** [with_hmap h cb] calls [cb()] in an ambient context in which [get_map()] will return
+      [h]. Once [cb()] returns, the storage is reset to its previous value. *)
+
+  val create_key : unit -> 'a key
+  (** Create a new storage key, guaranteed to be distinct from any previously created key. *)
+
+  val get : 'a key -> 'a option
+
+  val with_binding : 'a key -> 'a -> (unit -> 'b) -> 'b
+
+  val without_binding : 'a key -> (unit -> 'b) -> 'b
+end
+
+type storage = (module STORAGE)

From e8ed97100beec0b72578319f1317bf8ea2015229 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 6 Sep 2024 07:48:47 -0400
Subject: [PATCH 2/5] ocamlformat

---
 .ocamlformat | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.ocamlformat b/.ocamlformat
index e74cb7f7..87d74bf8 100644
--- a/.ocamlformat
+++ b/.ocamlformat
@@ -1,4 +1,4 @@
-version = 0.24.1
+version = 0.26.2
 profile=conventional
 margin=80
 if-then-else=k-r

From faa0808034f00b1290b2918c4c6a648fb3c29b24 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 6 Sep 2024 07:48:54 -0400
Subject: [PATCH 3/5] wip

---
 dune-project                                  |  3 +-
 .../eio/ambient_context_eio.ml                | 37 +++++++++++++++++++
 src/ambient-context/eio/dune.tmp              |  6 +++
 .../lwt/ambient_context_lwt.ml                | 35 ++++++++++++++++++
 src/ambient-context/lwt/dune.tmp              |  6 +++
 src/core/dune                                 |  2 +-
 src/core/opentelemetry.ml                     |  2 +
 7 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 src/ambient-context/eio/ambient_context_eio.ml
 create mode 100644 src/ambient-context/eio/dune.tmp
 create mode 100644 src/ambient-context/lwt/ambient_context_lwt.ml
 create mode 100644 src/ambient-context/lwt/dune.tmp

diff --git a/dune-project b/dune-project
index 65fef54e..a2153000 100644
--- a/dune-project
+++ b/dune-project
@@ -28,7 +28,8 @@
    (>= "4.08"))
   ptime
   hmap
-  ambient-context
+  atomic
+  (thread-local-storage (and (>= 0.2) (< 0.3)))
   (odoc :with-doc)
   (alcotest :with-test)
   (pbrt
diff --git a/src/ambient-context/eio/ambient_context_eio.ml b/src/ambient-context/eio/ambient_context_eio.ml
new file mode 100644
index 00000000..642be8c8
--- /dev/null
+++ b/src/ambient-context/eio/ambient_context_eio.ml
@@ -0,0 +1,37 @@
+module TLS = Ambient_context_thread_local.Thread_local
+module Hmap = Ambient_context.Hmap
+module Fiber = Eio.Fiber
+
+let _internal_key : Hmap.t Fiber.key = Fiber.create_key ()
+let ( let* ) = Option.bind
+
+module M = struct
+  let name = "Storage_eio"
+  let[@inline] get_map () = Fiber.get _internal_key
+  let[@inline] with_map m cb = Fiber.with_binding _internal_key m cb
+  let create_key = Hmap.Key.create
+
+  let get k =
+     let* context = get_map () in
+     Hmap.find k context
+
+
+  let with_binding k v cb =
+     let new_context =
+        match get_map () with
+        | None -> Hmap.singleton k v
+        | Some old_context -> Hmap.add k v old_context
+     in
+     with_map new_context cb
+
+
+  let without_binding k cb =
+     let new_context =
+        match get_map () with
+        | None -> Hmap.empty
+        | Some old_context -> Hmap.rem k old_context
+     in
+     with_map new_context cb
+end
+
+let storage () : Ambient_context.storage = (module M)
diff --git a/src/ambient-context/eio/dune.tmp b/src/ambient-context/eio/dune.tmp
new file mode 100644
index 00000000..cb379f9a
--- /dev/null
+++ b/src/ambient-context/eio/dune.tmp
@@ -0,0 +1,6 @@
+(library
+ (name ambient_context_eio)
+ (synopsis
+  "Storage backend for ambient-context using Eio's fibre-local storage")
+ (public_name ambient-context-eio)
+ (libraries eio ambient-context ambient-context.thread_local))
diff --git a/src/ambient-context/lwt/ambient_context_lwt.ml b/src/ambient-context/lwt/ambient_context_lwt.ml
new file mode 100644
index 00000000..6010d02a
--- /dev/null
+++ b/src/ambient-context/lwt/ambient_context_lwt.ml
@@ -0,0 +1,35 @@
+module Hmap = Ambient_context.Hmap
+
+let _internal_key : Hmap.t Lwt.key = Lwt.new_key ()
+let ( let* ) = Option.bind
+
+module M = struct
+  let name = "Storage_lwt"
+  let[@inline] get_map () = Lwt.get _internal_key
+  let[@inline] with_map m cb = Lwt.with_value _internal_key (Some m) cb
+  let create_key = Hmap.Key.create
+
+  let get k =
+     let* context = get_map () in
+     Hmap.find k context
+
+
+  let with_binding k v cb =
+     let new_context =
+        match get_map () with
+        | None -> Hmap.singleton k v
+        | Some old_context -> Hmap.add k v old_context
+     in
+     with_map new_context cb
+
+
+  let without_binding k cb =
+     let new_context =
+        match get_map () with
+        | None -> Hmap.empty
+        | Some old_context -> Hmap.rem k old_context
+     in
+     with_map new_context cb
+end
+
+let storage () : Ambient_context.storage = (module M)
diff --git a/src/ambient-context/lwt/dune.tmp b/src/ambient-context/lwt/dune.tmp
new file mode 100644
index 00000000..5dae2dd9
--- /dev/null
+++ b/src/ambient-context/lwt/dune.tmp
@@ -0,0 +1,6 @@
+(library
+ (name ambient_context_lwt)
+ (synopsis
+  "Storage backend for ambient-context using Lwt's sequence-associated storage")
+ (public_name ambient-context-lwt)
+ (libraries lwt ambient-context ambient-context.tls))
diff --git a/src/core/dune b/src/core/dune
index f442b5ac..235f8ad9 100644
--- a/src/core/dune
+++ b/src/core/dune
@@ -2,6 +2,6 @@
  (name opentelemetry)
  (synopsis "API for opentelemetry instrumentation")
  (flags :standard -warn-error -a+8)
- (libraries opentelemetry.proto ambient-context ptime ptime.clock.os pbrt threads
+ (libraries opentelemetry.proto opentelemetry.ambient-context ptime ptime.clock.os pbrt threads
    opentelemetry.atomic hmap)
  (public_name opentelemetry))
diff --git a/src/core/opentelemetry.ml b/src/core/opentelemetry.ml
index 7cf2e2c7..be4135f0 100644
--- a/src/core/opentelemetry.ml
+++ b/src/core/opentelemetry.ml
@@ -2,6 +2,8 @@
 
 open struct
   let spf = Printf.sprintf
+
+  module Ambient_context = Opentelemetry_ambient_context
 end
 
 module Lock = Lock

From 1bcea95ed91397cdc77f7214ae8a28c11d5f0f73 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 6 Sep 2024 09:43:56 -0400
Subject: [PATCH 4/5] feat: lwt backend

---
 dune-project                                  | 14 +++++--
 opentelemetry.opam                            |  5 ++-
 .../lwt/ambient_context_lwt.ml                | 35 ------------------
 src/ambient-context/lwt/dune                  |  7 ++++
 src/ambient-context/lwt/dune.tmp              |  6 ---
 .../lwt/opentelemetry_ambient_context_lwt.ml  | 37 +++++++++++++++++++
 .../lwt/opentelemetry_ambient_context_lwt.mli |  2 +
 7 files changed, 59 insertions(+), 47 deletions(-)
 delete mode 100644 src/ambient-context/lwt/ambient_context_lwt.ml
 create mode 100644 src/ambient-context/lwt/dune
 delete mode 100644 src/ambient-context/lwt/dune.tmp
 create mode 100644 src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml
 create mode 100644 src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli

diff --git a/dune-project b/dune-project
index a2153000..ffe64d74 100644
--- a/dune-project
+++ b/dune-project
@@ -29,20 +29,26 @@
   ptime
   hmap
   atomic
-  (thread-local-storage (and (>= 0.2) (< 0.3)))
+  (thread-local-storage
+   (and
+    (>= 0.2)
+    (< 0.3)))
   (odoc :with-doc)
   (alcotest :with-test)
   (pbrt
-   (and (>= 3.0) (< 4.0)))
+   (and
+    (>= 3.0)
+    (< 4.0)))
   (ocaml-lsp-server :with-dev-setup)
   (ocamlformat
    (and
     :with-dev-setup
     (>= 0.24)
     (< 0.25))))
- (depopts trace)
+ (depopts trace lwt eio)
  (conflicts
-  (trace (< 0.7)))
+  (trace
+   (< 0.7)))
  (tags
   (instrumentation tracing opentelemetry datadog jaeger)))
 
diff --git a/opentelemetry.opam b/opentelemetry.opam
index 916dd6ad..80da4b91 100644
--- a/opentelemetry.opam
+++ b/opentelemetry.opam
@@ -17,14 +17,15 @@ depends: [
   "ocaml" {>= "4.08"}
   "ptime"
   "hmap"
-  "ambient-context"
+  "atomic"
+  "thread-local-storage" {>= "0.2" & < "0.3"}
   "odoc" {with-doc}
   "alcotest" {with-test}
   "pbrt" {>= "3.0" & < "4.0"}
   "ocaml-lsp-server" {with-dev-setup}
   "ocamlformat" {with-dev-setup & >= "0.24" & < "0.25"}
 ]
-depopts: ["trace"]
+depopts: ["trace" "lwt" "eio"]
 conflicts: [
   "trace" {< "0.7"}
 ]
diff --git a/src/ambient-context/lwt/ambient_context_lwt.ml b/src/ambient-context/lwt/ambient_context_lwt.ml
deleted file mode 100644
index 6010d02a..00000000
--- a/src/ambient-context/lwt/ambient_context_lwt.ml
+++ /dev/null
@@ -1,35 +0,0 @@
-module Hmap = Ambient_context.Hmap
-
-let _internal_key : Hmap.t Lwt.key = Lwt.new_key ()
-let ( let* ) = Option.bind
-
-module M = struct
-  let name = "Storage_lwt"
-  let[@inline] get_map () = Lwt.get _internal_key
-  let[@inline] with_map m cb = Lwt.with_value _internal_key (Some m) cb
-  let create_key = Hmap.Key.create
-
-  let get k =
-     let* context = get_map () in
-     Hmap.find k context
-
-
-  let with_binding k v cb =
-     let new_context =
-        match get_map () with
-        | None -> Hmap.singleton k v
-        | Some old_context -> Hmap.add k v old_context
-     in
-     with_map new_context cb
-
-
-  let without_binding k cb =
-     let new_context =
-        match get_map () with
-        | None -> Hmap.empty
-        | Some old_context -> Hmap.rem k old_context
-     in
-     with_map new_context cb
-end
-
-let storage () : Ambient_context.storage = (module M)
diff --git a/src/ambient-context/lwt/dune b/src/ambient-context/lwt/dune
new file mode 100644
index 00000000..68a9de15
--- /dev/null
+++ b/src/ambient-context/lwt/dune
@@ -0,0 +1,7 @@
+(library
+ (name opentelemetry_ambient_context_lwt)
+ (public_name opentelemetry.ambient-context.lwt)
+ (optional) ; lwt
+ (synopsis
+  "Storage backend for ambient-context using Lwt's sequence-associated storage")
+ (libraries lwt opentelemetry.ambient-context thread-local-storage))
diff --git a/src/ambient-context/lwt/dune.tmp b/src/ambient-context/lwt/dune.tmp
deleted file mode 100644
index 5dae2dd9..00000000
--- a/src/ambient-context/lwt/dune.tmp
+++ /dev/null
@@ -1,6 +0,0 @@
-(library
- (name ambient_context_lwt)
- (synopsis
-  "Storage backend for ambient-context using Lwt's sequence-associated storage")
- (public_name ambient-context-lwt)
- (libraries lwt ambient-context ambient-context.tls))
diff --git a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml
new file mode 100644
index 00000000..b75105f2
--- /dev/null
+++ b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml
@@ -0,0 +1,37 @@
+open struct
+  let _internal_key : Hmap.t Lwt.key = Lwt.new_key ()
+
+  let ( let* ) = Option.bind
+end
+
+module M = struct
+  let name = "Storage_lwt"
+
+  let[@inline] get_map () = Lwt.get _internal_key
+
+  let[@inline] with_map m cb = Lwt.with_value _internal_key (Some m) cb
+
+  let create_key = Hmap.Key.create
+
+  let get k =
+    let* context = get_map () in
+    Hmap.find k context
+
+  let with_binding k v cb =
+    let new_context =
+      match get_map () with
+      | None -> Hmap.singleton k v
+      | Some old_context -> Hmap.add k v old_context
+    in
+    with_map new_context cb
+
+  let without_binding k cb =
+    let new_context =
+      match get_map () with
+      | None -> Hmap.empty
+      | Some old_context -> Hmap.rem k old_context
+    in
+    with_map new_context cb
+end
+
+let storage () : Opentelemetry_ambient_context.storage = (module M)
diff --git a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli
new file mode 100644
index 00000000..3c462a8d
--- /dev/null
+++ b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli
@@ -0,0 +1,2 @@
+val storage : unit -> Opentelemetry_ambient_context.storage
+(** Storage using Lwt keys *)

From a9971e4d416442aaf279970b684be0a2838317b6 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Fri, 6 Sep 2024 09:48:06 -0400
Subject: [PATCH 5/5] eio local storage

---
 .../eio/ambient_context_eio.ml                | 37 -----------------
 src/ambient-context/eio/dune                  |  7 ++++
 src/ambient-context/eio/dune.tmp              |  6 ---
 .../eio/opentelemetry_ambient_context_eio.ml  | 40 +++++++++++++++++++
 .../eio/opentelemetry_ambient_context_eio.mli |  2 +
 5 files changed, 49 insertions(+), 43 deletions(-)
 delete mode 100644 src/ambient-context/eio/ambient_context_eio.ml
 create mode 100644 src/ambient-context/eio/dune
 delete mode 100644 src/ambient-context/eio/dune.tmp
 create mode 100644 src/ambient-context/eio/opentelemetry_ambient_context_eio.ml
 create mode 100644 src/ambient-context/eio/opentelemetry_ambient_context_eio.mli

diff --git a/src/ambient-context/eio/ambient_context_eio.ml b/src/ambient-context/eio/ambient_context_eio.ml
deleted file mode 100644
index 642be8c8..00000000
--- a/src/ambient-context/eio/ambient_context_eio.ml
+++ /dev/null
@@ -1,37 +0,0 @@
-module TLS = Ambient_context_thread_local.Thread_local
-module Hmap = Ambient_context.Hmap
-module Fiber = Eio.Fiber
-
-let _internal_key : Hmap.t Fiber.key = Fiber.create_key ()
-let ( let* ) = Option.bind
-
-module M = struct
-  let name = "Storage_eio"
-  let[@inline] get_map () = Fiber.get _internal_key
-  let[@inline] with_map m cb = Fiber.with_binding _internal_key m cb
-  let create_key = Hmap.Key.create
-
-  let get k =
-     let* context = get_map () in
-     Hmap.find k context
-
-
-  let with_binding k v cb =
-     let new_context =
-        match get_map () with
-        | None -> Hmap.singleton k v
-        | Some old_context -> Hmap.add k v old_context
-     in
-     with_map new_context cb
-
-
-  let without_binding k cb =
-     let new_context =
-        match get_map () with
-        | None -> Hmap.empty
-        | Some old_context -> Hmap.rem k old_context
-     in
-     with_map new_context cb
-end
-
-let storage () : Ambient_context.storage = (module M)
diff --git a/src/ambient-context/eio/dune b/src/ambient-context/eio/dune
new file mode 100644
index 00000000..f3f76be7
--- /dev/null
+++ b/src/ambient-context/eio/dune
@@ -0,0 +1,7 @@
+(library
+ (name opentelemetry_ambient_context_eio)
+ (public_name opentelemetry.ambient-context.eio)
+ (synopsis
+  "Storage backend for ambient-context using Eio's fibre-local storage")
+ (optional) ; eio
+ (libraries eio hmap opentelemetry.ambient-context thread-local-storage))
diff --git a/src/ambient-context/eio/dune.tmp b/src/ambient-context/eio/dune.tmp
deleted file mode 100644
index cb379f9a..00000000
--- a/src/ambient-context/eio/dune.tmp
+++ /dev/null
@@ -1,6 +0,0 @@
-(library
- (name ambient_context_eio)
- (synopsis
-  "Storage backend for ambient-context using Eio's fibre-local storage")
- (public_name ambient-context-eio)
- (libraries eio ambient-context ambient-context.thread_local))
diff --git a/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml b/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml
new file mode 100644
index 00000000..41f682e5
--- /dev/null
+++ b/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml
@@ -0,0 +1,40 @@
+module TLS = Thread_local_storage
+module Fiber = Eio.Fiber
+
+open struct
+  let _internal_key : Hmap.t Fiber.key = Fiber.create_key ()
+
+  let ( let* ) = Option.bind
+end
+
+module M = struct
+  let name = "Storage_eio"
+
+  let[@inline] get_map () = Fiber.get _internal_key
+
+  let[@inline] with_map m cb = Fiber.with_binding _internal_key m cb
+
+  let create_key = Hmap.Key.create
+
+  let get k =
+    let* context = get_map () in
+    Hmap.find k context
+
+  let with_binding k v cb =
+    let new_context =
+      match get_map () with
+      | None -> Hmap.singleton k v
+      | Some old_context -> Hmap.add k v old_context
+    in
+    with_map new_context cb
+
+  let without_binding k cb =
+    let new_context =
+      match get_map () with
+      | None -> Hmap.empty
+      | Some old_context -> Hmap.rem k old_context
+    in
+    with_map new_context cb
+end
+
+let storage () : Opentelemetry_ambient_context.storage = (module M)
diff --git a/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli b/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli
new file mode 100644
index 00000000..ac5cf8ba
--- /dev/null
+++ b/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli
@@ -0,0 +1,2 @@
+val storage : unit -> Opentelemetry_ambient_context.storage
+(** Storage using Eio's fibers local storage *)