2 files changed, 88 insertions, 1 deletions

diff --git a/std/atomic/queue.zig b/std/atomic/queue.zig
index 54981d4a61..c7ce00d6cf 100644
--- a/std/atomic/queue.zig
+++ b/std/atomic/queue.zig
@@ -1,3 +1,7 @@
+const builtin = @import("builtin");
+const AtomicOrder = builtin.AtomicOrder;
+const AtomicRmwOp = builtin.AtomicRmwOp;
+
 /// Many reader, many writer, non-allocating, thread-safe, lock-free
 pub fn Queue(comptime T: type) type {
     return struct {
@@ -12,7 +16,7 @@ pub fn Queue(comptime T: type) type {
             data: T,
         };
 
-        // TODO: well defined copy elision
+        // TODO: well defined copy elision: https://github.com/zig-lang/zig/issues/287
         pub fn init(self: &Self) void {
             self.root.next = null;
             self.head = &self.root;
@@ -35,3 +39,82 @@ pub fn Queue(comptime T: type) type {
         }
     };
 }
+
+const std = @import("std");
+const Context = struct {
+    allocator: &std.mem.Allocator,
+    queue: &Queue(i32),
+    put_sum: isize,
+    get_sum: isize,
+    get_count: usize,
+    puts_done: u8, // TODO make this a bool
+};
+const puts_per_thread = 10000;
+const put_thread_count = 3;
+
+test "std.atomic.queue" {
+    var direct_allocator = std.heap.DirectAllocator.init();
+    defer direct_allocator.deinit();
+
+    var plenty_of_memory = try direct_allocator.allocator.alloc(u8, 64 * 1024 * 1024);
+    defer direct_allocator.allocator.free(plenty_of_memory);
+
+    var fixed_buffer_allocator = std.heap.ThreadSafeFixedBufferAllocator.init(plenty_of_memory);
+    var a = &fixed_buffer_allocator.allocator;
+
+    var queue: Queue(i32) = undefined;
+    queue.init();
+    var context = Context {
+        .allocator = a,
+        .queue = &queue,
+        .put_sum = 0,
+        .get_sum = 0,
+        .puts_done = 0,
+        .get_count = 0,
+    };
+
+    var putters: [put_thread_count]&std.os.Thread = undefined;
+    for (putters) |*t| {
+        *t = try std.os.spawnThreadAllocator(a, &context, startPuts);
+    }
+    var getters: [put_thread_count]&std.os.Thread = undefined;
+    for (getters) |*t| {
+        *t = try std.os.spawnThreadAllocator(a, &context, startGets);
+    }
+
+    for (putters) |t| t.wait();
+    _ = @atomicRmw(u8, &context.puts_done, builtin.AtomicRmwOp.Xchg, 1, AtomicOrder.SeqCst);
+    for (getters) |t| t.wait();
+
+    std.debug.assert(context.put_sum == context.get_sum);
+    std.debug.assert(context.get_count == puts_per_thread * put_thread_count);
+}
+
+fn startPuts(ctx: &Context) u8 {
+    var put_count: usize = puts_per_thread;
+    var r = std.rand.DefaultPrng.init(0xdeadbeef);
+    while (put_count != 0) : (put_count -= 1) {
+        std.os.time.sleep(0, 1); // let the os scheduler be our fuzz
+        const x = @bitCast(i32, r.random.scalar(u32));
+        const node = ctx.allocator.create(Queue(i32).Node) catch unreachable;
+        node.data = x;
+        ctx.queue.put(node);
+        _ = @atomicRmw(isize, &ctx.put_sum, builtin.AtomicRmwOp.Add, x, AtomicOrder.SeqCst);
+    }
+    return 0;
+}
+
+fn startGets(ctx: &Context) u8 {
+    while (true) {
+        while (ctx.queue.get()) |node| {
+            std.os.time.sleep(0, 1); // let the os scheduler be our fuzz
+            _ = @atomicRmw(isize, &ctx.get_sum, builtin.AtomicRmwOp.Add, node.data, builtin.AtomicOrder.SeqCst);
+            _ = @atomicRmw(usize, &ctx.get_count, builtin.AtomicRmwOp.Add, 1, builtin.AtomicOrder.SeqCst);
+        }
+
+        if (@atomicLoad(u8, &ctx.puts_done, builtin.AtomicOrder.SeqCst) == 1) {
+            break;
+        }
+    }
+    return 0;
+}
diff --git a/std/atomic/stack.zig b/std/atomic/stack.zig
index a1e686155c..a53f7682b0 100644
--- a/std/atomic/stack.zig
+++ b/std/atomic/stack.zig
@@ -53,6 +53,7 @@ const Context = struct {
     stack: &Stack(i32),
     put_sum: isize,
     get_sum: isize,
+    get_count: usize,
     puts_done: u8, // TODO make this a bool
 };
 const puts_per_thread = 1000;
@@ -75,6 +76,7 @@ test "std.atomic.stack" {
         .put_sum = 0,
         .get_sum = 0,
         .puts_done = 0,
+        .get_count = 0,
     };
 
     var putters: [put_thread_count]&std.os.Thread = undefined;
@@ -91,6 +93,7 @@ test "std.atomic.stack" {
     for (getters) |t| t.wait();
 
     std.debug.assert(context.put_sum == context.get_sum);
+    std.debug.assert(context.get_count == puts_per_thread * put_thread_count);
 }
 
 fn startPuts(ctx: &Context) u8 {
@@ -112,6 +115,7 @@ fn startGets(ctx: &Context) u8 {
         while (ctx.stack.pop()) |node| {
             std.os.time.sleep(0, 1); // let the os scheduler be our fuzz
             _ = @atomicRmw(isize, &ctx.get_sum, builtin.AtomicRmwOp.Add, node.data, builtin.AtomicOrder.SeqCst);
+            _ = @atomicRmw(usize, &ctx.get_count, builtin.AtomicRmwOp.Add, 1, builtin.AtomicOrder.SeqCst);
         }
 
         if (@atomicLoad(u8, &ctx.puts_done, builtin.AtomicOrder.SeqCst) == 1) {