Skip to content

Commit d64a573

Browse files
polytypicpolytypic
committed
Introduce Per_thread cache
1 parent fbc4124 commit d64a573

File tree

8 files changed

+92
-67
lines changed

8 files changed

+92
-67
lines changed

‎lib/picos/dune

+12-13
Original file line numberDiff line numberDiff line change
@@ -35,26 +35,25 @@
3535

3636
;;
3737

38-
(library
39-
(name picos_thread_ocaml4)
40-
(package picos)
41-
(modules)
38+
(rule
4239
(enabled_if
4340
(< %{ocaml_version} 5.0.0))
44-
(libraries
45-
(re_export picos.thread)))
41+
(action
42+
(copy intf.ocaml4.ml intf.ml)))
43+
44+
(rule
45+
(enabled_if
46+
(<= 5.0.0 %{ocaml_version}))
47+
(action
48+
(copy intf.ocaml5.ml intf.ml)))
49+
50+
;;
4651

4752
(library
4853
(name picos)
4954
(public_name picos)
5055
(modules picos intf)
51-
(libraries
52-
backoff
53-
(select
54-
intf.ml
55-
from
56-
(picos_thread_ocaml4 -> intf.ocaml4.ml)
57-
(-> intf.ocaml5.ml))))
56+
(libraries backoff picos.thread))
5857

5958
(mdx
6059
(package picos_meta)

‎lib/picos/intf.ocaml4.ml

+1
Original file line numberDiff line numberDiff line change
@@ -8,5 +8,6 @@ end
88

99
module type Fiber = sig
1010
type t
11+
type maybe
1112
type _ computation
1213
end

‎lib/picos/intf.ocaml5.ml

+11
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ end
6565

6666
module type Fiber = sig
6767
type t
68+
type maybe
6869
type _ computation
6970

7071
val resume :
@@ -152,4 +153,14 @@ module type Fiber = sig
152153
type _ Effect.t +=
153154
private
154155
| Spawn : { fiber : t; main : t -> unit } -> unit Effect.t
156+
157+
module Per_thread : sig
158+
(** *)
159+
160+
type t = { mutable current : maybe (** *) }
161+
(** *)
162+
163+
val get : unit -> t
164+
(** *)
165+
end
155166
end

‎lib/picos/picos.mli

+4-1
Original file line numberDiff line numberDiff line change
@@ -1192,7 +1192,10 @@ module Fiber : sig
11921192
@raise Invalid_argument if the trigger is not in the signaled state. *)
11931193

11941194
include
1195-
Intf.Fiber with type t := t with type 'a computation := 'a Computation.t
1195+
Intf.Fiber
1196+
with type t := t
1197+
with type maybe = Maybe.t
1198+
with type 'a computation := 'a Computation.t
11961199

