use a separate ambient-context library

dune-project

@@ -36,6 +36,8 @@
    (and
     (>= 4.0)
     (< 5.0)))
+  (ambient-context
+   (>= 0.2))
   (ocaml-lsp-server :with-dev-setup)
   (ocamlformat
    (and
@@ -47,7 +49,7 @@
  (depopts atomic trace thread-local-storage lwt eio picos)
  (conflicts
   (trace
-   (< 0.10)))
+   (< 0.99)))
  (tags
   (instrumentation tracing opentelemetry datadog jaeger)))
 
@@ -74,6 +76,7 @@
    (>= "4.08"))
   (opentelemetry
    (= :version))
+  ambient-context-lwt
   (cohttp-lwt-unix :with-test)
   (odoc :with-doc)
   (lwt
@@ -158,6 +161,7 @@
    (= :version))
   (opentelemetry-lwt
    (= :version))
+  ambient-context-lwt
   (odoc :with-doc)
   (lwt
    (>= "5.3"))
@@ -180,6 +184,7 @@
    (= :version))
   (opentelemetry-client
    (= :version))
+  ambient-context-lwt
   (odoc :with-doc)
   (lwt
    (>= "5.3"))
@@ -204,6 +209,7 @@
    (>= "1.4"))
   ca-certs
   mirage-crypto-rng-eio
