1//! A lock that supports one writer or many readers.
2const RwLock = @This();
3
4const builtin = @import("builtin");
5
6const std = @import("../std.zig");
7const Io = std.Io;
8const assert = std.debug.assert;
9const testing = std.testing;
10
11state: usize,
12mutex: Io.Mutex,
13semaphore: Io.Semaphore,
14
15pub const init: RwLock = .{
16 .state = 0,
17 .mutex = .init,
18 .semaphore = .{},
19};
20
21const is_writing: usize = 1;
22const writer: usize = 1 << 1;
23const reader: usize = 1 << (1 + @bitSizeOf(Count));
24const writer_mask: usize = std.math.maxInt(Count) << @ctz(writer);
25const reader_mask: usize = std.math.maxInt(Count) << @ctz(reader);
26const Count = @Int(.unsigned, @divFloor(@bitSizeOf(usize) - 1, 2));
27
28pub fn tryLock(rl: *RwLock, io: Io) bool {
29 if (rl.mutex.tryLock()) {
30 // Unlike `lock`, this never registers in `writer_mask`, so holding the mutex does
31 // not stop a reader from taking the fast path; the CAS catches one that raced in
32 // after `state` was loaded.
33 const state = @atomicLoad(usize, &rl.state, .seq_cst);
34 if (state & reader_mask == 0) {
35 _ = @cmpxchgStrong(
36 usize,
37 &rl.state,
38 state,
39 state | is_writing,
40 .seq_cst,
41 .seq_cst,
42 ) orelse return true;
43 }
44
45 rl.mutex.unlock(io);
46 }
47
48 return false;
49}
50
51pub fn lockUncancelable(rl: *RwLock, io: Io) void {
52 _ = @atomicRmw(usize, &rl.state, .Add, writer, .seq_cst);
53 rl.mutex.lockUncancelable(io);
54
55 const state = @atomicRmw(usize, &rl.state, .Add, is_writing -% writer, .seq_cst);
56 if (state & reader_mask != 0)
57 rl.semaphore.waitUncancelable(io);
58}
59
60pub fn lock(rl: *RwLock, io: Io) Io.Cancelable!void {
61 _ = @atomicRmw(usize, &rl.state, .Add, writer, .seq_cst);
62 rl.mutex.lock(io) catch |err| switch (err) {
63 error.Canceled => {
64 _ = @atomicRmw(usize, &rl.state, .Sub, writer, .seq_cst);
65 return error.Canceled;
66 },
67 };
68
69 const state = @atomicRmw(usize, &rl.state, .Add, is_writing -% writer, .seq_cst);
70 if (state & reader_mask != 0)
71 rl.semaphore.wait(io) catch |err| switch (err) {
72 error.Canceled => {
73 // Clearing `is_writing` while still holding the mutex means the last reader
74 // either saw it set, and posts a permit only we can consume, or did not and
75 // never posts. A stale permit would let the next writer in past the readers.
76 const prev_state = @atomicRmw(usize, &rl.state, .And, ~is_writing, .seq_cst);
77 if (prev_state & reader_mask == 0) rl.semaphore.waitUncancelable(io);
78 rl.mutex.unlock(io);
79 return error.Canceled;
80 },
81 };
82}
83
84pub fn unlock(rl: *RwLock, io: Io) void {
85 _ = @atomicRmw(usize, &rl.state, .And, ~is_writing, .seq_cst);
86 rl.mutex.unlock(io);
87}
88
89pub fn tryLockShared(rl: *RwLock, io: Io) bool {
90 const state = @atomicLoad(usize, &rl.state, .seq_cst);
91 if (state & (is_writing | writer_mask) == 0) {
92 _ = @cmpxchgStrong(
93 usize,
94 &rl.state,
95 state,
96 state + reader,
97 .seq_cst,
98 .seq_cst,
99 ) orelse return true;
100 }
101
102 if (rl.mutex.tryLock()) {
103 _ = @atomicRmw(usize, &rl.state, .Add, reader, .seq_cst);
104 rl.mutex.unlock(io);
105 return true;
106 }
107
108 return false;
109}
110
111pub fn lockSharedUncancelable(rl: *RwLock, io: Io) void {
112 var state = @atomicLoad(usize, &rl.state, .seq_cst);
113 while (state & (is_writing | writer_mask) == 0) {
114 state = @cmpxchgWeak(
115 usize,
116 &rl.state,
117 state,
118 state + reader,
119 .seq_cst,
120 .seq_cst,
121 ) orelse return;
122 }
123
124 rl.mutex.lockUncancelable(io);
125 _ = @atomicRmw(usize, &rl.state, .Add, reader, .seq_cst);
126 rl.mutex.unlock(io);
127}
128
129pub fn lockShared(rl: *RwLock, io: Io) Io.Cancelable!void {
130 var state = @atomicLoad(usize, &rl.state, .seq_cst);
131 while (state & (is_writing | writer_mask) == 0) {
132 state = @cmpxchgWeak(
133 usize,
134 &rl.state,
135 state,
136 state + reader,
137 .seq_cst,
138 .seq_cst,
139 ) orelse return;
140 }
141
142 try rl.mutex.lock(io);
143 _ = @atomicRmw(usize, &rl.state, .Add, reader, .seq_cst);
144 rl.mutex.unlock(io);
145}
146
147pub fn unlockShared(rl: *RwLock, io: Io) void {
148 const state = @atomicRmw(usize, &rl.state, .Sub, reader, .seq_cst);
149
150 if ((state & reader_mask == reader) and (state & is_writing != 0))
151 rl.semaphore.post(io);
152}
153
154test "internal state" {
155 const io = testing.io;
156
157 var rl: Io.RwLock = .init;
158
159 // The following failed prior to the fix for Issue #13163,
160 // where the WRITER flag was subtracted by the lock method.
161
162 rl.lockUncancelable(io);
163 rl.unlock(io);
164 try testing.expectEqual(rl, Io.RwLock.init);
165
166 try rl.lock(io);
167 rl.unlock(io);
168 try testing.expectEqual(rl, Io.RwLock.init);
169}
170
171test "smoke test" {
172 const io = testing.io;
173
174 var rl: Io.RwLock = .init;
175
176 rl.lockUncancelable(io);
177 try testing.expect(!rl.tryLock(io));
178 try testing.expect(!rl.tryLockShared(io));
179 rl.unlock(io);
180
181 try rl.lock(io);
182 try testing.expect(!rl.tryLock(io));
183 try testing.expect(!rl.tryLockShared(io));
184 rl.unlock(io);
185
186 try testing.expect(rl.tryLock(io));
187 try testing.expect(!rl.tryLock(io));
188 try testing.expect(!rl.tryLockShared(io));
189 rl.unlock(io);
190
191 rl.lockSharedUncancelable(io);
192 try testing.expect(!rl.tryLock(io));
193 try testing.expect(rl.tryLockShared(io));
194 rl.unlockShared(io);
195 rl.unlockShared(io);
196
197 try testing.expect(rl.tryLockShared(io));
198 try testing.expect(!rl.tryLock(io));
199 try testing.expect(rl.tryLockShared(io));
200 rl.unlockShared(io);
201 rl.unlockShared(io);
202
203 rl.lockUncancelable(io);
204 rl.unlock(io);
205}
206
207test "concurrent access" {
208 if (builtin.single_threaded) return;
209
210 const io = testing.io;
211 const num_writers: usize = 2;
212 const num_readers: usize = 4;
213 const num_writes: usize = 1000;
214 const num_reads: usize = 2000;
215
216 const Runner = struct {
217 const Runner = @This();
218
219 io: Io,
220
221 rl: Io.RwLock,
222 writes: usize,
223 reads: std.atomic.Value(usize),
224
225 val_a: usize,
226 val_b: usize,
227
228 fn reader(run: *Runner, thread_idx: usize) !void {
229 var prng = std.Random.DefaultPrng.init(thread_idx);
230 const rnd = prng.random();
231 while (true) {
232 run.rl.lockSharedUncancelable(run.io);
233 defer run.rl.unlockShared(run.io);
234
235 try testing.expect(run.writes <= num_writes);
236 if (run.reads.fetchAdd(1, .monotonic) >= num_reads) break;
237
238 // We use `volatile` accesses so that we can make sure the memory is accessed either
239 // side of a yield, maximising chances of a race.
240 const a_ptr: *const volatile usize = &run.val_a;
241 const b_ptr: *const volatile usize = &run.val_b;
242
243 const old_a = a_ptr.*;
244 if (rnd.boolean()) try std.Thread.yield();
245 const old_b = b_ptr.*;
246 try testing.expect(old_a == old_b);
247 }
248 }
249
250 fn writer(run: *Runner, thread_idx: usize) !void {
251 var prng = std.Random.DefaultPrng.init(thread_idx);
252 const rnd = prng.random();
253 while (true) {
254 run.rl.lockUncancelable(run.io);
255 defer run.rl.unlock(run.io);
256
257 try testing.expect(run.writes <= num_writes);
258 if (run.writes == num_writes) break;
259
260 // We use `volatile` accesses so that we can make sure the memory is accessed either
261 // side of a yield, maximising chances of a race.
262 const a_ptr: *volatile usize = &run.val_a;
263 const b_ptr: *volatile usize = &run.val_b;
264
265 const new_val = rnd.int(usize);
266
267 const old_a = a_ptr.*;
268 a_ptr.* = new_val;
269 if (rnd.boolean()) try std.Thread.yield();
270 const old_b = b_ptr.*;
271 b_ptr.* = new_val;
272 try testing.expect(old_a == old_b);
273
274 run.writes += 1;
275 }
276 }
277 };
278
279 var run: Runner = .{
280 .io = io,
281 .rl = .init,
282 .writes = 0,
283 .reads = .init(0),
284 .val_a = 0,
285 .val_b = 0,
286 };
287 var write_threads: [num_writers]std.Thread = undefined;
288 var read_threads: [num_readers]std.Thread = undefined;
289
290 for (&write_threads, 0..) |*t, i| t.* = try .spawn(.{}, Runner.writer, .{ &run, i });
291 for (&read_threads, num_writers..) |*t, i| t.* = try .spawn(.{}, Runner.reader, .{ &run, i });
292
293 for (write_threads) |t| t.join();
294 for (read_threads) |t| t.join();
295
296 try testing.expect(run.writes == num_writes);
297 try testing.expect(run.reads.raw >= num_reads);
298}
299
300test "lock canceling" {
301 const io = testing.io;
302
303 var rl: Io.RwLock = .init;
304
305 rl.lockSharedUncancelable(io);
306 var sfuture = io.concurrent(semaphoreLockCancel, .{ &rl, io }) catch |err| switch (err) {
307 error.ConcurrencyUnavailable => return error.SkipZigTest,
308 };
309 try std.testing.expectEqual(error.Canceled, sfuture.cancel(io));
310 rl.unlockShared(io);
311 try testing.expectEqual(rl, Io.RwLock.init);
312
313 rl.lockUncancelable(io);
314 var mfuture = io.concurrent(mutexLockCancel, .{ &rl, io }) catch |err| switch (err) {
315 error.ConcurrencyUnavailable => return error.SkipZigTest,
316 };
317 try std.testing.expectEqual(error.Canceled, mfuture.cancel(io));
318 rl.unlock(io);
319 try testing.expectEqual(rl, Io.RwLock.init);
320}
321
322test "tryLock does not race with readers" {
323 if (builtin.single_threaded) return error.SkipZigTest;
324
325 const Context = struct {
326 rl: Io.RwLock,
327
328 fn reader(ctx: *@This(), io: Io) !void {
329 while (true) {
330 if (ctx.rl.tryLockShared(io)) ctx.rl.unlockShared(io);
331 try io.checkCancel();
332 }
333 }
334 };
335
336 var ctx: Context = .{ .rl = .init };
337 const io = testing.io;
338
339 var future = io.concurrent(Context.reader, .{ &ctx, io }) catch |err| switch (err) {
340 error.ConcurrencyUnavailable => return error.SkipZigTest,
341 };
342 defer future.cancel(io) catch {};
343
344 for (0..1000) |_| {
345 if (!ctx.rl.tryLock(io)) continue;
346 defer ctx.rl.unlock(io);
347 const state = @atomicLoad(usize, &ctx.rl.state, .seq_cst);
348 try testing.expectEqual(0, state & reader_mask);
349 }
350
351 try testing.expectEqual(0, ctx.rl.semaphore.permits);
352}
353
354test "canceled writer does not leak a semaphore permit" {
355 if (builtin.single_threaded) return error.SkipZigTest;
356
357 const Writer = struct {
358 fn lockUnlock(rl: *Io.RwLock, io: Io) Io.Cancelable!void {
359 try rl.lock(io);
360 rl.unlock(io);
361 }
362 };
363
364 const io = testing.io;
365
366 var rl: Io.RwLock = .init;
367
368 for (0..1000) |_| {
369 rl.lockSharedUncancelable(io);
370
371 var wfuture = io.concurrent(Writer.lockUnlock, .{ &rl, io }) catch |err| switch (err) {
372 error.ConcurrencyUnavailable => {
373 rl.unlockShared(io);
374 return error.SkipZigTest;
375 },
376 };
377 var rfuture = io.concurrent(Io.RwLock.unlockShared, .{ &rl, io }) catch |err| switch (err) {
378 error.ConcurrencyUnavailable => {
379 rl.unlockShared(io);
380 wfuture.await(io) catch {};
381 return error.SkipZigTest;
382 },
383 };
384
385 // Races the last reader's `post` against the writer's cancelation.
386 wfuture.cancel(io) catch {};
387 rfuture.await(io);
388
389 try testing.expectEqual(0, rl.state);
390 try testing.expectEqual(0, rl.semaphore.permits);
391 }
392}
393
394fn semaphoreLockCancel(rl: *Io.RwLock, io: Io) !void {
395 try rl.lock(io); //tests semaphore cancelling
396}
397
398fn mutexLockCancel(rl: *Io.RwLock, io: Io) !void {
399 //tests mutex canceling
400 try std.testing.expectEqual(error.Canceled, rl.lockShared(io));
401 io.recancel();
402 try std.testing.expectEqual(error.Canceled, rl.lock(io));
403 return error.Canceled;
404}