From 9584a7426f6ad8c16e1cba2ed3f0f4a30d2ab848 Mon Sep 17 00:00:00 2001 From: Simon Cruanes Date: Thu, 5 Sep 2024 16:11:09 -0400 Subject: [PATCH] wip: inline ambient-context into opentelemetry --- src/ambient-context/dune | 11 ++ src/ambient-context/hmap_key_.new.ml | 1 + src/ambient-context/hmap_key_.rcontext.ml | 1 + .../opentelemetry_ambient_context.ml | 124 ++++++++++++++++++ .../opentelemetry_ambient_context.mli | 54 ++++++++ src/ambient-context/types/dune | 4 + .../opentelemetry_ambient_context_types.ml | 19 +++ .../opentelemetry_ambient_context_types.mli | 30 +++++ 8 files changed, 244 insertions(+) create mode 100644 src/ambient-context/dune create mode 100644 src/ambient-context/hmap_key_.new.ml create mode 100644 src/ambient-context/hmap_key_.rcontext.ml create mode 100644 src/ambient-context/opentelemetry_ambient_context.ml create mode 100644 src/ambient-context/opentelemetry_ambient_context.mli create mode 100644 src/ambient-context/types/dune create mode 100644 src/ambient-context/types/opentelemetry_ambient_context_types.ml create mode 100644 src/ambient-context/types/opentelemetry_ambient_context_types.mli diff --git a/src/ambient-context/dune b/src/ambient-context/dune new file mode 100644 index 00000000..f0f87c7d --- /dev/null +++ b/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)))) diff --git a/src/ambient-context/hmap_key_.new.ml b/src/ambient-context/hmap_key_.new.ml new file mode 100644 index 00000000..1925b70e --- /dev/null +++ b/src/ambient-context/hmap_key_.new.ml @@ -0,0 +1 @@ +let key : Hmap.t Thread_local_storage.t = Thread_local_storage.create () diff --git a/src/ambient-context/hmap_key_.rcontext.ml b/src/ambient-context/hmap_key_.rcontext.ml new file mode 100644 index 00000000..99f136e6 --- /dev/null +++ b/src/ambient-context/hmap_key_.rcontext.ml @@ -0,0 +1 @@ +let key : Hmap.t Thread_local_storage.t = Rcontext.Ambient_hmap.k_hmap diff --git a/src/ambient-context/opentelemetry_ambient_context.ml b/src/ambient-context/opentelemetry_ambient_context.ml new file mode 100644 index 00000000..7c622eb7 --- /dev/null +++ b/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 + ) diff --git a/src/ambient-context/opentelemetry_ambient_context.mli b/src/ambient-context/opentelemetry_ambient_context.mli new file mode 100644 index 00000000..4cb1cc51 --- /dev/null +++ b/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). *) diff --git a/src/ambient-context/types/dune b/src/ambient-context/types/dune new file mode 100644 index 00000000..b9e4146c --- /dev/null +++ b/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)) diff --git a/src/ambient-context/types/opentelemetry_ambient_context_types.ml b/src/ambient-context/types/opentelemetry_ambient_context_types.ml new file mode 100644 index 00000000..829f7789 --- /dev/null +++ b/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) diff --git a/src/ambient-context/types/opentelemetry_ambient_context_types.mli b/src/ambient-context/types/opentelemetry_ambient_context_types.mli new file mode 100644 index 00000000..cded6589 --- /dev/null +++ b/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)