lib/sys/src/atomic/sequence.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const builtin = @import("builtin");
  3 
  4 const Counter = if (builtin.single_threaded)
  5     Plain
  6 else if (@bitSizeOf(usize) >= @bitSizeOf(u64))
  7     Wide
  8 else
  9     Locked;
 10 
 11 const retry_limit = std.math.maxInt(u32);
 12 
 13 pub const Sequence64 = struct {
 14     counter: Counter,
 15 
 16     pub fn init(value: u64) Sequence64 {
 17         return .{ .counter = Counter.init(value) };
 18     }
 19 
 20     pub fn take(self: *Sequence64) ?u64 {
 21         return self.counter.take();
 22     }
 23 };
 24 
 25 const Plain = struct {
 26     value: u64,
 27 
 28     fn init(value: u64) Plain {
 29         return .{ .value = value };
 30     }
 31 
 32     fn take(self: *Plain) ?u64 {
 33         if (self.value == std.math.maxInt(u64)) return null;
 34         self.value += 1;
 35         return self.value;
 36     }
 37 };
 38 
 39 const Wide = struct {
 40     value: std.atomic.Value(u64),
 41 
 42     fn init(value: u64) Wide {
 43         return .{ .value = .init(value) };
 44     }
 45 
 46     fn take(self: *Wide) ?u64 {
 47         var current = self.value.load(.monotonic);
 48         var retries: u32 = 0;
 49         while (current != std.math.maxInt(u64)) : (retries += 1) {
 50             std.debug.assert(retries < retry_limit);
 51             const next = current + 1;
 52             if (self.value.cmpxchgWeak(
 53                 current,
 54                 next,
 55                 .monotonic,
 56                 .monotonic,
 57             )) |observed| {
 58                 current = observed;
 59             } else return next;
 60         }
 61         return null;
 62     }
 63 };
 64 
 65 const Locked = struct {
 66     guard: std.atomic.Value(u32),
 67     value: u64,
 68 
 69     fn init(value: u64) Locked {
 70         return .{ .guard = .init(0), .value = value };
 71     }
 72 
 73     fn take(self: *Locked) ?u64 {
 74         self.acquire();
 75         defer self.release();
 76         if (self.value == std.math.maxInt(u64)) return null;
 77         self.value += 1;
 78         return self.value;
 79     }
 80 
 81     fn acquire(self: *Locked) void {
 82         var retries: u32 = 0;
 83         while (self.guard.cmpxchgWeak(
 84             0,
 85             1,
 86             .acquire,
 87             .monotonic,
 88         ) != null) : (retries += 1) {
 89             std.debug.assert(retries < retry_limit);
 90             std.atomic.spinLoopHint();
 91         }
 92     }
 93 
 94     fn release(self: *Locked) void {
 95         self.guard.store(0, .release);
 96     }
 97 };
 98 
 99 test "64 bit sequence issues nonzero identities and never wraps" {
100     const testing = std.testing;
101     var sequence = Sequence64.init(0);
102     try testing.expectEqual(@as(?u64, 1), sequence.take());
103     try testing.expectEqual(@as(?u64, 2), sequence.take());
104 
105     var exhausted = Sequence64.init(std.math.maxInt(u64) - 1);
106     try testing.expectEqual(@as(?u64, std.math.maxInt(u64)), exhausted.take());
107     try testing.expectEqual(@as(?u64, null), exhausted.take());
108     try testing.expectEqual(@as(?u64, null), exhausted.take());
109 }
110 
111 fn expectSaturation(comptime CounterType: type) !void {
112     var counter = CounterType.init(std.math.maxInt(u64) - 1);
113     try std.testing.expectEqual(@as(?u64, std.math.maxInt(u64)), counter.take());
114     try std.testing.expectEqual(@as(?u64, null), counter.take());
115     try std.testing.expectEqual(@as(?u64, null), counter.take());
116 }
117 
118 test "64 bit sequence backends saturate without wrapping" {
119     try expectSaturation(Plain);
120     try expectSaturation(Wide);
121     try expectSaturation(Locked);
122 }
123 
124 const concurrent_workers = 8;
125 const values_per_worker = 64;
126 const concurrent_values = concurrent_workers * values_per_worker;
127 
128 fn takeConcurrent(counter: anytype, values: []u64) void {
129     for (values) |*value| value.* = counter.take() orelse unreachable;
130 }
131 
132 fn expectConcurrentSequence(comptime CounterType: type) !void {
133     if (builtin.single_threaded) return error.SkipZigTest;
134 
135     var counter = CounterType.init(0);
136     var values: [concurrent_values]u64 = undefined;
137     var threads: [concurrent_workers]std.Thread = undefined;
138     for (&threads, 0..) |*thread, index| {
139         const begin = index * values_per_worker;
140         const end = begin + values_per_worker;
141         thread.* = try std.Thread.spawn(
142             .{},
143             takeConcurrent,
144             .{ &counter, values[begin..end] },
145         );
146     }
147     for (threads) |thread| thread.join();
148 
149     std.mem.sort(u64, &values, {}, std.sort.asc(u64));
150     for (values, 0..) |value, index| {
151         try std.testing.expectEqual(@as(u64, index + 1), value);
152     }
153 }
154 
155 test "64 bit sequence wide counter is unique under contention" {
156     try expectConcurrentSequence(Wide);
157 }
158 
159 test "64 bit sequence locked counter is unique under contention" {
160     try expectConcurrentSequence(Locked);
161 }