wip: inline ambient-context into opentelemetry

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

src/ambient-context/hmap_key_.new.ml

@@ -0,0 +1 @@
+let key : Hmap.t Thread_local_storage.t = Thread_local_storage.create ()

src/ambient-context/hmap_key_.rcontext.ml

@@ -0,0 +1 @@
+let key : Hmap.t Thread_local_storage.t = Rcontext.Ambient_hmap.k_hmap

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

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

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

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)

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)