Skip to content

Commit 22b5a1e

Browse files
polytypicpolytypic
committed
Introduce Per_thread cache
1 parent a5e27af commit 22b5a1e

File tree

8 files changed

+78
-53
lines changed

8 files changed

+78
-53
lines changed

‎lib/picos/bootstrap/dune

+1-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
11
(library
22
(name picos_bootstrap)
33
(package picos)
4-
(libraries backoff))
4+
(libraries picos.thread backoff))

‎lib/picos/bootstrap/picos_bootstrap.ml

+18
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,10 @@ module Fiber = struct
376376

377377
type t = [ `Fiber ] tdt
378378

379+
module Maybe = struct
380+
type t = T : [< `Nothing | `Fiber ] tdt -> t [@@unboxed]
381+
end
382+
379383
let create_packed ~forbid packed = Fiber { forbid; packed; fls = [||] }
380384

381385
let create ~forbid computation =
@@ -526,3 +530,17 @@ module Handler = struct
526530
await : 'c -> Trigger.t -> (exn * Printexc.raw_backtrace) option;
527531
}
528532
end
533+
534+
module Per_thread = struct
535+
type t = { mutable current : Fiber.Maybe.t (** *) }
536+
537+
let key = Picos_thread.TLS.create ()
538+
539+
let get () =
540+
match Picos_thread.TLS.get_exn key with
541+
| p -> p
542+
| exception Picos_thread.TLS.Not_set ->
543+
let p = { current = T Nothing } in
544+
Picos_thread.TLS.set key p;
545+
p
546+
end

‎lib/picos/ocaml5/picos_ocaml.ml

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

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

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

2832
type _ Effect.t +=
2933
| 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
@@ -27,9 +27,9 @@ module Fiber = struct
2727
include Picos_ocaml.Fiber
2828

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

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

3434
let[@inline] to_fiber_or_current = function
3535
| T Nothing -> current ()
@@ -83,3 +83,7 @@ module Handler = struct
8383
include Picos_bootstrap.Handler
8484
include Picos_ocaml.Handler
8585
end
86+
87+
module Per_thread = struct
88+
include Picos_bootstrap.Per_thread
89+
end

‎lib/picos/picos.mli

+10
Original file line numberDiff line numberDiff line change
@@ -1254,3 +1254,13 @@ module Handler : sig
12541254
implements an effect handler directly, because that is likely to perform
12551255
better. *)
12561256
end
1257+
1258+
module Per_thread : sig
1259+
(** *)
1260+
1261+
type t = { mutable current : Fiber.Maybe.t (** *) }
1262+
(** *)
1263+
1264+
val get : unit -> t
1265+
(** *)
1266+
end

‎lib/picos_mux.fifo/picos_mux_fifo.ml

+14-14
Original file line numberDiff line numberDiff line change
@@ -27,32 +27,32 @@ type t = {
2727
return : ((unit, unit) Effect.Deep.continuation -> unit) option;
2828
discontinue : ((unit, unit) Effect.Deep.continuation -> unit) option;
2929
handler : (unit, unit) Effect.Deep.handler;
30+
per_thread : Per_thread.t;
3031
quota : int;
31-
mutable fiber : Fiber.Maybe.t;
3232
mutable remaining_quota : int;
3333
mutable num_alive_fibers : int;
3434
}
3535

3636
let rec next t =
3737
match Mpscq.pop_exn t.ready with
3838
| Spawn (fiber, main) ->
39-
t.fiber <- Fiber.Maybe.of_fiber fiber;
39+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
4040
t.remaining_quota <- t.quota;
4141
Effect.Deep.match_with main fiber t.handler
4242
| Return (fiber, k) ->
43-
t.fiber <- fiber;
43+
t.per_thread.current <- fiber;
4444
t.remaining_quota <- t.quota;
4545
Effect.Deep.continue k ()
4646
| Continue (fiber, k) ->
47-
t.fiber <- Fiber.Maybe.of_fiber fiber;
47+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
4848
t.remaining_quota <- t.quota;
4949
Fiber.continue fiber k ()
5050
| Resume (fiber, k) ->
51-
t.fiber <- Fiber.Maybe.of_fiber fiber;
51+
t.per_thread.current <- Fiber.Maybe.of_fiber fiber;
5252
t.remaining_quota <- t.quota;
5353
Fiber.resume fiber k
5454
| exception Mpscq.Empty ->
55-
t.fiber <- Fiber.Maybe.nothing;
55+
t.per_thread.current <- Fiber.Maybe.nothing;
5656
if t.num_alive_fibers <> 0 then begin
5757
if Atomic.get t.needs_wakeup then begin
5858
Mutex.lock t.mutex;
@@ -84,7 +84,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
8484
let rec t =
8585
{
8686
ready;
87-
fiber = Fiber.Maybe.nothing;
87+
per_thread = Per_thread.get ();
8888
needs_wakeup;
8989
mutex;
9090
condition;
@@ -101,12 +101,12 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
101101
and current =
102102
Some
103103
(fun k ->
104-
let fiber = Fiber.Maybe.to_fiber t.fiber in
104+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
105105
Effect.Deep.continue k fiber)
106106
and yield =
107107
Some
108108
(fun k ->
109-
let fiber = Fiber.Maybe.to_fiber t.fiber in
109+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
110110
Mpscq.push t.ready (Continue (fiber, k));
111111
next t)
112112
and return =
@@ -118,21 +118,21 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
118118
Effect.Deep.continue k ()
119119
end
120120
else begin
121-
Mpscq.push t.ready (Return (t.fiber, k));
121+
Mpscq.push t.ready (Return (t.per_thread.current, k));
122122
next t
123123
end)
124124
and discontinue =
125125
Some
126126
(fun k ->
127-
let fiber = Fiber.Maybe.to_fiber t.fiber in
127+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
128128
Fiber.continue fiber k ())
129129
and handler = { retc; exnc; effc }
130130
and[@alert "-handler"] effc :
131131
type a. a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option =
132132
function
133133
| Fiber.Current -> t.current
134134
| Fiber.Spawn r ->
135-
let fiber = Fiber.Maybe.to_fiber t.fiber in
135+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
136136
if Fiber.is_canceled fiber then t.discontinue
137137
else begin
138138
t.num_alive_fibers <- t.num_alive_fibers + 1;
@@ -141,7 +141,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
141141
end
142142
| Fiber.Yield -> t.yield
143143
| Computation.Cancel_after r -> begin
144-
let fiber = Fiber.Maybe.to_fiber t.fiber in
144+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
145145
if Fiber.is_canceled fiber then t.discontinue
146146
else
147147
match
@@ -155,7 +155,7 @@ let run_fiber ?quota ?fatal_exn_handler:(exnc : _ = raise) fiber main =
155155
| Trigger.Await trigger ->
156156
Some
157157
(fun k ->
158-
let fiber = Fiber.Maybe.to_fiber t.fiber in
158+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
159159
if Fiber.try_suspend fiber trigger fiber k t.resume then next t
160160
else
161161
let remaining_quota = t.remaining_quota - 1 in

‎lib/picos_mux.random/picos_mux_random.ml

+23-34
Original file line numberDiff line numberDiff line change
@@ -58,17 +58,6 @@ type t = {
5858
mutable run : bool;
5959
}
6060

61-
let fiber_key : Fiber.Maybe.t ref Picos_thread.TLS.t =
62-
Picos_thread.TLS.create ()
63-
64-
let get () =
65-
match Picos_thread.TLS.get_exn fiber_key with
66-
| p -> p
67-
| exception Picos_thread.TLS.Not_set ->
68-
let p = ref Fiber.Maybe.nothing in
69-
Picos_thread.TLS.set fiber_key p;
70-
p
71-
7261
let[@inline] relaxed_wakeup t ~known_not_empty =
7362
if
7463
t.num_waiters_non_zero
@@ -79,24 +68,24 @@ let[@inline] relaxed_wakeup t ~known_not_empty =
7968
Condition.signal t.condition
8069
end
8170

82-
let exec p t = function
71+
let exec (p : Per_thread.t) t = function
8372
| Spawn (fiber, main) ->
84-
p := Fiber.Maybe.of_fiber fiber;
73+
p.current <- Fiber.Maybe.of_fiber fiber;
8574
Effect.Deep.match_with main fiber t.handler
8675
| Raise (fiber, k, exn, bt) ->
87-
p := Fiber.Maybe.of_fiber fiber;
76+
p.current <- Fiber.Maybe.of_fiber fiber;
8877
Effect.Deep.discontinue_with_backtrace k exn bt
8978
| Return (fiber, k) ->
90-
p := Fiber.Maybe.of_fiber fiber;
79+
p.current <- Fiber.Maybe.of_fiber fiber;
9180
Effect.Deep.continue k ()
9281
| Current (fiber, k) ->
93-
p := Fiber.Maybe.of_fiber fiber;
82+
p.current <- Fiber.Maybe.of_fiber fiber;
9483
Effect.Deep.continue k fiber
9584
| Continue (fiber, k) ->
96-
p := Fiber.Maybe.of_fiber fiber;
85+
p.current <- Fiber.Maybe.of_fiber fiber;
9786
Fiber.continue fiber k ()
9887
| Resume (fiber, k) ->
99-
p := Fiber.Maybe.of_fiber fiber;
88+
p.current <- Fiber.Maybe.of_fiber fiber;
10089
Fiber.resume fiber k
10190

10291
let rec next p t =
@@ -105,7 +94,7 @@ let rec next p t =
10594
relaxed_wakeup t ~known_not_empty:false;
10695
exec p t ready
10796
| exception Not_found ->
108-
p := Fiber.Maybe.nothing;
97+
p.current <- Fiber.Maybe.nothing;
10998
if Atomic.get t.num_alive_fibers <> 0 then begin
11099
Mutex.lock t.mutex;
111100
let n = !(t.num_waiters) + 1 in
@@ -190,22 +179,22 @@ let context ?fatal_exn_handler () =
190179
and current =
191180
Some
192181
(fun k ->
193-
let p = Picos_thread.TLS.get_exn fiber_key in
194-
let fiber = Fiber.Maybe.to_fiber !p in
182+
let p = Per_thread.get () in
183+
let fiber = Fiber.Maybe.to_fiber p.current in
195184
Collection.push t.ready (Current (fiber, k));
196185
next p t)
197186
and yield =
198187
Some
199188
(fun k ->
200-
let p = Picos_thread.TLS.get_exn fiber_key in
201-
let fiber = Fiber.Maybe.to_fiber !p in
189+
let p = Per_thread.get () in
190+
let fiber = Fiber.Maybe.to_fiber p.current in
202191
Collection.push t.ready (Continue (fiber, k));
203192
next p t)
204193
and return =
205194
Some
206195
(fun k ->
207-
let p = Picos_thread.TLS.get_exn fiber_key in
208-
let fiber = Fiber.Maybe.to_fiber !p in
196+
let p = Per_thread.get () in
197+
let fiber = Fiber.Maybe.to_fiber p.current in
209198
Collection.push t.ready (Return (fiber, k));
210199
next p t)
211200
and handler = { retc; exnc; effc }
@@ -214,8 +203,8 @@ let context ?fatal_exn_handler () =
214203
function
215204
| Fiber.Current -> t.current
216205
| Fiber.Spawn r ->
217-
let p = Picos_thread.TLS.get_exn fiber_key in
218-
let fiber = Fiber.Maybe.to_fiber !p in
206+
let p = Per_thread.get () in
207+
let fiber = Fiber.Maybe.to_fiber p.current in
219208
if Fiber.is_canceled fiber then t.yield
220209
else begin
221210
Atomic.incr t.num_alive_fibers;
@@ -225,8 +214,8 @@ let context ?fatal_exn_handler () =
225214
end
226215
| Fiber.Yield -> t.yield
227216
| Computation.Cancel_after r -> begin
228-
let p = Picos_thread.TLS.get_exn fiber_key in
229-
let fiber = Fiber.Maybe.to_fiber !p in
217+
let p = Per_thread.get () in
218+
let fiber = Fiber.Maybe.to_fiber p.current in
230219
if Fiber.is_canceled fiber then t.yield
231220
else
232221
match
@@ -243,8 +232,8 @@ let context ?fatal_exn_handler () =
243232
| Trigger.Await trigger ->
244233
Some
245234
(fun k ->
246-
let p = Picos_thread.TLS.get_exn fiber_key in
247-
let fiber = Fiber.Maybe.to_fiber !p in
235+
let p = Per_thread.get () in
236+
let fiber = Fiber.Maybe.to_fiber p.current in
248237
if Fiber.try_suspend fiber trigger fiber k t.resume then next p t
249238
else begin
250239
Collection.push t.ready (Resume (fiber, k));
@@ -253,14 +242,14 @@ let context ?fatal_exn_handler () =
253242
| _ -> None
254243
and retc () =
255244
Atomic.decr t.num_alive_fibers;
256-
let p = Picos_thread.TLS.get_exn fiber_key in
245+
let p = Per_thread.get () in
257246
next p t
258247
in
259248
t
260249

261250
let runner_on_this_thread t =
262251
Select.check_configured ();
263-
next (get ()) t
252+
next (Per_thread.get ()) t
264253

265254
let rec await t =
266255
if t.num_waiters_non_zero then begin
@@ -289,7 +278,7 @@ let run_fiber ?context:t_opt fiber main =
289278
t.run <- true;
290279
Mutex.unlock t.mutex;
291280
Collection.push t.ready (Spawn (fiber, main));
292-
next (get ()) t;
281+
next (Per_thread.get ()) t;
293282
Mutex.lock t.mutex;
294283
await t
295284
end

‎test/dune

+1-1
Original file line numberDiff line numberDiff line change
@@ -88,7 +88,7 @@
8888
(modules test_picos_dscheck picos_bootstrap)
8989
(build_if
9090
(>= %{ocaml_version} 5))
91-
(libraries picos backoff traced_atomic dscheck alcotest))
91+
(libraries picos.thread backoff traced_atomic dscheck alcotest))
9292

9393
;;
9494

0 commit comments

Comments
 (0)