Source file thread_local_storage.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
(* original idea:
   https://discuss.ocaml.org/t/a-hack-to-implement-efficient-tls-thread-local-storage/13264 *)

(* sanity check *)
let () = assert (Obj.field (Obj.repr (Thread.self ())) 1 = Obj.repr ())

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 = Atomic.make 0

(** 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

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[@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

exception Not_set

let[@inline] get_exn slot =
  let v = get_raw slot in
  if v != sentinel_value_for_uninit_tls then
    Obj.obj v
  else
    raise_notrace Not_set

let[@inline] get_opt slot =
  let v = get_raw slot in
  if v != sentinel_value_for_uninit_tls then
    Some (Obj.obj v)
  else
    None

(** Allocating and setting *)

(** Grow the array so that [index] is valid. *)
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 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 [||] index in
    thread.tls <- Obj.repr new_tls;
    new_tls
  ) else (
    let tls = (Obj.obj tls : Obj.t array) in
    if index < Array.length tls then
      tls
    else (
      let new_tls = grow tls index in
      thread.tls <- Obj.repr new_tls;
      new_tls
    )
  )

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))

let[@inline] get_default ~default slot =
  let v = get_raw slot in
  if v != sentinel_value_for_uninit_tls then
    Obj.obj v
  else (
    let v = default () in
    set slot v;
    v
  )