add picos backend to ambient-context

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

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"}
 ]

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

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 }