+  ambient-context-eio
   (opentelemetry
    (= :version))
   (opentelemetry-client

opentelemetry-client-cohttp-eio.opam

@@ -17,6 +17,7 @@ depends: [
   "mtime" {>= "1.4"}
   "ca-certs"
   "mirage-crypto-rng-eio"
+  "ambient-context-eio"
   "opentelemetry" {= version}
   "opentelemetry-client" {= version}
   "odoc" {with-doc}

opentelemetry-client-cohttp-lwt.opam

@@ -17,6 +17,7 @@ depends: [
   "mtime" {>= "1.4"}
   "opentelemetry" {= version}
   "opentelemetry-client" {= version}
+  "ambient-context-lwt"
   "odoc" {with-doc}
   "lwt" {>= "5.3"}
   "lwt_ppx" {>= "2.0"}

opentelemetry-cohttp-lwt.opam

@@ -16,6 +16,7 @@ depends: [
   "ocaml" {>= "4.08"}
   "opentelemetry" {= version}
   "opentelemetry-lwt" {= version}
+  "ambient-context-lwt"
   "odoc" {with-doc}
   "lwt" {>= "5.3"}
   "cohttp" {>= "6.0.0"}

opentelemetry-lwt.opam

@@ -16,6 +16,7 @@ depends: [
   "dune" {>= "2.9"}
   "ocaml" {>= "4.08"}
   "opentelemetry" {= version}
+  "ambient-context-lwt"
   "cohttp-lwt-unix" {with-test}
   "odoc" {with-doc}
   "lwt" {>= "5.3"}

opentelemetry.opam

@@ -20,13 +20,14 @@ depends: [
   "odoc" {with-doc}
   "alcotest" {with-test}
   "pbrt" {>= "4.0" & < "5.0"}
+  "ambient-context" {>= "0.2"}
   "ocaml-lsp-server" {with-dev-setup}
   "ocamlformat" {with-dev-setup & >= "0.27" & < "0.28"}
   "mtime" {>= "1.4"}
 ]
 depopts: ["atomic" "trace" "thread-local-storage" "lwt" "eio" "picos"]
 conflicts: [
-  "trace" {< "0.10"}
+  "trace" {< "0.99"}
 ]
 build: [
   ["dune" "subst"] {dev}

src/ambient-context/basic_map.ml

@@ -1,49 +0,0 @@
-(** Extremely basic storage using a map from thread id to context *)
-
-open Opentelemetry_ambient_context_core
-
-open struct
-  module Atomic = Opentelemetry_atomic.Atomic
-
-  module Int_map = Map.Make (struct
-    type t = int
-
-    let compare : t -> t -> int = Stdlib.compare
-  end)
-
-  type st = { m: Context.t ref Int_map.t Atomic.t } [@@unboxed]
-
-  let get (self : st) : Context.t =
-    let tid = Thread.id @@ Thread.self () in
-    match Int_map.find tid (Atomic.get self.m) with
-    | exception Not_found -> Context.empty
-    | ctx_ref -> !ctx_ref
-
-  let with_context (self : st) ctx f =
-    let tid = Thread.id @@ Thread.self () in
-
-    let ctx_ref =
-      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
-        r, Int_map.add tid r m
-    in
-
-    let old_ctx = !ctx_ref in
-    ctx_ref := ctx;
-
-    let finally () = ctx_ref := old_ctx in
-    Fun.protect ~finally f
-end
-
-let create_storage () : Storage.t =
-  let st = { m = Atomic.make Int_map.empty } in
-  {
-    name = "basic-map";
-    get_context = (fun () -> get st);
-    with_context = (fun ctx f -> with_context st ctx f);
-  }
-
-(** Default storage *)
-let storage : Storage.t = create_storage ()

src/ambient-context/core/context.ml

@@ -1,7 +0,0 @@
-type t = Hmap.t
-
-type 'a key = 'a Hmap.key
-
-let empty : t = Hmap.empty
-
-let[@inline] new_key () : _ key = Hmap.Key.create ()

src/ambient-context/core/dune

@@ -1,6 +0,0 @@
-(library
- (name opentelemetry_ambient_context_core)
- (public_name opentelemetry.ambient-context.core)
- (synopsis "Core definitions for ambient-context")
- (libraries
-  (re_export hmap)))

src/ambient-context/core/storage.ml

@@ -1,44 +0,0 @@
-(** Storage implementation.
-
-    There is a singleton storage for a given program, responsible for providing
-    ambient context to the rest of the program. *)
-
-type t = {
-  name: string;
-  get_context: unit -> Context.t;
-  with_context: 'a. Context.t -> (unit -> 'a) -> 'a;
-}
-(** Storage type *)
-
-(** Name of the storage implementation. *)
-let[@inline] name self = self.name
-
-(** Get the context from the current storage, or [Hmap.empty] if there is no
-    ambient context. *)
-let[@inline] get_context self = self.get_context ()
-
-(** [with_context storage ctx f] calls [f()] in an ambient context in which
-    [get_context()] will return [ctx]. Once [f()] returns, the storage is reset
-    to its previous value. *)
-let[@inline] with_context self ctx f = self.with_context ctx f
-
-(** Get the ambient context and then look up [k] in it *)
-let[@inline] get self (k : 'a Context.key) : 'a option =
-  Hmap.find k (get_context self)
-
-(** [with_key_bound_to storage k v f] calls [f()] in a context updated to have
-    [k] map to [v]. *)
-let with_key_bound_to self k v f =
-  let ctx = get_context self in
-  let new_ctx = Hmap.add k v ctx in
-  self.with_context new_ctx f
-
-(** [with_key_unbound storage k f] calls [f()] in a context updated to have [k]
-    bound to no value. *)
-let with_key_unbound self k f =
-  let ctx = get_context self in
-  if Hmap.mem k ctx then (
-    let new_ctx = Hmap.rem k ctx in
-    self.with_context new_ctx f
-  ) else
-    f ()

src/ambient-context/default_.map.ml

@@ -1 +0,0 @@
-let storage = Basic_map.storage

src/ambient-context/default_.mli

@@ -1,2 +0,0 @@
-val storage : Storage.t
-(** Default storage. *)

src/ambient-context/default_.tls.ml

@@ -1 +0,0 @@
-let storage = Opentelemetry_ambient_context_tls.storage

src/ambient-context/dune

@@ -1,22 +1,7 @@
 (library
  (name opentelemetry_ambient_context)
  (public_name opentelemetry.ambient-context)
- (synopsis
-  "Abstraction over thread-local storage and fiber-local storage mechanisms")
- (flags
-  :standard
-  -open
-  Opentelemetry_ambient_context_core
-  -open
-  Opentelemetry_atomic)
+ (synopsis "re-export ambient-context")
  (libraries
-  hmap
-  threads
-  (re_export opentelemetry.ambient-context.core)
-  (re_export opentelemetry.atomic)
-  (re_export opentelemetry.util)
-  (select
-   default_.ml
-   from
-   (opentelemetry.ambient-context.tls -> default_.tls.ml)
-   (-> default_.map.ml))))
+  (re_export ambient-context.core)
+  (re_export ambient-context)))

src/ambient-context/eio/dune

@@ -1,11 +0,0 @@
-(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
-  (re_export eio)
-  (re_export eio.core)
-  hmap
-  (re_export opentelemetry.ambient-context.core)))

src/ambient-context/eio/opentelemetry_ambient_context_eio.ml

@@ -1,15 +0,0 @@
-open Opentelemetry_ambient_context_core
-module Fiber = Eio.Fiber
-
-open struct
-  let fiber_context_key : Context.t Fiber.key = Fiber.create_key ()
-end
-
-let storage : Storage.t =
-  {
-    name = "eio";
-    get_context =
-      (fun () ->
-        Fiber.get fiber_context_key |> Option.value ~default:Hmap.empty);
-    with_context = (fun ctx f -> Fiber.with_binding fiber_context_key ctx f);
-  }

src/ambient-context/lwt/dune

@@ -1,7 +0,0 @@
-(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.core))

src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml

@@ -1,15 +0,0 @@
-(** Storage using Lwt keys *)
-
-open Opentelemetry_ambient_context_core
-
-open struct
-  let lwt_context_key : Context.t Lwt.key = Lwt.new_key ()
-end
-
-let storage : Storage.t =
-  {
-    name = "lwt";
-    get_context =
-      (fun () -> Lwt.get lwt_context_key |> Option.value ~default:Hmap.empty);
-    with_context = (fun ctx f -> Lwt.with_value lwt_context_key (Some ctx) f);
-  }

src/ambient-context/opentelemetry_ambient_context.ml

@@ -1,49 +1,3 @@
-include Opentelemetry_ambient_context_core
+(** Just forward to the [ambient-context] library *)
 
-let default_storage = Default_.storage
-
-open struct
-  (** The current ambient-context storage. *)
-  let cur_storage : Storage.t Atomic.t = Atomic.make Default_.storage
-end
-
-let[@inline] get_current_storage () = Atomic.get cur_storage
-
-(* NOTE: we can't really "map" each local context from the old
-   to the new. Maybe the old storage is TLS based and the new one
-   is per-lwt-task. *)
-let set_current_storage (storage : Storage.t) = Atomic.set cur_storage storage
-
-(** {2 Functions operating with the current storage} *)
-
-(** Get the context from the current storage, or [Hmap.empty] if there is no
-    ambient context. *)
-let[@inline] get_context () = Storage.get_context (Atomic.get cur_storage)
-
-(** [with_context ctx f] calls [f()] in an ambient context in which
-    [get_context()] will return [ctx]. Once [f()] returns, the storage is reset
-    to its previous value. *)
-let[@inline] with_context ctx f =
-  Storage.with_context (Atomic.get cur_storage) ctx f
-
-(** Get the ambient context and then look up [k] in it *)
-let[@inline] get (k : 'a Context.key) : 'a option = Hmap.find k (get_context ())
-
-(** [with_key_bound_to storage k v f] calls [f()] in a context updated to have
-    [k] map to [v]. *)
-let with_key_bound_to k v f =
-  let storage = get_current_storage () in
-  let ctx = Storage.get_context storage in
-  let new_ctx = Hmap.add k v ctx in
-  Storage.with_context storage new_ctx f
-
-(** [with_key_unbound k f] calls [f()] in a context updated to have [k] bound to
-    no value. *)
-let with_key_unbound k f =
-  let storage = Atomic.get cur_storage in
-  let ctx = Storage.get_context storage in
-  if Hmap.mem k ctx then (
-    let new_ctx = Hmap.rem k ctx in
-    Storage.with_context storage new_ctx f
-  ) else
-    f ()
+include Ambient_context

src/ambient-context/picos/dune

@@ -1,6 +0,0 @@
-(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

@@ -1,34 +0,0 @@
-(** 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 }

src/ambient-context/tls/dune

@@ -1,8 +0,0 @@
-(library
- (name opentelemetry_ambient_context_tls)
- (public_name opentelemetry.ambient-context.tls)
- (synopsis "Implementation of ambient-context from thread-local-storage")
- (optional) ; TLS
- (libraries
-  (re_export opentelemetry.ambient-context.core)
-  thread-local-storage))

src/ambient-context/tls/opentelemetry_ambient_context_tls.ml

@@ -1,23 +0,0 @@
-open Opentelemetry_ambient_context_core
-
-open struct
-  module TLS = Thread_local_storage
-
-  (* key used to access the context *)
-  let tls_k_context : Context.t TLS.t = TLS.create ()
-end
-
-let storage : Storage.t =
-  {
-    name = "tls";
-    get_context =
-      (fun () -> try TLS.get_exn tls_k_context with TLS.Not_set -> Hmap.empty);
-    with_context =
-      (fun ctx f ->
-        let old =
-          try TLS.get_exn tls_k_context with TLS.Not_set -> Hmap.empty
-        in
-        let finally () = TLS.set tls_k_context old in
-        TLS.set tls_k_context ctx;
-        Fun.protect ~finally f);
-  }

src/client-cohttp-eio/dune

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

src/client-cohttp-eio/opentelemetry_client_cohttp_eio.ml

@@ -180,8 +180,7 @@ 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;
+  Opentelemetry_ambient_context.set_current_storage Ambient_context_eio.storage;
   let exp = create_exporter ?config ~sw ~env () in
   Main_exporter.set exp
 

src/client/lwt/dune

@@ -17,7 +17,7 @@
   (re_export opentelemetry.core)
   (re_export opentelemetry)
   (re_export opentelemetry.ambient-context)
-  opentelemetry.ambient-context.lwt
+  ambient-context-lwt
   (re_export opentelemetry-client)
   (re_export lwt)
   threads

src/client/lwt/util_ambient_context.ml

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

src/lwt/dune

@@ -4,5 +4,5 @@
  (synopsis "Lwt frontend for opentelemetry")
  (libraries
   (re_export lwt)
-  opentelemetry.ambient-context.lwt
+  ambient-context-lwt
   (re_export opentelemetry)))

src/lwt/opentelemetry_lwt.ml

@@ -2,8 +2,7 @@ include Opentelemetry
 
 (** Setup Lwt as the ambient context *)
 let setup_ambient_context () =
-  Opentelemetry_ambient_context.set_current_storage
-    Opentelemetry_ambient_context_lwt.storage
+  Opentelemetry_ambient_context.set_current_storage Ambient_context_lwt.storage
 
 module Main_exporter = struct
   include Main_exporter