lib/coz/src/delay.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 pub const ThreadState = struct {
4 in_use: std.atomic.Value(bool) = .init(false),
5 local_delay_ns: std.atomic.Value(u64) = .init(0),
6 pre_block_time_ns: u64 = 0,
7 is_blocked: std.atomic.Value(bool) = .init(false),
8
9 pub fn setInUse(self: *ThreadState, value: bool) void {
10 self.in_use.store(value, .seq_cst);
11 }
12
13 pub fn acquireUse(self: *ThreadState) bool {
14 return self.in_use.cmpxchgStrong(false, true, .seq_cst, .seq_cst) == null;
15 }
16
17 pub fn checkInUse(self: *const ThreadState) bool {
18 return self.in_use.load(.seq_cst);
19 }
20
21 pub fn localDelay(self: *const ThreadState) u64 {
22 return self.local_delay_ns.load(.monotonic);
23 }
24 };
25
26 pub const Coordinator = struct {
27 experiment_active: std.atomic.Value(bool) = .init(false),
28 global_delay_ns: std.atomic.Value(u64) = .init(0),
29 delay_size_ns: std.atomic.Value(u64) = .init(0),
30 overshoot_ns: std.atomic.Value(u64) = .init(0),
31 capped_waits: bool = false,
32
33 pub fn startExperiment(self: *Coordinator, delay_size_ns: u64) void {
34 self.delay_size_ns.store(delay_size_ns, .monotonic);
35 self.overshoot_ns.store(0, .monotonic);
36 self.experiment_active.store(true, .release);
37 }
38
39 pub fn finishExperiment(self: *Coordinator) void {
40 self.experiment_active.store(false, .release);
41 }
42
43 pub fn active(self: *const Coordinator) bool {
44 return self.experiment_active.load(.acquire);
45 }
46
47 pub fn globalDelay(self: *const Coordinator) u64 {
48 return self.global_delay_ns.load(.monotonic);
49 }
50
51 pub fn delaySize(self: *const Coordinator) u64 {
52 return self.delay_size_ns.load(.monotonic);
53 }
54
55 pub fn overshoot(self: *const Coordinator) u64 {
56 return self.overshoot_ns.load(.monotonic);
57 }
58
59 pub fn preBlock(self: *Coordinator, thread: *ThreadState) void {
60 thread.is_blocked.store(true, .release);
61 thread.pre_block_time_ns = self.globalDelay();
62 }
63
64 pub fn postBlock(self: *Coordinator, thread: *ThreadState, skip_delays: bool) void {
65 thread.setInUse(true);
66 defer thread.setInUse(false);
67
68 if (skip_delays) {
69 const delta = self.globalDelay() -| thread.pre_block_time_ns;
70 _ = thread.local_delay_ns.fetchAdd(delta, .monotonic);
71 }
72 thread.is_blocked.store(false, .release);
73 }
74
75 pub fn creditSelectedHit(self: *Coordinator, thread: *ThreadState) void {
76 if (!self.active()) return;
77 _ = thread.local_delay_ns.fetchAdd(self.delaySize(), .monotonic);
78 }
79
80 pub fn creditSelectedThreadAndPush(self: *Coordinator, thread: *ThreadState) u64 {
81 if (!self.active()) return self.globalDelay();
82
83 const delay_size = self.delaySize();
84 const new_global = self.global_delay_ns.fetchAdd(delay_size, .monotonic) + delay_size;
85 var local = thread.localDelay();
86 while (local < new_global) {
87 if (thread.local_delay_ns.cmpxchgWeak(local, new_global, .monotonic, .monotonic) == null) break;
88 local = thread.localDelay();
89 }
90 return new_global;
91 }
92
93 pub fn addDelays(self: *Coordinator, thread: *ThreadState, wait: anytype) u64 {
94 if (self.active()) {
95 if (thread.is_blocked.load(.acquire)) return 0;
96
97 const global = self.globalDelay();
98 const local = thread.localDelay();
99 if (local > global) {
100 _ = self.global_delay_ns.fetchAdd(local - global, .monotonic);
101 return 0;
102 }
103 if (local < global) {
104 const needed = global - local;
105 const waited = callWait(wait, needed);
106 if (self.capped_waits) {
107 _ = thread.local_delay_ns.fetchAdd(needed, .monotonic);
108 if (waited > needed) _ = self.overshoot_ns.fetchAdd(waited - needed, .monotonic);
109 } else {
110 _ = thread.local_delay_ns.fetchAdd(waited, .monotonic);
111 }
112 return waited;
113 }
114 return 0;
115 }
116
117 thread.local_delay_ns.store(self.globalDelay(), .monotonic);
118 return 0;
119 }
120
121 pub fn inheritDelay(_: *Coordinator, child: *ThreadState, parent_delay_ns: u64) void {
122 child.local_delay_ns.store(parent_delay_ns, .monotonic);
123 }
124 };
125
126 fn callWait(wait: anytype, ns: u64) u64 {
127 switch (@typeInfo(@TypeOf(wait))) {
128 .@"struct", .@"enum", .@"union", .@"opaque" => if (comptime @hasDecl(@TypeOf(wait), "wait")) return wait.wait(ns),
129 else => {},
130 }
131 return wait(ns);
132 }
133
134 fn exactWait(ns: u64) u64 {
135 return ns;
136 }
137
138 fn overshootingWait(ns: u64) u64 {
139 return ns + 3;
140 }
141
142 const CountingWait = struct {
143 calls: *u32,
144
145 pub fn wait(self: CountingWait, ns: u64) u64 {
146 self.calls.* += 1;
147 return ns;
148 }
149 };
150
151 test "thread state publishes in-use changes" {
152 var thread: ThreadState = .{};
153
154 try std.testing.expect(!thread.checkInUse());
155 try std.testing.expect(thread.acquireUse());
156 try std.testing.expect(!thread.acquireUse());
157 thread.setInUse(true);
158 try std.testing.expect(thread.checkInUse());
159 thread.setInUse(false);
160 try std.testing.expect(!thread.checkInUse());
161 }
162
163 test "selected hit advances sampled thread then pushes global delay" {
164 var coordinator: Coordinator = .{};
165 var sampled: ThreadState = .{};
166 var other: ThreadState = .{};
167
168 coordinator.startExperiment(10);
169 coordinator.creditSelectedHit(&sampled);
170 try std.testing.expectEqual(@as(u64, 10), sampled.localDelay());
171
172 try std.testing.expectEqual(@as(u64, 0), coordinator.addDelays(&sampled, exactWait));
173 try std.testing.expectEqual(@as(u64, 10), coordinator.globalDelay());
174
175 try std.testing.expectEqual(@as(u64, 10), coordinator.addDelays(&other, exactWait));
176 try std.testing.expectEqual(@as(u64, 10), other.localDelay());
177 }
178
179 test "inactive coordinator catches thread local delay up without waiting" {
180 var coordinator: Coordinator = .{};
181 var thread: ThreadState = .{};
182 coordinator.global_delay_ns.store(25, .monotonic);
183
184 try std.testing.expectEqual(@as(u64, 0), coordinator.addDelays(&thread, exactWait));
185 try std.testing.expectEqual(@as(u64, 25), thread.localDelay());
186 }
187
188 test "blocked threads do not wait until post block clears the flag" {
189 var coordinator: Coordinator = .{};
190 var thread: ThreadState = .{};
191
192 coordinator.startExperiment(0);
193 coordinator.global_delay_ns.store(20, .monotonic);
194 coordinator.preBlock(&thread);
195
196 try std.testing.expectEqual(@as(u64, 0), coordinator.addDelays(&thread, exactWait));
197 try std.testing.expectEqual(@as(u64, 0), thread.localDelay());
198
199 coordinator.postBlock(&thread, false);
200 try std.testing.expectEqual(@as(u64, 20), coordinator.addDelays(&thread, exactWait));
201 try std.testing.expectEqual(@as(u64, 20), thread.localDelay());
202 }
203
204 test "post block can skip delays inserted while blocked" {
205 var coordinator: Coordinator = .{};
206 var thread: ThreadState = .{};
207
208 coordinator.startExperiment(0);
209 coordinator.global_delay_ns.store(10, .monotonic);
210 coordinator.preBlock(&thread);
211 coordinator.global_delay_ns.store(35, .monotonic);
212 coordinator.postBlock(&thread, true);
213
214 try std.testing.expectEqual(@as(u64, 25), thread.localDelay());
215 try std.testing.expectEqual(@as(u64, 10), thread.pre_block_time_ns);
216 try std.testing.expectEqual(@as(u64, 10), coordinator.addDelays(&thread, exactWait));
217 try std.testing.expectEqual(@as(u64, 35), thread.localDelay());
218 }
219
220 test "selected sampled thread credit pushes global delay without delaying sampled thread" {
221 var coordinator: Coordinator = .{};
222 var sampled: ThreadState = .{};
223 var other: ThreadState = .{};
224
225 coordinator.startExperiment(12);
226 try std.testing.expectEqual(@as(u64, 12), coordinator.creditSelectedThreadAndPush(&sampled));
227 try std.testing.expectEqual(@as(u64, 12), sampled.localDelay());
228 try std.testing.expectEqual(@as(u64, 12), coordinator.globalDelay());
229 try std.testing.expectEqual(@as(u64, 12), coordinator.addDelays(&other, exactWait));
230 }
231
232 test "capped waits charge requested delay and record overshoot separately" {
233 var coordinator: Coordinator = .{ .capped_waits = true };
234 var thread: ThreadState = .{};
235
236 coordinator.startExperiment(0);
237 coordinator.global_delay_ns.store(10, .monotonic);
238
239 try std.testing.expectEqual(@as(u64, 13), coordinator.addDelays(&thread, overshootingWait));
240 try std.testing.expectEqual(@as(u64, 10), thread.localDelay());
241 try std.testing.expectEqual(@as(u64, 3), coordinator.overshoot());
242 }
243
244 test "delay coordinator accepts wait objects" {
245 var coordinator: Coordinator = .{};
246 var thread: ThreadState = .{};
247 var calls: u32 = 0;
248
249 coordinator.startExperiment(0);
250 coordinator.global_delay_ns.store(11, .monotonic);
251
252 try std.testing.expectEqual(@as(u64, 11), coordinator.addDelays(&thread, CountingWait{ .calls = &calls }));
253 try std.testing.expectEqual(@as(u32, 1), calls);
254 try std.testing.expectEqual(@as(u64, 11), thread.localDelay());
255 }
256
257 test "child threads inherit parent delay credit" {
258 var coordinator: Coordinator = .{};
259 var child: ThreadState = .{};
260
261 coordinator.inheritDelay(&child, 17);
262
263 try std.testing.expectEqual(@as(u64, 17), child.localDelay());
264 }