Skip to content

Commit b4e984b

Browse files
GabrielBuicaGabrielBuica
committed
CP-52745: Add ThreadLocalStorage in Threadext
Signed-off-by: Gabriel Buica <[email protected]>
1 parent 03c1780 commit b4e984b

File tree

4 files changed

+221
-25
lines changed

4 files changed

+221
-25
lines changed

‎ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/dune

+38-24
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,43 @@
11
(library
2-
(public_name xapi-stdext-threads)
3-
(name xapi_stdext_threads)
4-
(modules :standard \ ipq scheduler threadext_test ipq_test scheduler_test)
5-
(libraries
6-
mtime
7-
mtime.clock.os
8-
threads.posix
9-
unix
10-
xapi-stdext-unix
11-
xapi-stdext-pervasives)
12-
(foreign_stubs
13-
(language c)
14-
(names delay_stubs))
15-
)
2+
(public_name xapi-stdext-threads)
3+
(name xapi_stdext_threads)
4+
(modules :standard \ ipq scheduler threadext_test ipq_test scheduler_test)
5+
(libraries
6+
ambient-context.thread_local
7+
mtime
8+
mtime.clock.os
9+
threads.posix
10+
unix
11+
tgroup
12+
xapi-stdext-unix
13+
xapi-stdext-pervasives)
14+
(foreign_stubs
15+
(language c)
16+
(names delay_stubs)))
1617

1718
(library
18-
(public_name xapi-stdext-threads.scheduler)
19-
(name xapi_stdext_threads_scheduler)
20-
(modules ipq scheduler)
21-
(libraries mtime mtime.clock.os threads.posix unix xapi-log xapi-stdext-threads clock)
22-
)
19+
(public_name xapi-stdext-threads.scheduler)
20+
(name xapi_stdext_threads_scheduler)
21+
(modules ipq scheduler)
22+
(libraries
23+
mtime
24+
mtime.clock.os
25+
threads.posix
26+
unix
27+
xapi-log
28+
xapi-stdext-threads
29+
clock))
2330

2431
(tests
25-
(names threadext_test ipq_test scheduler_test)
26-
(package xapi-stdext-threads)
27-
(modules threadext_test ipq_test scheduler_test)
28-
(libraries xapi_stdext_threads alcotest mtime.clock.os mtime fmt threads.posix xapi_stdext_threads_scheduler)
29-
)
32+
(names threadext_test ipq_test scheduler_test)
33+
(package xapi-stdext-threads)
34+
(modules threadext_test ipq_test scheduler_test)
35+
(libraries
36+
xapi_stdext_threads
37+
alcotest
38+
mtime.clock.os
39+
mtime
40+
fmt
41+
tgroup
42+
threads.posix
43+
xapi_stdext_threads_scheduler))

‎ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/threadext.ml

+64
Original file line numberDiff line numberDiff line change
@@ -97,3 +97,67 @@ let wait_timed_write fd timeout =
9797
true
9898
| _ ->
9999
assert false
100+
101+
module ThreadLocalStorage = struct
102+
type t = {
103+
ocaml_tid: int
104+
; thread_name: string
105+
; time_running: Mtime.span
106+
; time_last_yield: Mtime_clock.counter
107+
; tepoch: Mtime.t
108+
; tgroup: Tgroup.Group.t
109+
}
110+
111+
let thread_local_storage = Ambient_context_thread_local.Thread_local.create ()
112+
113+
let create () =
114+
let ocaml_tid = Thread.self () |> Thread.id in
115+
let thread_name = "" in
116+
let time_running = Mtime.Span.zero in
117+
let time_last_yield = Mtime_clock.counter () in
118+
let tepoch = Mtime_clock.now () in
119+
let tgroup =
120+
Tgroup.Group.(
121+
of_creator (Creator.make ~identity:Identity.root_identity ())
122+
)
123+
in
124+
let tls =
125+
{thread_name; tgroup; ocaml_tid; time_running; time_last_yield; tepoch}
126+
in
127+
let () =
128+
Ambient_context_thread_local.Thread_local.set thread_local_storage tls
129+
in
130+
tls
131+
132+
let get () =
133+
Ambient_context_thread_local.Thread_local.get_or_create ~create
134+
thread_local_storage
135+
136+
let set ?thread_name ?time_running ?time_last_yield ?tepoch ?tgroup () =
137+
let tls = get () in
138+
let tls =
139+
Option.fold ~none:tls
140+
~some:(fun thread_name -> {tls with thread_name})
141+
thread_name
142+
in
143+
let tls =
144+
Option.fold ~none:tls
145+
~some:(fun time_running -> {tls with time_running})
146+
time_running
147+
in
148+
let tls =
149+
Option.fold ~none:tls
150+
~some:(fun time_last_yield -> {tls with time_last_yield})
151+
time_last_yield
152+
in
153+
let tls =
154+
Option.fold ~none:tls ~some:(fun tepoch -> {tls with tepoch}) tepoch
155+
in
156+
let tls =
157+
Option.fold ~none:tls ~some:(fun tgroup -> {tls with tgroup}) tgroup
158+
in
159+
Ambient_context_thread_local.Thread_local.set thread_local_storage tls
160+
161+
let remove () =
162+
Ambient_context_thread_local.Thread_local.remove thread_local_storage
163+
end

