Skip to content

Commit f88a519

Browse files
polytypicpolytypic
committed
Introduce Per_thread cache
1 parent bce1098 commit f88a519

File tree

8 files changed

+71
-43
lines changed

8 files changed

+71
-43
lines changed

‎lib/picos/bootstrap/dune

+1
Original file line numberDiff line numberDiff line change
@@ -3,4 +3,5 @@
33
(package picos)
44
(libraries
55
(re_export picos.exn_bt)
6+
picos.thread
67
backoff))

‎lib/picos/bootstrap/picos_bootstrap.ml

+11
Original file line numberDiff line numberDiff line change
@@ -230,6 +230,10 @@ module Fiber = struct
230230

231231
type t = [ `Fiber ] tdt
232232

233+
module Maybe = struct
234+
type t = T : [< `Nothing | `Fiber ] tdt -> t [@@unboxed]
235+
end
236+
233237
let create_packed ~forbid packed = Fiber { forbid; packed; fls = [||] }
234238

235239
let create ~forbid computation =
@@ -392,3 +396,10 @@ module Handler = struct
392396
await : 'c -> Trigger.t -> Exn_bt.t option;
393397
}
394398
end
399+
400+
module Per_thread = struct
401+
type t = { mutable current : Fiber.Maybe.t (** *) }
402+
403+
let key = Picos_thread.TLS.new_key @@ fun () -> { current = T Nothing }
404+
let get () = Picos_thread.TLS.get key
405+
end

‎lib/picos/ocaml5/picos_ocaml.ml

+5-1
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,11 @@ module Fiber = struct
2222

2323
type _ Effect.t += Current : Fiber.t Effect.t
2424

25-
let current () = Effect.perform Current
25+
let current () =
26+
let p = Per_thread.get () in
27+
match p.current with
28+
| T Nothing -> Effect.perform Current
29+
| T (Fiber _ as fiber) -> fiber
2630

2731
type _ Effect.t +=
2832
| Spawn : { fiber : Fiber.t; main : Fiber.t -> unit } -> unit Effect.t

‎lib/picos/picos.common.ml

+6-2
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@ module Fiber = struct
2929
include Picos_ocaml.Fiber
3030

3131
module Maybe = struct
32-
let[@inline never] not_a_fiber () = invalid_arg "not a fiber"
32+
include Maybe
3333

34-
type t = T : [< `Nothing | `Fiber ] tdt -> t [@@unboxed]
34+
let[@inline never] not_a_fiber () = invalid_arg "not a fiber"
3535

3636
let[@inline] to_fiber_or_current = function
3737
| T Nothing -> current ()
@@ -85,3 +85,7 @@ module Handler = struct
8585
include Picos_bootstrap.Handler
8686
include Picos_ocaml.Handler
8787
end
88+
89+
module Per_thread = struct
90+
include Picos_bootstrap.Per_thread
91+
end

‎lib/picos/picos.mli

+10
Original file line numberDiff line numberDiff line change
@@ -1152,3 +1152,13 @@ module Handler : sig
11521152
implements an effect handler directly, because that is likely to perform
11531153
better. *)
11541154
end
1155+
1156+
module Per_thread : sig
1157+
(** *)
1158+
1159+
type t = { mutable current : Fiber.Maybe.t (** *) }
1160+
(** *)
1161+
1162+
val get : unit -> t
1163+
(** *)
1164+
end

‎lib/picos_fifos/picos_fifos.ml

+14-14
Original file line numberDiff line numberDiff line change
@@ -26,31 +26,31 @@ type t = {
2626
return : ((unit, unit) Effect.Deep.continuation -> unit) option;
2727
discontinue : ((unit, unit) Effect.Deep.continuation -> unit) option;
2828
handler : (unit, unit) Effect.Deep.handler;
29+
per_thread : Per_thread.t;
2930
quota : int;
30-
mutable fiber : Fiber.Maybe.t;
3131
mutable remaining_quota : int;
3232
}
3333

3434
let rec next t =
3535
match Picos_mpscq.pop_exn t.ready with
3636
| Spawn (fiber, main) ->
37-
t.fiber <- Fiber.Maybe.of_fiber fiber;
37+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
3838
t.remaining_quota <- t.quota;
3939
Effect.Deep.match_with main fiber t.handler
4040
| Return (fiber, k) ->
41-
t.fiber <- fiber;
41+
t.per_thread.current <- fiber;
4242
t.remaining_quota <- t.quota;
4343
Effect.Deep.continue k ()
4444
| Continue (fiber, k) ->
45-
t.fiber <- Fiber.Maybe.of_fiber fiber;
45+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
4646
t.remaining_quota <- t.quota;
4747
Fiber.continue fiber k ()
4848
| Resume (fiber, k) ->
49-
t.fiber <- Fiber.Maybe.of_fiber fiber;
49+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
5050
t.remaining_quota <- t.quota;
5151
Fiber.resume fiber k
5252
| exception Picos_mpscq.Empty ->
53-
t.fiber <- Fiber.Maybe.nothing;
53+
t.per_thread.current <- Fiber.Maybe.nothing;
5454
if Atomic.get t.num_alive_fibers <> 0 then begin
5555
if Atomic.get t.needs_wakeup then begin
5656
Mutex.lock t.mutex;
@@ -87,7 +87,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
8787
let rec t =
8888
{
8989
ready;
90-
fiber = Fiber.Maybe.nothing;
90+
per_thread = Per_thread.get ();
9191
needs_wakeup;
9292
num_alive_fibers;
9393
mutex;
@@ -107,12 +107,12 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
107107
later. *)
108108
Some
109109
(fun k ->
110-
let fiber = Fiber.Maybe.to_fiber t.fiber in
110+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
111111
Effect.Deep.continue k fiber)
112112
and yield =
113113
Some
114114
(fun k ->
115-
let fiber = Fiber.Maybe.to_fiber t.fiber in
115+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
116116
Picos_mpscq.push t.ready (Continue (fiber, k));
117117
next t)
118118
and return =
@@ -124,13 +124,13 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
124124
Effect.Deep.continue k ()
125125
end
126126
else begin
127-
Picos_mpscq.push t.ready (Return (t.fiber, k));
127+
Picos_mpscq.push t.ready (Return (t.per_thread.current, k));
128128
next t
129129
end)
130130
and discontinue =
131131
Some
132132
(fun k ->
133-
let fiber = Fiber.Maybe.to_fiber t.fiber in
133+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
134134
Fiber.continue fiber k ())
135135
and handler = { retc; exnc; effc }
136136
and[@alert "-handler"] effc :
@@ -143,7 +143,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
143143
| Fiber.Spawn r ->
144144
(* We check cancelation status once and then either perform the
145145
whole operation or discontinue the fiber. *)
146-
let fiber = Fiber.Maybe.to_fiber t.fiber in
146+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
147147
if Fiber.is_canceled fiber then t.discontinue
148148
else begin
149149
Atomic.incr t.num_alive_fibers;
@@ -154,7 +154,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
154154
| Computation.Cancel_after r -> begin
155155
(* We check cancelation status once and then either perform the
156156
whole operation or discontinue the fiber. *)
157-
let fiber = Fiber.Maybe.to_fiber t.fiber in
157+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
158158
if Fiber.is_canceled fiber then t.discontinue
159159
else
160160
match
@@ -168,7 +168,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
168168
| Trigger.Await trigger ->
169169
Some
170170
(fun k ->
171-
let fiber = Fiber.Maybe.to_fiber t.fiber in
171+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
172172
if Fiber.try_suspend fiber trigger fiber k t.resume then next t
173173
else
174174
let remaining_quota = t.remaining_quota - 1 in

‎lib/picos_randos/picos_randos.ml

+23-25
Original file line numberDiff line numberDiff line change
@@ -53,30 +53,28 @@ type t = {
5353
mutable run : bool;
5454
}
5555

56-
let fiber_key = Picos_thread.TLS.new_key @@ fun () -> ref Fiber.Maybe.nothing
57-
58-
let rec next p t =
56+
let rec next (p : Per_thread.t) t =
5957
match Collection.pop_exn t.ready with
6058
| Spawn (fiber, main) ->
61-
p := Fiber.Maybe.of_fiber fiber;
59+
p.current <- Fiber.Maybe.of_fiber fiber;
6260
Effect.Deep.match_with main fiber t.handler
6361
| Raise (fiber, k, exn_bt) ->
64-
p := Fiber.Maybe.of_fiber fiber;
62+
p.current <- Fiber.Maybe.of_fiber fiber;
6563
Exn_bt.discontinue k exn_bt
6664
| Return (fiber, k) ->
67-
p := Fiber.Maybe.of_fiber fiber;
65+
p.current <- Fiber.Maybe.of_fiber fiber;
6866
Effect.Deep.continue k ()
6967
| Current (fiber, k) ->
70-
p := Fiber.Maybe.of_fiber fiber;
68+
p.current <- Fiber.Maybe.of_fiber fiber;
7169
Effect.Deep.continue k fiber
7270
| Continue (fiber, k) ->
73-
p := Fiber.Maybe.of_fiber fiber;
71+
p.current <- Fiber.Maybe.of_fiber fiber;
7472
Fiber.continue fiber k ()
7573
| Resume (fiber, k) ->
76-
p := Fiber.Maybe.of_fiber fiber;
74+
p.current <- Fiber.Maybe.of_fiber fiber;
7775
Fiber.resume fiber k
7876
| exception Not_found ->
79-
p := Fiber.Maybe.nothing;
77+
p.current <- Fiber.Maybe.nothing;
8078
if Atomic.get t.num_alive_fibers <> 0 then begin
8179
Mutex.lock t.mutex;
8280
if Collection.is_empty t.ready && Atomic.get t.num_alive_fibers <> 0
@@ -158,22 +156,22 @@ let context ?fatal_exn_handler () =
158156
and current =
159157
Some
160158
(fun k ->
161-
let p = Picos_thread.TLS.get fiber_key in
162-
let fiber = Fiber.Maybe.to_fiber !p in
159+
let p = Per_thread.get () in
160+
let fiber = Fiber.Maybe.to_fiber p.current in
163161
Collection.push t.ready (Current (fiber, k));
164162
next p t)
165163
and yield =
166164
Some
167165
(fun k ->
168-
let p = Picos_thread.TLS.get fiber_key in
169-
let fiber = Fiber.Maybe.to_fiber !p in
166+
let p = Per_thread.get () in
167+
let fiber = Fiber.Maybe.to_fiber p.current in
170168
Collection.push t.ready (Continue (fiber, k));
171169
next p t)
172170
and return =
173171
Some
174172
(fun k ->
175-
let p = Picos_thread.TLS.get fiber_key in
176-
let fiber = Fiber.Maybe.to_fiber !p in
173+
let p = Per_thread.get () in
174+
let fiber = Fiber.Maybe.to_fiber p.current in
177175
Collection.push t.ready (Return (fiber, k));
178176
next p t)
179177
and handler = { retc; exnc; effc }
@@ -182,8 +180,8 @@ let context ?fatal_exn_handler () =
182180
function
183181
| Fiber.Current -> t.current
184182
| Fiber.Spawn r ->
185-
let p = Picos_thread.TLS.get fiber_key in
186-
let fiber = Fiber.Maybe.to_fiber !p in
183+
let p = Per_thread.get () in
184+
let fiber = Fiber.Maybe.to_fiber p.current in
187185
if Fiber.is_canceled fiber then t.yield
188186
else begin
189187
Atomic.incr t.num_alive_fibers;
@@ -193,8 +191,8 @@ let context ?fatal_exn_handler () =
193191
end
194192
| Fiber.Yield -> t.yield
195193
| Computation.Cancel_after r -> begin
196-
let p = Picos_thread.TLS.get fiber_key in
197-
let fiber = Fiber.Maybe.to_fiber !p in
194+
let p = Per_thread.get () in
195+
let fiber = Fiber.Maybe.to_fiber p.current in
198196
if Fiber.is_canceled fiber then t.yield
199197
else
200198
match
@@ -211,8 +209,8 @@ let context ?fatal_exn_handler () =
211209
| Trigger.Await trigger ->
212210
Some
213211
(fun k ->
214-
let p = Picos_thread.TLS.get fiber_key in
215-
let fiber = Fiber.Maybe.to_fiber !p in
212+
let p = Per_thread.get () in
213+
let fiber = Fiber.Maybe.to_fiber p.current in
216214
if Fiber.try_suspend fiber trigger fiber k t.resume then next p t
217215
else begin
218216
Collection.push t.ready (Resume (fiber, k));
@@ -221,14 +219,14 @@ let context ?fatal_exn_handler () =
221219
| _ -> None
222220
and retc () =
223221
Atomic.decr t.num_alive_fibers;
224-
let p = Picos_thread.TLS.get fiber_key in
222+
let p = Per_thread.get () in
225223
next p t
226224
in
227225
t
228226

229227
let runner_on_this_thread t =
230228
Select.check_configured ();
231-
next (Picos_thread.TLS.get fiber_key) t
229+
next (Per_thread.get ()) t
232230

233231
let rec await t =
234232
if !(t.num_waiters_non_zero) then begin
@@ -257,7 +255,7 @@ let run_fiber ?context:t_opt fiber main =
257255
t.run <- true;
258256
Mutex.unlock t.mutex;
259257
Collection.push t.ready (Spawn (fiber, main));
260-
next (Picos_thread.TLS.get fiber_key) t;
258+
next (Per_thread.get ()) t;
261259
Mutex.lock t.mutex;
262260
await t
263261
end

‎test/dune

+1-1
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@
6666
(modules test_picos_dscheck picos_bootstrap)
6767
(build_if
6868
(>= %{ocaml_version} 5))
69-
(libraries picos.exn_bt backoff traced_atomic dscheck alcotest))
69+
(libraries picos.exn_bt picos.thread backoff traced_atomic dscheck alcotest))
7070

7171
;;
7272

0 commit comments

Comments
 (0)