From fe8316d1e88a560353e952a3c8dcfe0e96ee3904 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 10 Dec 2025 14:55:52 -0500
Subject: [PATCH] add picos backend to ambient-context

---
 dune-project                                  |  2 +-
 opentelemetry.opam                            |  2 +-
 src/ambient-context/picos/dune                |  6 ++++
 .../opentelemetry_ambient_context_picos.ml    | 34 +++++++++++++++++++
 4 files changed, 42 insertions(+), 2 deletions(-)
 create mode 100644 src/ambient-context/picos/dune
 create mode 100644 src/ambient-context/picos/opentelemetry_ambient_context_picos.ml

diff --git a/dune-project b/dune-project
index 41888c90..aaf19572 100644
--- a/dune-project
+++ b/dune-project
@@ -42,7 +42,7 @@
     (< 0.28)))
   (mtime
    (>= "1.4")))
- (depopts atomic trace thread-local-storage lwt eio)
+ (depopts atomic trace thread-local-storage lwt eio picos)
  (conflicts
   (trace
    (< 0.10)))
diff --git a/opentelemetry.opam b/opentelemetry.opam
index 5d2d150c..79127ff3 100644
--- a/opentelemetry.opam
+++ b/opentelemetry.opam
@@ -24,7 +24,7 @@ depends: [
   "ocamlformat" {with-dev-setup & >= "0.27" & < "0.28"}
   "mtime" {>= "1.4"}
 ]
-depopts: ["atomic" "trace" "thread-local-storage" "lwt" "eio"]
+depopts: ["atomic" "trace" "thread-local-storage" "lwt" "eio" "picos"]
 conflicts: [
   "trace" {< "0.10"}
 ]
diff --git a/src/ambient-context/picos/dune b/src/ambient-context/picos/dune
new file mode 100644
index 00000000..d9f1243b
--- /dev/null
+++ b/src/ambient-context/picos/dune
@@ -0,0 +1,6 @@
+(library
+ (name opentelemetry_ambient_context_picos)
+ (public_name opentelemetry.ambient-context.picos)
+ (optional) ; picos
+ (synopsis "Storage backend for ambient-context using Picos' FLS")
+ (libraries picos opentelemetry.ambient-context.core))
diff --git a/src/ambient-context/picos/opentelemetry_ambient_context_picos.ml b/src/ambient-context/picos/opentelemetry_ambient_context_picos.ml
new file mode 100644
index 00000000..30244d42
--- /dev/null
+++ b/src/ambient-context/picos/opentelemetry_ambient_context_picos.ml
@@ -0,0 +1,34 @@
+(** Storage using Lwt keys *)
+
+open Opentelemetry_ambient_context_core
+
+open struct
+  module FLS = Picos.Fiber.FLS
+
+  let fls_context_key : Context.t FLS.t = FLS.create ()
+
+  let get_context () =
+    try FLS.get_exn (Picos.Fiber.current ()) fls_context_key
+    with _ -> Hmap.empty
+
+  let with_context ctx f =
+    match Picos.Fiber.current () with
+    | exception _ ->
+      (* if run outside a fiber, do nothing *)
+      f ()
+    | fiber ->
+      let old =
+        try FLS.get_exn fiber fls_context_key with FLS.Not_set -> Hmap.empty
+      in
+      FLS.set fiber fls_context_key ctx;
+      (match f () with
+      | res ->
+        FLS.set fiber fls_context_key old;
+        res
+      | exception exn ->
+        let bt = Printexc.get_raw_backtrace () in
+        FLS.set fiber fls_context_key old;
+        Printexc.raise_with_backtrace exn bt)
+end
+
+let storage : Storage.t = { name = "picos_fls"; get_context; with_context }