diff --git a/src/ambient-context/basic_map.ml b/src/ambient-context/basic_map.ml new file mode 100644 index 00000000..3575835a --- /dev/null +++ b/src/ambient-context/basic_map.ml @@ -0,0 +1,55 @@ +(** 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 = + try Int_map.find tid (Atomic.get self.m) + with Not_found -> + let r = ref Context.empty in + while + let m = Atomic.get self.m in + let m' = Int_map.add tid r m in + not (Atomic.compare_and_set self.m m m') + do + () + done; + r + 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 () diff --git a/src/ambient-context/core/context.ml b/src/ambient-context/core/context.ml new file mode 100644 index 00000000..658a83e6 --- /dev/null +++ b/src/ambient-context/core/context.ml @@ -0,0 +1,7 @@ +type t = Hmap.t + +type 'a key = 'a Hmap.key + +let empty : t = Hmap.empty + +let[@inline] new_key () : _ key = Hmap.Key.create () diff --git a/src/ambient-context/core/dune b/src/ambient-context/core/dune new file mode 100644 index 00000000..97961205 --- /dev/null +++ b/src/ambient-context/core/dune @@ -0,0 +1,5 @@ +(library + (name opentelemetry_ambient_context_core) + (public_name opentelemetry.ambient-context.core) + (synopsis "Core definitions for ambient-context") + (libraries hmap)) diff --git a/src/ambient-context/core/storage.ml b/src/ambient-context/core/storage.ml new file mode 100644 index 00000000..89bc3931 --- /dev/null +++ b/src/ambient-context/core/storage.ml @@ -0,0 +1,44 @@ +(** 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 () diff --git a/src/ambient-context/default_.map.ml b/src/ambient-context/default_.map.ml new file mode 100644 index 00000000..aecceb56 --- /dev/null +++ b/src/ambient-context/default_.map.ml @@ -0,0 +1 @@ +let storage = Basic_map.storage diff --git a/src/ambient-context/default_.mli b/src/ambient-context/default_.mli new file mode 100644 index 00000000..9f14c9bd --- /dev/null +++ b/src/ambient-context/default_.mli @@ -0,0 +1,2 @@ +val storage : Storage.t +(** Default storage. *) diff --git a/src/ambient-context/default_.tls.ml b/src/ambient-context/default_.tls.ml new file mode 100644 index 00000000..7411b452 --- /dev/null +++ b/src/ambient-context/default_.tls.ml @@ -0,0 +1 @@ +let storage = Opentelemetry_ambient_context_tls.storage diff --git a/src/ambient-context/dls.ml.tmp b/src/ambient-context/dls.ml.tmp new file mode 100644 index 00000000..dee2dc49 --- /dev/null +++ b/src/ambient-context/dls.ml.tmp @@ -0,0 +1,46 @@ +(* TODO: conditional compilation, and use Basic_map in each DLS *) + +(** Storage using DLS. *) + +open Opentelemetry_ambient_context_core + +open struct + module DLS = Domain.DLS + + module Int_map = Map.Make (struct + type t = int + + let compare : t -> t -> int = Stdlib.compare + end) + + (* key used to access the context *) + let dls_k_context : Context.t ref Int_map.t DLS.key = + DLS.new_key + ~split_from_parent:(fun _ -> Int_map.empty) + (fun _ -> Int_map.empty) + + let dls_get () = + let tid = Thread.id @@ Thread.self () in + let map_ref = DLS.get dls_k_context in + try !(Int_map.find tid map_ref) with Not_found -> Hmap.empty + + let dls_with ctx f = + let tid = Thread.id @@ Thread.self () in + let map = DLS.get dls_k_context in + let ctx_ref = + try Int_map.find tid map + with Not_found -> + let r = ref Context.empty in + DLS.set dls_k_context (Int_map.add tid r map); + r + in + + let old_ctx = !ctx_ref in + ctx_ref := ctx; + + let finally () = ctx_ref := old_ctx in + Fun.protect ~finally f +end + +let storage : Storage.t = + { name = "dls-int-map"; get_context = dls_get; with_context = dls_with } diff --git a/src/ambient-context/dune b/src/ambient-context/dune index 68e7acf5..de1f5b26 100644 --- a/src/ambient-context/dune +++ b/src/ambient-context/dune @@ -3,13 +3,19 @@ (public_name opentelemetry.ambient-context) (synopsis "Abstraction over thread-local storage and fiber-local storage mechanisms") - (private_modules hmap_key_) + (flags + :standard + -open + Opentelemetry_ambient_context_core + -open + Opentelemetry_atomic) (libraries - thread-local-storage - threads + hmap atomic - opentelemetry.ambient-context.types + opentelemetry.ambient-context.core + opentelemetry.atomic (select - hmap_key_.ml + default_.ml from - (-> hmap_key_.new.ml)))) + (opentelemetry.ambient-context.tls -> default_.tls.ml) + (-> default_.map.ml)))) diff --git a/src/ambient-context/eio/dune b/src/ambient-context/eio/dune index f3f76be7..fa56edd8 100644 --- a/src/ambient-context/eio/dune +++ b/src/ambient-context/eio/dune @@ -4,4 +4,4 @@ (synopsis "Storage backend for ambient-context using Eio's fibre-local storage") (optional) ; eio - (libraries eio hmap opentelemetry.ambient-context thread-local-storage)) + (libraries eio hmap opentelemetry.ambient-context.core)) diff --git a/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml b/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml index 56451a1a..1da61600 100644 --- a/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml +++ b/src/ambient-context/eio/opentelemetry_ambient_context_eio.ml @@ -1,39 +1,15 @@ +open Opentelemetry_ambient_context_core module Fiber = Eio.Fiber open struct - let _internal_key : Hmap.t Fiber.key = Fiber.create_key () - - let ( let* ) = Option.bind + let fiber_context_key : Context.t Fiber.key = Fiber.create_key () end -module M = struct - let name = "Storage_eio" - - let[@inline] get_map () = Fiber.get _internal_key - - let[@inline] with_map m cb = Fiber.with_binding _internal_key m 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 cb - - let without_binding k cb = - let new_context = - match get_map () with - | None -> Hmap.empty - | Some old_context -> Hmap.rem k old_context - in - with_map new_context cb -end - -let storage () : Opentelemetry_ambient_context.storage = (module M) +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); + } diff --git a/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli b/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli deleted file mode 100644 index ac5cf8ba..00000000 --- a/src/ambient-context/eio/opentelemetry_ambient_context_eio.mli +++ /dev/null @@ -1,2 +0,0 @@ -val storage : unit -> Opentelemetry_ambient_context.storage -(** Storage using Eio's fibers local storage *) diff --git a/src/ambient-context/hmap_key_.new.ml b/src/ambient-context/hmap_key_.new.ml deleted file mode 100644 index 1925b70e..00000000 --- a/src/ambient-context/hmap_key_.new.ml +++ /dev/null @@ -1 +0,0 @@ -let key : Hmap.t Thread_local_storage.t = Thread_local_storage.create () diff --git a/src/ambient-context/lwt/dune b/src/ambient-context/lwt/dune index 68a9de15..fb7398fb 100644 --- a/src/ambient-context/lwt/dune +++ b/src/ambient-context/lwt/dune @@ -4,4 +4,4 @@ (optional) ; lwt (synopsis "Storage backend for ambient-context using Lwt's sequence-associated storage") - (libraries lwt opentelemetry.ambient-context thread-local-storage)) + (libraries lwt opentelemetry.ambient-context.core)) diff --git a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml index b75105f2..d7187670 100644 --- a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml +++ b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml @@ -1,37 +1,15 @@ +(** Storage using Lwt keys *) + +open Opentelemetry_ambient_context_core + open struct - let _internal_key : Hmap.t Lwt.key = Lwt.new_key () - - let ( let* ) = Option.bind + let lwt_context_key : Context.t Lwt.key = Lwt.new_key () end -module M = struct - let name = "Storage_lwt" - - let[@inline] get_map () = Lwt.get _internal_key - - let[@inline] with_map m cb = Lwt.with_value _internal_key (Some m) 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 cb - - let without_binding k cb = - let new_context = - match get_map () with - | None -> Hmap.empty - | Some old_context -> Hmap.rem k old_context - in - with_map new_context cb -end - -let storage () : Opentelemetry_ambient_context.storage = (module M) +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); + } diff --git a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli b/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli deleted file mode 100644 index 3c462a8d..00000000 --- a/src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli +++ /dev/null @@ -1,2 +0,0 @@ -val storage : unit -> Opentelemetry_ambient_context.storage -(** Storage using Lwt keys *) diff --git a/src/ambient-context/opentelemetry_ambient_context.ml b/src/ambient-context/opentelemetry_ambient_context.ml index 7c622eb7..d8a18c57 100644 --- a/src/ambient-context/opentelemetry_ambient_context.ml +++ b/src/ambient-context/opentelemetry_ambient_context.ml @@ -1,124 +1,49 @@ -module TLS = Thread_local_storage -include Opentelemetry_ambient_context_types +include Opentelemetry_ambient_context_core -type 'a key = int * 'a Hmap.key +let default_storage = Default_.storage -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 () +open struct + (** The current ambient-context storage. *) + let cur_storage : Storage.t Atomic.t = Atomic.make Default_.storage end -let default_storage : storage = (module Storage_tls_hmap) +let[@inline] get_current_storage () = Atomic.get cur_storage -let k_current_storage : storage TLS.t = TLS.create () +(* 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 -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 +(** {2 Functions operating with the current storage} *) -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 () - ) +(** 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) -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 - ) +(** [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 -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 - ) +(** 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 ()) -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 - ) +(** [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 -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 - ) +(** [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 () diff --git a/src/ambient-context/opentelemetry_ambient_context.mli b/src/ambient-context/opentelemetry_ambient_context.mli deleted file mode 100644 index 8f19ff0c..00000000 --- a/src/ambient-context/opentelemetry_ambient_context.mli +++ /dev/null @@ -1,55 +0,0 @@ -(** 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/tls/dune b/src/ambient-context/tls/dune new file mode 100644 index 00000000..c0a285bc --- /dev/null +++ b/src/ambient-context/tls/dune @@ -0,0 +1,6 @@ +(library + (name opentelemetry_ambient_context_tls) + (public_name opentelemetry.ambient-context.tls) + (synopsis "Implementation of ambient-context from thread-local-storage") + (optional) ; TLS + (libraries opentelemetry.ambient-context.core thread-local-storage)) diff --git a/src/ambient-context/tls/opentelemetry_ambient_context_tls.ml b/src/ambient-context/tls/opentelemetry_ambient_context_tls.ml new file mode 100644 index 00000000..7d40387c --- /dev/null +++ b/src/ambient-context/tls/opentelemetry_ambient_context_tls.ml @@ -0,0 +1,23 @@ +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); + } diff --git a/src/ambient-context/types/dune b/src/ambient-context/types/dune deleted file mode 100644 index b9e4146c..00000000 --- a/src/ambient-context/types/dune +++ /dev/null @@ -1,4 +0,0 @@ -(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 deleted file mode 100644 index 829f7789..00000000 --- a/src/ambient-context/types/opentelemetry_ambient_context_types.ml +++ /dev/null @@ -1,19 +0,0 @@ -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 deleted file mode 100644 index 738b7520..00000000 --- a/src/ambient-context/types/opentelemetry_ambient_context_types.mli +++ /dev/null @@ -1,32 +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 '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)