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] 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 *)