queueConcurrent.zig

const std = @import("std");
const Allocator = std.mem.Allocator;
const atomic = std.atomic;

/// Queue as ring buffer. There is one `T` of margin.
pub fn QueueConcurrent(T: type) type {
    return struct {
        buf: []atomic.Value(T),
        head: atomic.Value(usize),
        tail: atomic.Value(usize),

        const Self = @This();

        pub fn cap(self: *Self) usize {
            return self.buf.len;
        }

        pub fn enqueue(self: *Self, item: T) error{OutOfCapacity}!void {
            while (true) {
                const head_original = self.head.load(.acquire);
                const head_next = (head_original + 1) % self.buf.len;
                const tail = self.tail.load(.acquire);
                if (head_next == tail) return error.OutOfCapacity;
                if (self.head.cmpxchgWeak(head_original, head_next, .release, .monotonic) == null) {
                    self.buf[head_original].store(item, .monotonic);
                    break;
                }
            }
        }

        pub fn dequeue(self: *Self) ?T {
            while (true) {
                const tail_original = self.tail.load(.acquire);
                const tail_next = (tail_original + 1) % self.buf.len;
                const head = self.head.load(.acquire);
                if (tail_original == head) return null;
                if (self.tail.cmpxchgWeak(tail_original, tail_next, .release, .monotonic) == null) {
                    const item = self.buf[tail_original].load(.monotonic);
                    return item;
                }
            }
        }

        pub fn init(alloc: Allocator, size: usize) !Self {
            const buf = try alloc.alloc(atomic.Value(T), size);
            return Self{
                .buf = buf,
                .head = atomic.Value(usize).init(0),
                .tail = atomic.Value(usize).init(0),
            };
        }

        pub fn deinit(self: *Self, alloc: Allocator) void {
            alloc.free(self.buf);
        }
    };
}

const expect = std.testing.expect;
const expectEqual = std.testing.expectEqual;
const expectError = std.testing.expectError;
const Thread = std.Thread;

test "enqueue" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u8) = try .init(alloc, 100);
    try queue.enqueue(5);
    try queue.enqueue(6);
    try queue.enqueue(7);
    try queue.enqueue(8);

    try expectEqual(5, queue.buf[0].load(.monotonic));
    try expectEqual(8, queue.buf[3].load(.monotonic));
}

test "dequeue" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u8) = try .init(alloc, 100);

    try queue.enqueue(5);
    try queue.enqueue(6);
    try queue.enqueue(7);
    try queue.enqueue(8);

    try expectEqual(5, queue.dequeue().?);
    try expectEqual(6, queue.dequeue().?);
    try expectEqual(7, queue.dequeue().?);
    try expectEqual(8, queue.dequeue().?);
    try expectEqual(null, queue.dequeue());
}

test "error enqueue with not enough space" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u8) = try .init(alloc, 8);

    for (0..7) |_| {
        try queue.enqueue(0);
    }

    try expectError(error.OutOfCapacity, queue.enqueue(0));
}

test "queue wraps" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u8) = try .init(alloc, 4);

    try queue.enqueue(1);
    try queue.enqueue(2);
    try queue.enqueue(3);
    try expectError(error.OutOfCapacity, queue.enqueue(4));
    try expectEqual(1, queue.dequeue().?);
    try expectEqual(2, queue.dequeue().?);
    try queue.enqueue(5);
    try expectEqual(3, queue.dequeue().?);
    try expectEqual(5, queue.dequeue().?);
}

const q = 1_000_000;
const n = 200_000;
fn enqueue_thread(queue: *QueueConcurrent(u32)) !void {
    for (0..n) |i| {
        try queue.enqueue(@intCast(i));
    }
}

fn dequeue_thread(queue: *QueueConcurrent(u32), black_hole: *atomic.Value(u32)) !void {
    for (0..n) |_| {
        const res = queue.dequeue().?;
        _ = black_hole.fetchXor(res, .monotonic);
    }
}

fn dequeue_same_time(queue: *QueueConcurrent(u32), success: *atomic.Value(usize), fail: *atomic.Value(usize)) !void {
    for (0..n) |_| {
        if (queue.dequeue() != null) {
            _ = success.fetchAdd(1, .monotonic);
        } else {
            _ = fail.fetchAdd(1, .monotonic);
        }
    }
}

test "multithreading" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u32) = try .init(alloc, q);
    defer queue.deinit(alloc);

    var a = try Thread.spawn(.{}, enqueue_thread, .{&queue});
    var b = try Thread.spawn(.{}, enqueue_thread, .{&queue});
    var c = try Thread.spawn(.{}, enqueue_thread, .{&queue});
    var d = try Thread.spawn(.{}, enqueue_thread, .{&queue});
    a.join();
    b.join();
    c.join();
    d.join();

    try expectEqual(n * 4, queue.head.load(.monotonic));

    var black_hole = atomic.Value(u32).init(0);
    var e = try Thread.spawn(.{}, dequeue_thread, .{ &queue, &black_hole });
    var f = try Thread.spawn(.{}, dequeue_thread, .{ &queue, &black_hole });
    var g = try Thread.spawn(.{}, dequeue_thread, .{ &queue, &black_hole });
    var h = try Thread.spawn(.{}, dequeue_thread, .{ &queue, &black_hole });
    e.join();
    f.join();
    g.join();
    h.join();
    try expectEqual(0, black_hole.load(.monotonic));
    try expectEqual(n * 4, queue.tail.load(.monotonic));
}

test "all at once" {
    var gpa = std.heap.DebugAllocator(.{}).init;
    const alloc = gpa.allocator();
    var queue: QueueConcurrent(u32) = try .init(alloc, q);
    defer queue.deinit(alloc);

    var a = try Thread.spawn(.{}, enqueue_thread, .{&queue});
    var fail = atomic.Value(usize).init(0);
    var success = atomic.Value(usize).init(0);
    var e = try Thread.spawn(.{}, dequeue_same_time, .{ &queue, &success, &fail });
    var f = try Thread.spawn(.{}, dequeue_same_time, .{ &queue, &success, &fail });
    var g = try Thread.spawn(.{}, dequeue_same_time, .{ &queue, &success, &fail });
    var h = try Thread.spawn(.{}, dequeue_same_time, .{ &queue, &success, &fail });
    e.join();
    f.join();
    a.join();
    g.join();
    h.join();

    try expectEqual(n * 4, fail.load(.monotonic) + success.load(.monotonic));
    try expectEqual(n, queue.head.load(.monotonic));
}