-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
move to thread-local-storage 0.2 with get/set API
- Loading branch information
Showing
10 changed files
with
112 additions
and
89 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,21 +1,13 @@ | ||
| (** Thread local storage *) | ||
|
|
||
| (* TODO: alias this to the library if present *) | ||
| type 'a t | ||
| (** A TLS slot for values of type ['a]. This allows the storage of a | ||
| single value of type ['a] per thread. *) | ||
|
|
||
| type 'a key | ||
| (** A TLS key for values of type ['a]. This allows the | ||
| storage of a single value of type ['a] per thread. *) | ||
| val create : unit -> 'a t | ||
|
|
||
| val new_key : (unit -> 'a) -> 'a key | ||
| (** Allocate a new, generative key. | ||
| When the key is used for the first time on a thread, | ||
| the function is called to produce it. | ||
| val get : 'a t -> 'a | ||
| (** @raise Failure if not present *) | ||
|
|
||
| This should only ever be called at toplevel to produce | ||
| constants, do not use it in a loop. *) | ||
|
|
||
| val get : 'a key -> 'a | ||
| (** Get the value for the current thread. *) | ||
|
|
||
| val set : 'a key -> 'a -> unit | ||
| (** Set the value for the current thread. *) | ||
| val get_opt : 'a t -> 'a option | ||
| val set : 'a t -> 'a -> unit |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,82 +1,111 @@ | ||
| (* see: https://discuss.ocaml.org/t/a-hack-to-implement-efficient-tls-thread-local-storage/13264 *) | ||
|
|
||
| module A = Atomic_ | ||
| (* vendored from https://github.com/c-cube/thread-local-storage *) | ||
|
|
||
| (* sanity check *) | ||
| let () = assert (Obj.field (Obj.repr (Thread.self ())) 1 = Obj.repr ()) | ||
|
|
||
| type 'a key = { | ||
| index: int; (** Unique index for this key. *) | ||
| compute: unit -> 'a; | ||
| (** Initializer for values for this key. Called at most | ||
| once per thread. *) | ||
| } | ||
| type 'a t = int | ||
| (** Unique index for this TLS slot. *) | ||
|
|
||
| let tls_length index = | ||
| let ceil_pow_2_minus_1 (n : int) : int = | ||
| let n = n lor (n lsr 1) in | ||
| let n = n lor (n lsr 2) in | ||
| let n = n lor (n lsr 4) in | ||
| let n = n lor (n lsr 8) in | ||
| let n = n lor (n lsr 16) in | ||
| if Sys.int_size > 32 then | ||
| n lor (n lsr 32) | ||
| else | ||
| n | ||
| in | ||
| let size = ceil_pow_2_minus_1 (index + 1) in | ||
| assert (size > index); | ||
| size | ||
|
|
||
| (** Counter used to allocate new keys *) | ||
| let counter = A.make 0 | ||
| let counter = Atomic.make 0 | ||
|
|
||
| (** Value used to detect a TLS slot that was not initialized yet *) | ||
| let[@inline] sentinel_value_for_uninit_tls_ () : Obj.t = Obj.repr counter | ||
| (** Value used to detect a TLS slot that was not initialized yet. | ||
| Because [counter] is private and lives forever, no other | ||
| object the user can see will have the same address. *) | ||
| let sentinel_value_for_uninit_tls : Obj.t = Obj.repr counter | ||
|
|
||
| let new_key compute : _ key = | ||
| let index = A.fetch_and_add counter 1 in | ||
| { index; compute } | ||
| external max_wosize : unit -> int = "caml_sys_const_max_wosize" | ||
|
|
||
| let max_word_size = max_wosize () | ||
|
|
||
| let create () : _ t = | ||
| let index = Atomic.fetch_and_add counter 1 in | ||
| if tls_length index <= max_word_size then | ||
| index | ||
| else ( | ||
| (* Some platforms have a small max word size. *) | ||
| ignore (Atomic.fetch_and_add counter (-1)); | ||
| failwith "Thread_local_storage.create: out of TLS slots" | ||
| ) | ||
|
|
||
| type thread_internal_state = { | ||
| _id: int; (** Thread ID (here for padding reasons) *) | ||
| mutable tls: Obj.t; (** Our data, stowed away in this unused field *) | ||
| _other: Obj.t; | ||
| (** Here to avoid lying to ocamlopt/flambda about the size of [Thread.t] *) | ||
| } | ||
| (** A partial representation of the internal type [Thread.t], allowing | ||
| us to access the second field (unused after the thread | ||
| has started) and stash TLS data in it. *) | ||
|
|
||
| let ceil_pow_2_minus_1 (n : int) : int = | ||
| let n = n lor (n lsr 1) in | ||
| let n = n lor (n lsr 2) in | ||
| let n = n lor (n lsr 4) in | ||
| let n = n lor (n lsr 8) in | ||
| let n = n lor (n lsr 16) in | ||
| if Sys.int_size > 32 then | ||
| n lor (n lsr 32) | ||
| let[@inline] get_raw index : Obj.t = | ||
| let thread : thread_internal_state = Obj.magic (Thread.self ()) in | ||
| let tls = thread.tls in | ||
| if Obj.is_block tls && index < Array.length (Obj.obj tls : Obj.t array) then | ||
| Array.unsafe_get (Obj.obj tls : Obj.t array) index | ||
| else | ||
| sentinel_value_for_uninit_tls | ||
|
|
||
| let[@inline never] tls_error () = | ||
| failwith "Thread_local_storage.get: TLS entry not initialised" | ||
|
|
||
| let[@inline] get slot = | ||
| let v = get_raw slot in | ||
| if v != sentinel_value_for_uninit_tls then | ||
| Obj.obj v | ||
| else | ||
| tls_error () | ||
|
|
||
| let[@inline] get_opt slot = | ||
| let v = get_raw slot in | ||
| if v != sentinel_value_for_uninit_tls then | ||
| Some (Obj.obj v) | ||
| else | ||
| n | ||
| None | ||
|
|
||
| (** Allocating and setting *) | ||
|
|
||
| (** Grow the array so that [index] is valid. *) | ||
| let[@inline never] grow_tls (old : Obj.t array) (index : int) : Obj.t array = | ||
| let new_length = ceil_pow_2_minus_1 (index + 1) in | ||
| let new_ = Array.make new_length (sentinel_value_for_uninit_tls_ ()) in | ||
| let grow (old : Obj.t array) (index : int) : Obj.t array = | ||
| let new_length = tls_length index in | ||
| let new_ = Array.make new_length sentinel_value_for_uninit_tls in | ||
| Array.blit old 0 new_ 0 (Array.length old); | ||
| new_ | ||
|
|
||
| let[@inline] get_tls_ (index : int) : Obj.t array = | ||
| let get_tls_with_capacity index : Obj.t array = | ||
| let thread : thread_internal_state = Obj.magic (Thread.self ()) in | ||
| let tls = thread.tls in | ||
| if Obj.is_int tls then ( | ||
| let new_tls = grow_tls [||] index in | ||
| thread.tls <- Obj.magic new_tls; | ||
| let new_tls = grow [||] index in | ||
| thread.tls <- Obj.repr new_tls; | ||
| new_tls | ||
| ) else ( | ||
| let tls = (Obj.magic tls : Obj.t array) in | ||
| let tls = (Obj.obj tls : Obj.t array) in | ||
| if index < Array.length tls then | ||
| tls | ||
| else ( | ||
| let new_tls = grow_tls tls index in | ||
| thread.tls <- Obj.magic new_tls; | ||
| let new_tls = grow tls index in | ||
| thread.tls <- Obj.repr new_tls; | ||
| new_tls | ||
| ) | ||
| ) | ||
|
|
||
| let get key = | ||
| let tls = get_tls_ key.index in | ||
| let value = Array.unsafe_get tls key.index in | ||
| if value != sentinel_value_for_uninit_tls_ () then | ||
| Obj.magic value | ||
| else ( | ||
| let value = key.compute () in | ||
| Array.unsafe_set tls key.index (Obj.repr (Sys.opaque_identity value)); | ||
| value | ||
| ) | ||
|
|
||
| let set key value = | ||
| let tls = get_tls_ key.index in | ||
| Array.unsafe_set tls key.index (Obj.repr (Sys.opaque_identity value)) | ||
| let[@inline] set slot value : unit = | ||
| let tls = get_tls_with_capacity slot in | ||
| Array.unsafe_set tls slot (Obj.repr (Sys.opaque_identity value)) |