Merge pull request #59 from imandra-ai/simon/inline-ambient-context

simon/inline ambient context

.ocamlformat

@@ -1,4 +1,4 @@
-version = 0.24.1
+version = 0.26.2
 profile=conventional
 margin=80
 if-then-else=k-r

dune-project

@@ -28,20 +28,27 @@
    (>= "4.08"))
   ptime
   hmap
-  ambient-context
+  atomic
+  (thread-local-storage
+   (and
+    (>= 0.2)
+    (< 0.3)))
   (odoc :with-doc)
   (alcotest :with-test)
   (pbrt
-   (and (>= 3.0) (< 4.0)))
+   (and
+    (>= 3.0)
+    (< 4.0)))
   (ocaml-lsp-server :with-dev-setup)
   (ocamlformat
    (and
     :with-dev-setup
     (>= 0.24)
     (< 0.25))))
- (depopts trace)
+ (depopts trace lwt eio)
  (conflicts
-  (trace (< 0.7)))
+  (trace
+   (< 0.7)))
  (tags
   (instrumentation tracing opentelemetry datadog jaeger)))
 

opentelemetry.opam

@@ -17,14 +17,15 @@ depends: [
   "ocaml" {>= "4.08"}
   "ptime"
   "hmap"
-  "ambient-context"
+  "atomic"
+  "thread-local-storage" {>= "0.2" & < "0.3"}
   "odoc" {with-doc}
   "alcotest" {with-test}
   "pbrt" {>= "3.0" & < "4.0"}
   "ocaml-lsp-server" {with-dev-setup}
   "ocamlformat" {with-dev-setup & >= "0.24" & < "0.25"}
 ]
-depopts: ["trace"]
+depopts: ["trace" "lwt" "eio"]
 conflicts: [
   "trace" {< "0.7"}
 ]

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/eio/dune

@@ -0,0 +1,7 @@
+(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 eio hmap opentelemetry.ambient-context thread-local-storage))

src/ambient-context/eio/opentelemetry_ambient_context_eio.ml

@@ -0,0 +1,40 @@
+module TLS = Thread_local_storage
+module Fiber = Eio.Fiber
+
+open struct
+  let _internal_key : Hmap.t Fiber.key = Fiber.create_key ()
+
+  let ( let* ) = Option.bind
+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)

src/ambient-context/eio/opentelemetry_ambient_context_eio.mli

@@ -0,0 +1,2 @@
+val storage : unit -> Opentelemetry_ambient_context.storage
+(** Storage using Eio's fibers local storage *)

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/lwt/dune

@@ -0,0 +1,7 @@
+(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 thread-local-storage))

src/ambient-context/lwt/opentelemetry_ambient_context_lwt.ml

@@ -0,0 +1,37 @@
+open struct
+  let _internal_key : Hmap.t Lwt.key = Lwt.new_key ()
+
+  let ( let* ) = Option.bind
+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)

src/ambient-context/lwt/opentelemetry_ambient_context_lwt.mli

@@ -0,0 +1,2 @@
+val storage : unit -> Opentelemetry_ambient_context.storage
+(** Storage using Lwt keys *)

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)