11971200
(** {2 Design rationale}
11981201

‎lib/picos/picos.ocaml5.ml

+22-2
Original file line numberDiff line numberDiff line change
@@ -588,6 +588,8 @@ module Fiber = struct
588588
if key < Array.length fls then Array.unsafe_set fls key unique
589589
end
590590

591+
type maybe = T : [< `Nothing | `Fiber ] tdt -> maybe [@@unboxed]
592+
591593
(* END FIBER BOOTSTRAP *)
592594

593595
let resume t k = Effect.Deep.continue k (canceled t)
@@ -603,9 +605,27 @@ module Fiber = struct
603605
| None -> Effect.Shallow.continue_with k v h
604606
| Some (exn, bt) -> Effect.Shallow.discontinue_with_backtrace k exn bt h
605607

608+
module Per_thread = struct
609+
type t = { mutable current : maybe (** *) }
610+
611+
let key = Picos_thread.TLS.create ()
612+
613+
let get () =
614+
match Picos_thread.TLS.get_exn key with
615+
| p -> p
616+
| exception Picos_thread.TLS.Not_set ->
617+
let p = { current = T Nothing } in
618+
Picos_thread.TLS.set key p;
619+
p
620+
end
621+
606622
type _ Effect.t += Current : t Effect.t
607623

608-
let current () = Effect.perform Current
624+
let current () =
625+
let p = Per_thread.get () in
626+
match p.current with
627+
| T (Fiber _ as fiber) -> fiber
628+
| T Nothing -> Effect.perform Current
609629

610630
type _ Effect.t += Spawn : { fiber : t; main : t -> unit } -> unit Effect.t
611631

@@ -620,7 +640,7 @@ module Fiber = struct
620640
module Maybe = struct
621641
let[@inline never] not_a_fiber _ = invalid_arg "not a fiber"
622642

623-
type t = T : [< `Nothing | `Fiber ] tdt -> t [@@unboxed]
643+
type t = maybe
624644

625645
let[@inline] to_fiber_or_current = function
626646
| T Nothing -> current ()

‎lib/picos_mux.fifo/picos_mux_fifo.ml

+14-11
Original file line numberDiff line numberDiff line change
@@ -27,8 +27,8 @@ type t = {
2727
mutable return : ((unit, unit) Effect.Deep.continuation -> unit) option;
2828
mutable discontinue : ((unit, unit) Effect.Deep.continuation -> unit) option;
2929
mutable handler : (unit, unit) Effect.Deep.handler;
30+
per_thread : Fiber.Per_thread.t;
3031
quota : int;
31-
mutable fiber : Fiber.Maybe.t;
3232
mutable remaining_quota : int;
3333
mutable num_alive_fibers : int;
3434
}
@@ -37,7 +37,7 @@ let rec next t =
3737
match Mpscq.pop_exn t.ready with
3838
| ready -> begin
3939
t.remaining_quota <- t.quota;
40-
t.fiber <-
40+
t.per_thread.current <-
4141
(match ready with
4242
| Spawn (fiber, _)
4343
| Continue (fiber, _)
@@ -51,7 +51,7 @@ let rec next t =
5151
| Resume (fiber, k) -> Fiber.resume fiber k
5252
end
5353
| exception Mpscq.Empty ->
54-
t.fiber <- Fiber.Maybe.nothing;
54+
t.per_thread.current <- Fiber.Maybe.nothing;
5555
if t.num_alive_fibers <> 0 then begin
5656
if Atomic.get t.needs_wakeup then begin
5757
Mutex.lock t.mutex;
@@ -75,6 +75,8 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
7575
| None -> Int.max_int
7676
| Some quota -> if quota <= 0 then quota_non_positive quota else quota
7777
in
78+
let per_thread = Fiber.Per_thread.get () in
79+
per_thread.current <- Fiber.Maybe.of_fiber fiber;
7880
{
7981
ready = Mpscq.create ~padded:true ();
8082
needs_wakeup = Atomic.make false |> Multicore_magic.copy_as_padded;
@@ -86,8 +88,8 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
8688
return = Obj.magic ();
8789
discontinue = Obj.magic ();
8890
handler = Obj.magic ();
91+
per_thread;
8992
quota;
90-
fiber = Fiber.Maybe.of_fiber fiber;
9193
remaining_quota = quota;
9294
num_alive_fibers = 1;
9395
}
@@ -101,7 +103,7 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
101103
match e with
102104
| Fiber.Current -> t.current
103105
| Fiber.Spawn r ->
104-
let fiber = Fiber.Maybe.to_fiber t.fiber in
106+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
105107
if Fiber.is_canceled fiber then t.discontinue
106108
else begin
107109
t.num_alive_fibers <- t.num_alive_fibers + 1;
@@ -110,7 +112,7 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
110112
end
111113
| Fiber.Yield -> t.yield
112114
| Computation.Cancel_after r -> begin
113-
let fiber = Fiber.Maybe.to_fiber t.fiber in
115+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
114116
if Fiber.is_canceled fiber then t.discontinue
115117
else
116118
match
@@ -126,7 +128,7 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
126128
| Trigger.Await trigger ->
127129
Some
128130
(fun k ->
129-
let fiber = Fiber.Maybe.to_fiber t.fiber in
131+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
130132
if Fiber.try_suspend fiber trigger fiber k t.resume then
131133
next t
132134
else
@@ -164,12 +166,12 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
164166
t.current <-
165167
Some
166168
(fun k ->
167-
let fiber = Fiber.Maybe.to_fiber t.fiber in
169+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
168170
Effect.Deep.continue k fiber);
169171
t.yield <-
170172
Some
171173
(fun k ->
172-
let fiber = Fiber.Maybe.to_fiber t.fiber in
174+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
173175
Mpscq.push t.ready (Continue (fiber, k));
174176
next t);
175177
t.return <-
@@ -181,13 +183,14 @@ let run_fiber ?quota ?fatal_exn_handler fiber main =
181183
Effect.Deep.continue k ()
182184
end
183185
else begin
184-
Mpscq.push t.ready (Return (Fiber.Maybe.to_fiber t.fiber, k));
186+
Mpscq.push t.ready
187+
(Return (Fiber.Maybe.to_fiber t.per_thread.current, k));
185188
next t
186189
end);
187190
t.discontinue <-
188191
Some
189192
(fun k ->
190-
let fiber = Fiber.Maybe.to_fiber t.fiber in
193+
let fiber = Fiber.Maybe.to_fiber t.per_thread.current in
191194
Fiber.continue fiber k ());
192195
Effect.Deep.match_with main fiber t.handler
193196

‎lib/picos_mux.random/picos_mux_random.ml

+27-39
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,18 +68,16 @@ let[@inline] relaxed_wakeup t ~known_not_empty =
7968
Condition.signal t.condition
8069
end
8170

82-
let exec p t ready =
83-
begin
84-
p :=
85-
match ready with
86-
| Spawn (fiber, _)
87-
| Raise (fiber, _, _, _)
88-
| Return (fiber, _)
89-
| Current (fiber, _)
90-
| Continue (fiber, _)
91-
| Resume (fiber, _) ->
92-
Fiber.Maybe.of_fiber fiber
93-
end;
71+
let exec (p : Fiber.Per_thread.t) t ready =
72+
p.current <-
73+
(match ready with
74+
| Spawn (fiber, _)
75+
| Raise (fiber, _, _, _)
76+
| Return (fiber, _)
77+
| Current (fiber, _)
78+
| Continue (fiber, _)
79+
| Resume (fiber, _) ->
80+
Fiber.Maybe.of_fiber fiber);
9481
match ready with
9582
| Spawn (fiber, main) -> Effect.Deep.match_with main fiber t.handler
9683
| Raise (_, k, exn, bt) -> Effect.Deep.discontinue_with_backtrace k exn bt
@@ -105,7 +92,7 @@ let rec next p t =
10592
relaxed_wakeup t ~known_not_empty:false;
10693
exec p t ready
10794
| exception Not_found ->
108-
p := Fiber.Maybe.nothing;
95+
p.current <- Fiber.Maybe.nothing;
10996
if Atomic.get t.num_alive_fibers <> 0 then begin
11097
Mutex.lock t.mutex;
11198
let n = !(t.num_waiters) + 1 in
@@ -192,30 +179,30 @@ let context ?fatal_exn_handler () =
192179
t.current <-
193180
Some
194181
(fun k ->
195-
let p = Picos_thread.TLS.get_exn fiber_key in
196-
let fiber = Fiber.Maybe.to_fiber !p in
182+
let p = Fiber.Per_thread.get () in
183+
let fiber = Fiber.Maybe.to_fiber p.current in
197184
Collection.push t.ready (Current (fiber, k));
198185
next p t);
199186
t.yield <-
200187
Some
201188
(fun k ->
202-
let p = Picos_thread.TLS.get_exn fiber_key in
203-
let fiber = Fiber.Maybe.to_fiber !p in
189+
let p = Fiber.Per_thread.get () in
190+
let fiber = Fiber.Maybe.to_fiber p.current in
204191
Collection.push t.ready (Continue (fiber, k));
205192
next p t);
206193
t.return <-
207194
Some
208195
(fun k ->
209-
let p = Picos_thread.TLS.get_exn fiber_key in
210-
let fiber = Fiber.Maybe.to_fiber !p in
196+
let p = Fiber.Per_thread.get () in
197+
let fiber = Fiber.Maybe.to_fiber p.current in
211198
Collection.push t.ready (Return (fiber, k));
212199
next p t);
213200
t.handler <-
214201
{
215202
retc =
216203
(fun () ->
217204
Atomic.decr t.num_alive_fibers;
218-
let p = Picos_thread.TLS.get_exn fiber_key in
205+
let p = Fiber.Per_thread.get () in
219206
next p t);
220207
exnc;
221208
effc =
@@ -224,8 +211,8 @@ let context ?fatal_exn_handler () =
224211
match e with
225212
| Fiber.Current -> t.current
226213
| Fiber.Spawn r ->
227-
let p = Picos_thread.TLS.get_exn fiber_key in
228-
let fiber = Fiber.Maybe.to_fiber !p in
214+
let p = Fiber.Per_thread.get () in
215+
let fiber = Fiber.Maybe.to_fiber p.current in
229216
if Fiber.is_canceled fiber then t.yield
230217
else begin
231218
Atomic.incr t.num_alive_fibers;
@@ -235,8 +222,8 @@ let context ?fatal_exn_handler () =
235222
end
236223
| Fiber.Yield -> t.yield
237224
| Computation.Cancel_after r -> begin
238-
let p = Picos_thread.TLS.get_exn fiber_key in
239-
let fiber = Fiber.Maybe.to_fiber !p in
225+
let p = Fiber.Per_thread.get () in
226+
let fiber = Fiber.Maybe.to_fiber p.current in
240227
if Fiber.is_canceled fiber then t.yield
241228
else
242229
match
@@ -254,8 +241,8 @@ let context ?fatal_exn_handler () =
254241
| Trigger.Await trigger ->
255242
Some
256243
(fun k ->
257-
let p = Picos_thread.TLS.get_exn fiber_key in
258-
let fiber = Fiber.Maybe.to_fiber !p in
244+
let p = Fiber.Per_thread.get () in
245+
let fiber = Fiber.Maybe.to_fiber p.current in
259246
if Fiber.try_suspend fiber trigger fiber k t.resume then
260247
next p t
261248
else begin
@@ -268,7 +255,7 @@ let context ?fatal_exn_handler () =
268255

269256
let runner_on_this_thread t =
270257
Select.check_configured ();
271-
next (get ()) t
258+
next (Fiber.Per_thread.get ()) t
272259

273260
let rec await t =
274261
if t.num_waiters_non_zero then begin
@@ -296,7 +283,8 @@ let run_fiber ?context:t_opt fiber main =
296283
else begin
297284
t.run <- true;
298285
Mutex.unlock t.mutex;
299-
get () := Fiber.Maybe.of_fiber fiber;
286+
let p = Fiber.Per_thread.get () in
287+
p.current <- Fiber.Maybe.of_fiber fiber;
300288
Effect.Deep.match_with main fiber t.handler;
301289
Mutex.lock t.mutex;
302290
await t

‎test/dune

+1-1
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@
8787
(modules test_picos_dscheck picos)
8888
(build_if
8989
(>= %{ocaml_version} 5))
90-
(libraries backoff traced_atomic dscheck alcotest))
90+
(libraries picos.thread backoff traced_atomic dscheck alcotest))
9191

9292
;;
9393

0 commit comments

Comments
 (0)