mirror of
https://github.com/c-cube/moonpool.git
synced 2025-12-14 23:06:06 -05:00
68 lines
2.1 KiB
OCaml
68 lines
2.1 KiB
OCaml
open Types_
|
|
|
|
open struct
|
|
module FLS = Picos.Fiber.FLS
|
|
end
|
|
|
|
(** A local hmap, inherited in children fibers *)
|
|
let k_local_hmap : Hmap.t FLS.t = FLS.create ()
|
|
|
|
(** Access the local [hmap], or an empty one if not set *)
|
|
let[@inline] get_local_hmap () : Hmap.t =
|
|
match TLS.get_exn k_cur_fiber with
|
|
| exception TLS.Not_set -> Hmap.empty
|
|
| fiber -> FLS.get fiber ~default:Hmap.empty k_local_hmap
|
|
|
|
let[@inline] set_local_hmap (h : Hmap.t) : unit =
|
|
match TLS.get_exn k_cur_fiber with
|
|
| exception TLS.Not_set -> ()
|
|
| fiber -> FLS.set fiber k_local_hmap h
|
|
|
|
let[@inline] update_local_hmap (f : Hmap.t -> Hmap.t) : unit =
|
|
match TLS.get_exn k_cur_fiber with
|
|
| exception TLS.Not_set -> ()
|
|
| fiber ->
|
|
let h = FLS.get fiber ~default:Hmap.empty k_local_hmap in
|
|
let h = f h in
|
|
FLS.set fiber k_local_hmap h
|
|
|
|
(** @raise Invalid_argument if not present *)
|
|
let get_in_local_hmap_exn (k : 'a Hmap.key) : 'a =
|
|
let h = get_local_hmap () in
|
|
Hmap.get k h
|
|
|
|
let get_in_local_hmap_opt (k : 'a Hmap.key) : 'a option =
|
|
let h = get_local_hmap () in
|
|
Hmap.find k h
|
|
|
|
(** Remove given key from the local hmap *)
|
|
let[@inline] remove_in_local_hmap (k : _ Hmap.key) : unit =
|
|
update_local_hmap (Hmap.rem k)
|
|
|
|
let[@inline] set_in_local_hmap (k : 'a Hmap.key) (v : 'a) : unit =
|
|
update_local_hmap (Hmap.add k v)
|
|
|
|
(** [with_in_local_hmap k v f] calls [f()] in a context where [k] is bound to
|
|
[v] in the local hmap. Then it restores the previous binding for [k]. *)
|
|
let with_in_local_hmap (k : 'a Hmap.key) (v : 'a) f =
|
|
let h = get_local_hmap () in
|
|
match Hmap.find k h with
|
|
| None ->
|
|
set_in_local_hmap k v;
|
|
Fun.protect ~finally:(fun () -> remove_in_local_hmap k) f
|
|
| Some old_v ->
|
|
set_in_local_hmap k v;
|
|
Fun.protect ~finally:(fun () -> set_in_local_hmap k old_v) f
|
|
|
|
(**/**)
|
|
|
|
(* private functions, to be used by the rest of moonpool *)
|
|
module Private_hmap_ls_ = struct
|
|
(** Copy the hmap from f1.fls to f2.fls *)
|
|
let copy_fls (f1 : Picos.Fiber.t) (f2 : Picos.Fiber.t) : unit =
|
|
match FLS.get_exn f1 k_local_hmap with
|
|
| exception FLS.Not_set -> ()
|
|
| hmap -> FLS.set f2 k_local_hmap hmap
|
|
end
|
|
|
|
(**/**)
|