‎ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/threadext.mli

+32
Original file line numberDiff line numberDiff line change
@@ -43,3 +43,35 @@ end
4343
val wait_timed_read : Unix.file_descr -> float -> bool
4444

4545
val wait_timed_write : Unix.file_descr -> float -> bool
46+
47+
module ThreadLocalStorage : sig
48+
type t = {
49+
ocaml_tid: int
50+
; thread_name: string
51+
; time_running: Mtime.span
52+
; time_last_yield: Mtime_clock.counter
53+
; tepoch: Mtime.t
54+
; tgroup: Tgroup.Group.t
55+
}
56+
57+
val create : unit -> t
58+
(** [create ()] creates and returns an initial thread local strorage for the
59+
current thread. *)
60+
61+
val get : unit -> t
62+
(** [get ()] returns the current thread local storage. *)
63+
64+
val set :
65+
?thread_name:string
66+
-> ?time_running:Mtime.span
67+
-> ?time_last_yield:Mtime_clock.counter
68+
-> ?tepoch:Mtime.t
69+
-> ?tgroup:Tgroup.Group.t
70+
-> unit
71+
-> unit
72+
(** [set ?thread_name ?time_running ?time_last_yield ?tepoch ?tgroup ()]
73+
updates the thread local storage based on the supplied arguments. *)
74+
75+
val remove : unit -> unit
76+
(** [remove ()] removes the thread local storage of the current thread. *)
77+
end

‎ocaml/libs/xapi-stdext/lib/xapi-stdext-threads/threadext_test.ml

+87-1
Original file line numberDiff line numberDiff line change
@@ -74,4 +74,90 @@ let tests =
7474
; ("other_thread", `Quick, other_thread)
7575
]
7676

77-
let () = Alcotest.run "Threadext" [("Delay", tests)]
77+
let test_create_ambient_storage () =
78+
let open Xapi_stdext_threads.Threadext in
79+
let _ : Thread.t =
80+
Thread.create
81+
(fun () ->
82+
let storage = ThreadLocalStorage.create () in
83+
let storage_tid = storage.ocaml_tid in
84+
let ocaml_tid = Thread.self () |> Thread.id in
85+
Alcotest.(check int)
86+
"Ocaml thread id matches the thread id stored" ocaml_tid storage_tid
87+
)
88+
()
89+
in
90+
()
91+
92+
let test_thread_storage_set_and_get () =
93+
let open Xapi_stdext_threads.Threadext in
94+
let _ : Thread.t =
95+
Thread.create
96+
(fun () ->
97+
let _ : ThreadLocalStorage.t = ThreadLocalStorage.create () in
98+
99+
let expected_name = "thread_1" in
100+
ThreadLocalStorage.set ~thread_name:expected_name () ;
101+
let storage = ThreadLocalStorage.get () in
102+
Alcotest.(check string)
103+
"Check if correct value is set in storage" expected_name
104+
storage.thread_name
105+
)
106+
()
107+
in
108+
()
109+
110+
let test_storage_locality () =
111+
let open Xapi_stdext_threads.Threadext in
112+
let r1 = ref None in
113+
let r2 = ref None in
114+
115+
let thread1_expected_name = "thread_1" in
116+
let thread2_expected_name = "thread_2" in
117+
118+
let thread1 =
119+
Thread.create
120+
(fun () ->
121+
let _ : ThreadLocalStorage.t = ThreadLocalStorage.create () in
122+
ThreadLocalStorage.set ~thread_name:thread1_expected_name () ;
123+
Thread.delay 1. ;
124+
r1 := Some (ThreadLocalStorage.get ())
125+
)
126+
()
127+
in
128+
let thread2 =
129+
Thread.create
130+
(fun () ->
131+
let _ : ThreadLocalStorage.t = ThreadLocalStorage.create () in
132+
ThreadLocalStorage.set ~thread_name:thread2_expected_name () ;
133+
134+
r2 := Some (ThreadLocalStorage.get ())
135+
)
136+
()
137+
in
138+
Thread.join thread1 ;
139+
Thread.join thread2 ;
140+
Alcotest.(check bool)
141+
"Check thread local storage is set for thread1" true (Option.is_some !r1) ;
142+
Alcotest.(check bool)
143+
"Check thread local storage is set for thread2" true (Option.is_some !r2) ;
144+
let thread1_name =
145+
let r1 = Option.get !r1 in
146+
r1.thread_name
147+
in
148+
let thread2_name =
149+
let r2 = Option.get !r2 in
150+
r2.thread_name
151+
in
152+
Alcotest.(check string) "Thread1 name" thread1_expected_name thread1_name ;
153+
Alcotest.(check string) "Thread2 name" thread2_expected_name thread2_name
154+
155+
let tls_tests =
156+
[
157+
("create storage", `Quick, test_create_ambient_storage)
158+
; ("storage set and get", `Quick, test_thread_storage_set_and_get)
159+
; ("thread local storage", `Quick, test_storage_locality)
160+
]
161+
162+
let () =
163+
Alcotest.run "Threadext" [("Delay", tests); ("ThreadLocalStorage", tls_tests)]

0 commit comments

Comments
 (0)