Backlog/lib/p2/src/structures/concurrent-queue.zig

274 lines
8.7 KiB
Zig

const std = @import("std");
const utils = @import("utils.zig");
const ArrayList = std.ArrayList;
const ArrayListUnmanaged = std.ArrayListUnmanaged;
const asserts = utils.asserts;
const Atomic = std.atomic.Value;
pub const ConcurrentQueueError = error{
CorruptedState,
QueueIsEmpty,
QueueIsFull,
};
pub fn ConcurrentQueueU(comptime T: type) type {
return ConcurrentQueueUnmanagedAdvanced(T, .{});
}
pub fn ConcurrentQueueAssert(comptime T: type) type {
return ConcurrentQueueUnmanagedAdvanced(T, .{ .allowAsserts = true });
}
// lock-free concurrent queue, fixed capacity,
// will never resize.
//
// new version based on
// https://github.com/rigtorp/MPMCQueue
pub const ConcurrentStatus = packed struct(usize) {
alive: bool = false,
generation: u63 = 0,
};
pub fn ConcurrentQueueUnmanagedAdvanced(comptime T: type, comptime opts: struct {
allowAsserts: bool = false,
debug: bool = false,
}) type {
return struct {
data: []align(64) T align(64),
status: []align(64) Atomic(ConcurrentStatus) align(64), // 1 = valid, 0 = invalid,
pushId: Atomic(usize) align(64),
popId: Atomic(usize) align(64),
// tail points to next free slot
// head points to the next one to pop
pub fn initCapacity(allocator: std.mem.Allocator, cap: usize) !@This() {
const new = @This(){
.data = try allocator.alignedAlloc(T, .@"64", cap),
.status = try allocator.alignedAlloc(Atomic(ConcurrentStatus), .@"64", cap),
.pushId = Atomic(usize).init(0),
.popId = Atomic(usize).init(0),
};
for (new.status) |*s| {
s.* = Atomic(ConcurrentStatus).init(std.mem.zeroes(ConcurrentStatus));
}
return new;
}
pub fn push(self: *@This(), value: T) !void {
const pushId = self.pushId.fetchAdd(1, .acq_rel);
const slot = @mod(pushId, self.data.len);
const newStatus = ConcurrentStatus{
.generation = @intCast(@divTrunc(pushId, self.data.len)),
.alive = true,
};
if (opts.debug) std.log.warn("pushing {d}", .{pushId});
while (true) {
const status = self.status[slot].load(.acquire);
if (status.generation == newStatus.generation and status.alive == false) {
break;
}
std.atomic.spinLoopHint();
}
if (opts.debug) std.log.warn("pushed {d}", .{pushId});
self.status[slot].store(newStatus, .release);
self.data[slot] = value;
}
// pop the value from the queue, moves the head forward
pub fn pop(self: *@This()) ?T { // things are empty
if (self.count() == 0)
return null;
const popId = self.popId.fetchAdd(1, .acq_rel);
const slot = @mod(popId, self.data.len);
const expectedStatus = ConcurrentStatus{
.generation = @intCast(@divTrunc(popId, self.data.len)),
.alive = true,
};
if (opts.debug) std.log.warn("popping {d}", .{popId});
while (true) {
const status = self.status[slot].load(.acquire);
if (status.generation == expectedStatus.generation and status.alive) {
break;
}
std.atomic.spinLoopHint();
}
if (opts.debug) std.log.warn("popped {d}", .{popId});
self.status[slot].store(.{ .alive = false, .generation = expectedStatus.generation +% 1 }, .release);
return self.data[slot];
}
pub fn count(self: @This()) usize {
const head = self.pushId.load(.acquire);
const tail = self.popId.load(.acquire);
if (opts.debug)
std.log.warn("count {d}", .{head - tail});
if (opts.allowAsserts) {
asserts(tail <= head, "tail > head in concurrent queue, this shouldnt ever happen", .{}, "concurrent queue assert");
}
return head - tail;
}
pub fn capacity(self: @This()) usize {
return self.data.len;
}
pub fn deinit(self: *@This(), allocator: std.mem.Allocator) void {
allocator.free(self.data);
allocator.free(self.status);
}
};
}
test "concurrent queue basic correctness test" {
const Info = struct {
x: u32 = 0,
};
const allocator = std.testing.allocator;
var y = try ConcurrentQueueUnmanagedAdvanced(Info, .{ .allowAsserts = true, .debug = true }).initCapacity(allocator, 420);
defer y.deinit(allocator);
try y.push(.{});
try y.push(.{ .x = 1 });
try y.push(.{ .x = 2 });
try y.push(.{ .x = 1 });
try utils.assertf(y.count() == 4, "expected there to be {d} elements in queue, we saw {d}", .{ 4, y.count() });
_ = y.pop();
_ = y.pop();
_ = y.pop();
_ = y.pop();
try utils.assertf(y.count() == 0, "expected there to be {d} elements in queue, we saw {d}", .{ 0, y.count() });
var x = try ConcurrentQueueU(Info).initCapacity(allocator, 12);
defer x.deinit(allocator);
try x.push(.{ .x = 0 });
try x.push(.{ .x = 1 });
try x.push(.{ .x = 2 });
try x.push(.{ .x = 3 });
try x.push(.{ .x = 4 });
try x.push(.{ .x = 5 });
try x.push(.{ .x = 6 });
try x.push(.{ .x = 7 });
try x.push(.{ .x = 8 });
try x.push(.{ .x = 9 });
try x.push(.{ .x = 10 });
// const maybeError = x.push(.{ .x = 11 });
// try utils.assertf(maybeError == ConcurrentQueueError.QueueIsFull, "Expected queue to have an error", .{});
}
test "concurrent queue multiple producer single consumer" {
// 1. create multiple threads
const threadCount = 24;
const Payload = struct {
x: i64 = 0,
arb: [4096]u8 = undefined,
};
const QueueType = ConcurrentQueueUnmanagedAdvanced(Payload, .{ .debug = false, .allowAsserts = true });
const Wrap = struct {
pub fn threadFunc(queueRef: *QueueType, id: i64, exitSignal: *Atomic(bool), pushedCountResults: *Atomic(i64)) void {
var pushedCount: i64 = 0;
while (!exitSignal.load(.monotonic)) {
queueRef.push(.{
.x = id + pushedCount,
}) catch unreachable;
pushedCount += 1;
}
_ = pushedCountResults.fetchAdd(pushedCount, .seq_cst);
}
};
var threads: [threadCount]std.Thread = undefined;
var testQueue = try QueueType.initCapacity(std.testing.allocator, 4096 * 4);
defer testQueue.deinit(std.testing.allocator);
var exitSignalAtomic = Atomic(bool).init(false);
var pushedCountResults = Atomic(i64).init(0);
for (0..threadCount) |i| {
threads[i] = try std.Thread.spawn(.{}, Wrap.threadFunc, .{ &testQueue, @as(i64, @intCast(i * 10000)), &exitSignalAtomic, &pushedCountResults });
}
// 5 second message pump test
var oldTime: f64 = test_getTime();
// 10 second test, 5 seconds of input, 5 seconds of drain
const startTime: f64 = 2.2;
var timeLeft: f64 = startTime;
var poppedCount: i64 = 0;
var signaled: bool = false;
while (timeLeft > 0 or testQueue.count() > 0) {
const newTime = test_getTime();
const deltaTime = newTime - oldTime;
if (timeLeft - deltaTime < 0.0 and !signaled) {
signaled = true;
exitSignalAtomic.store(true, .seq_cst);
}
timeLeft -= deltaTime;
oldTime = newTime;
if (testQueue.pop()) |x| {
_ = x;
poppedCount += 1;
}
std.atomic.spinLoopHint();
}
for (0..threadCount) |i| {
threads[i].join();
}
std.debug.print("popped {d} entries in {d} seconds payloadSize: {d} dataRate: {d:.3} MiB/s time/event {d:.3} us \n", .{
poppedCount,
startTime,
@sizeOf(Payload),
@as(f64, @floatFromInt(@as(usize, @intCast(poppedCount)) * @sizeOf(Payload))) / startTime / 1024 / 1024,
startTime / @as(f64, @floatFromInt(@as(usize, @intCast(poppedCount)))) * 1000 * 1000,
});
try utils.assertf(
poppedCount == pushedCountResults.load(.seq_cst),
"mismatched, we popped {d} records while the workers pushed {d}",
.{
poppedCount,
pushedCountResults.load(.seq_cst),
},
);
}
fn test_getTime() f64 {
return @as(f64, @floatFromInt(std.time.milliTimestamp())) / 1000;
}