mirror of
https://github.com/c-cube/moonpool.git
synced 2025-12-12 22:10:46 -05:00
81 lines
2 KiB
OCaml
81 lines
2 KiB
OCaml
open Types_
|
|
module A = Atomic_
|
|
|
|
type 'a key = 'a ls_key
|
|
|
|
let key_count_ = A.make 0
|
|
|
|
type t = local_storage
|
|
type ls_value += Dummy
|
|
|
|
let dummy : t = ref [||]
|
|
|
|
(** Resize array of TLS values *)
|
|
let[@inline never] resize_ (cur : ls_value array ref) n =
|
|
if n > Sys.max_array_length then failwith "too many task local storage keys";
|
|
let len = Array.length !cur in
|
|
let new_ls =
|
|
Array.make (min Sys.max_array_length (max n ((len * 2) + 2))) Dummy
|
|
in
|
|
Array.blit !cur 0 new_ls 0 len;
|
|
cur := new_ls
|
|
|
|
module Direct = struct
|
|
type nonrec t = t
|
|
|
|
let create = create_local_storage
|
|
let[@inline] copy (self : t) = ref (Array.copy !self)
|
|
|
|
let get (type a) (self : t) ((module K) : a key) : a =
|
|
if K.offset >= Array.length !self then resize_ self (K.offset + 1);
|
|
match !self.(K.offset) with
|
|
| K.V x -> (* common case first *) x
|
|
| Dummy ->
|
|
(* first time we access this *)
|
|
let v = K.init () in
|
|
!self.(K.offset) <- K.V v;
|
|
v
|
|
| _ -> assert false
|
|
|
|
let set (type a) (self : t) ((module K) : a key) (v : a) : unit =
|
|
assert (self != dummy);
|
|
if K.offset >= Array.length !self then resize_ self (K.offset + 1);
|
|
!self.(K.offset) <- K.V v;
|
|
()
|
|
end
|
|
|
|
let new_key (type t) ~init () : t key =
|
|
let offset = A.fetch_and_add key_count_ 1 in
|
|
(module struct
|
|
type nonrec t = t
|
|
type ls_value += V of t
|
|
|
|
let offset = offset
|
|
let init = init
|
|
end : LS_KEY
|
|
with type t = t)
|
|
|
|
let[@inline] get_cur_ () : ls_value array ref =
|
|
match get_current_storage () with
|
|
| Some r -> r
|
|
| None -> failwith "Task local storage must be accessed from within a runner."
|
|
|
|
let[@inline] get (key : 'a key) : 'a =
|
|
let cur = get_cur_ () in
|
|
Direct.get cur key
|
|
|
|
let[@inline] get_opt key =
|
|
match get_current_storage () with
|
|
| None -> None
|
|
| Some cur -> Some (Direct.get cur key)
|
|
|
|
let[@inline] set key v : unit =
|
|
let cur = get_cur_ () in
|
|
Direct.set cur key v
|
|
|
|
let with_value key x f =
|
|
let old = get key in
|
|
set key x;
|
|
Fun.protect ~finally:(fun () -> set key old) f
|
|
|
|
let get_current = get_current_storage
|