| ... | @@ -1,413 +0,0 @@ |
| 1 | const std = @import("../std.zig"); |
| 2 | const builtin = @import("builtin"); |
| 3 | const assert = std.debug.assert; |
| 4 | const expect = std.testing.expect; |
| 5 | |
| 6 | /// Many producer, many consumer, non-allocating, thread-safe. |
| 7 | /// Uses a mutex to protect access. |
| 8 | /// The queue does not manage ownership and the user is responsible to |
| 9 | /// manage the storage of the nodes. |
| 10 | pub fn Queue(comptime T: type) type { |
| 11 | return struct { |
| 12 | head: ?*Node, |
| 13 | tail: ?*Node, |
| 14 | mutex: std.Thread.Mutex, |
| 15 | |
| 16 | pub const Self = @This(); |
| 17 | pub const Node = std.DoublyLinkedList(T).Node; |
| 18 | |
| 19 | /// Initializes a new queue. The queue does not provide a `deinit()` |
| 20 | /// function, so the user must take care of cleaning up the queue elements. |
| 21 | pub fn init() Self { |
| 22 | return Self{ |
| 23 | .head = null, |
| 24 | .tail = null, |
| 25 | .mutex = std.Thread.Mutex{}, |
| 26 | }; |
| 27 | } |
| 28 | |
| 29 | /// Appends `node` to the queue. |
| 30 | /// The lifetime of `node` must be longer than the lifetime of the queue. |
| 31 | pub fn put(self: *Self, node: *Node) void { |
| 32 | node.next = null; |
| 33 | |
| 34 | self.mutex.lock(); |
| 35 | defer self.mutex.unlock(); |
| 36 | |
| 37 | node.prev = self.tail; |
| 38 | self.tail = node; |
| 39 | if (node.prev) |prev_tail| { |
| 40 | prev_tail.next = node; |
| 41 | } else { |
| 42 | assert(self.head == null); |
| 43 | self.head = node; |
| 44 | } |
| 45 | } |
| 46 | |
| 47 | /// Gets a previously inserted node or returns `null` if there is none. |
| 48 | /// It is safe to `get()` a node from the queue while another thread tries |
| 49 | /// to `remove()` the same node at the same time. |
| 50 | pub fn get(self: *Self) ?*Node { |
| 51 | self.mutex.lock(); |
| 52 | defer self.mutex.unlock(); |
| 53 | |
| 54 | const head = self.head orelse return null; |
| 55 | self.head = head.next; |
| 56 | if (head.next) |new_head| { |
| 57 | new_head.prev = null; |
| 58 | } else { |
| 59 | self.tail = null; |
| 60 | } |
| 61 | // This way, a get() and a remove() are thread-safe with each other. |
| 62 | head.prev = null; |
| 63 | head.next = null; |
| 64 | return head; |
| 65 | } |
| 66 | |
| 67 | /// Prepends `node` to the front of the queue. |
| 68 | /// The lifetime of `node` must be longer than the lifetime of the queue. |
| 69 | pub fn unget(self: *Self, node: *Node) void { |
| 70 | node.prev = null; |
| 71 | |
| 72 | self.mutex.lock(); |
| 73 | defer self.mutex.unlock(); |
| 74 | |
| 75 | const opt_head = self.head; |
| 76 | self.head = node; |
| 77 | if (opt_head) |old_head| { |
| 78 | node.next = old_head; |
| 79 | } else { |
| 80 | assert(self.tail == null); |
| 81 | self.tail = node; |
| 82 | } |
| 83 | } |
| 84 | |
| 85 | /// Removes a node from the queue, returns whether node was actually removed. |
| 86 | /// It is safe to `remove()` a node from the queue while another thread tries |
| 87 | /// to `get()` the same node at the same time. |
| 88 | pub fn remove(self: *Self, node: *Node) bool { |
| 89 | self.mutex.lock(); |
| 90 | defer self.mutex.unlock(); |
| 91 | |
| 92 | if (node.prev == null and node.next == null and self.head != node) { |
| 93 | return false; |
| 94 | } |
| 95 | |
| 96 | if (node.prev) |prev| { |
| 97 | prev.next = node.next; |
| 98 | } else { |
| 99 | self.head = node.next; |
| 100 | } |
| 101 | if (node.next) |next| { |
| 102 | next.prev = node.prev; |
| 103 | } else { |
| 104 | self.tail = node.prev; |
| 105 | } |
| 106 | node.prev = null; |
| 107 | node.next = null; |
| 108 | return true; |
| 109 | } |
| 110 | |
| 111 | /// Returns `true` if the queue is currently empty. |
| 112 | /// Note that in a multi-consumer environment a return value of `false` |
| 113 | /// does not mean that `get` will yield a non-`null` value! |
| 114 | pub fn isEmpty(self: *Self) bool { |
| 115 | self.mutex.lock(); |
| 116 | defer self.mutex.unlock(); |
| 117 | return self.head == null; |
| 118 | } |
| 119 | |
| 120 | /// Dumps the contents of the queue to `stderr`. |
| 121 | pub fn dump(self: *Self) void { |
| 122 | self.dumpToStream(std.io.getStdErr().writer()) catch return; |
| 123 | } |
| 124 | |
| 125 | /// Dumps the contents of the queue to `stream`. |
| 126 | /// Up to 4 elements from the head are dumped and the tail of the queue is |
| 127 | /// dumped as well. |
| 128 | pub fn dumpToStream(self: *Self, stream: anytype) !void { |
| 129 | const S = struct { |
| 130 | fn dumpRecursive( |
| 131 | s: anytype, |
| 132 | optional_node: ?*Node, |
| 133 | indent: usize, |
| 134 | comptime depth: comptime_int, |
| 135 | ) !void { |
| 136 | try s.writeByteNTimes(' ', indent); |
| 137 | if (optional_node) |node| { |
| 138 | try s.print("0x{x}={}\n", .{ @intFromPtr(node), node.data }); |
| 139 | if (depth == 0) { |
| 140 | try s.print("(max depth)\n", .{}); |
| 141 | return; |
| 142 | } |
| 143 | try dumpRecursive(s, node.next, indent + 1, depth - 1); |
| 144 | } else { |
| 145 | try s.print("(null)\n", .{}); |
| 146 | } |
| 147 | } |
| 148 | }; |
| 149 | self.mutex.lock(); |
| 150 | defer self.mutex.unlock(); |
| 151 | |
| 152 | try stream.print("head: ", .{}); |
| 153 | try S.dumpRecursive(stream, self.head, 0, 4); |
| 154 | try stream.print("tail: ", .{}); |
| 155 | try S.dumpRecursive(stream, self.tail, 0, 4); |
| 156 | } |
| 157 | }; |
| 158 | } |
| 159 | |
| 160 | const Context = struct { |
| 161 | allocator: std.mem.Allocator, |
| 162 | queue: *Queue(i32), |
| 163 | put_sum: isize, |
| 164 | get_sum: isize, |
| 165 | get_count: usize, |
| 166 | puts_done: bool, |
| 167 | }; |
| 168 | |
| 169 | // TODO add lazy evaluated build options and then put puts_per_thread behind |
| 170 | // some option such as: "AggressiveMultithreadedFuzzTest". In the AppVeyor |
| 171 | // CI we would use a less aggressive setting since at 1 core, while we still |
| 172 | // want this test to pass, we need a smaller value since there is so much thrashing |
| 173 | // we would also use a less aggressive setting when running in valgrind |
| 174 | const puts_per_thread = 500; |
| 175 | const put_thread_count = 3; |
| 176 | |
| 177 | test "std.atomic.Queue" { |
| 178 | const plenty_of_memory = try std.heap.page_allocator.alloc(u8, 300 * 1024); |
| 179 | defer std.heap.page_allocator.free(plenty_of_memory); |
| 180 | |
| 181 | var fixed_buffer_allocator = std.heap.FixedBufferAllocator.init(plenty_of_memory); |
| 182 | const a = fixed_buffer_allocator.threadSafeAllocator(); |
| 183 | |
| 184 | var queue = Queue(i32).init(); |
| 185 | var context = Context{ |
| 186 | .allocator = a, |
| 187 | .queue = &queue, |
| 188 | .put_sum = 0, |
| 189 | .get_sum = 0, |
| 190 | .puts_done = false, |
| 191 | .get_count = 0, |
| 192 | }; |
| 193 | |
| 194 | if (builtin.single_threaded) { |
| 195 | try expect(context.queue.isEmpty()); |
| 196 | { |
| 197 | var i: usize = 0; |
| 198 | while (i < put_thread_count) : (i += 1) { |
| 199 | try expect(startPuts(&context) == 0); |
| 200 | } |
| 201 | } |
| 202 | try expect(!context.queue.isEmpty()); |
| 203 | context.puts_done = true; |
| 204 | { |
| 205 | var i: usize = 0; |
| 206 | while (i < put_thread_count) : (i += 1) { |
| 207 | try expect(startGets(&context) == 0); |
| 208 | } |
| 209 | } |
| 210 | try expect(context.queue.isEmpty()); |
| 211 | } else { |
| 212 | try expect(context.queue.isEmpty()); |
| 213 | |
| 214 | var putters: [put_thread_count]std.Thread = undefined; |
| 215 | for (&putters) |*t| { |
| 216 | t.* = try std.Thread.spawn(.{}, startPuts, .{&context}); |
| 217 | } |
| 218 | var getters: [put_thread_count]std.Thread = undefined; |
| 219 | for (&getters) |*t| { |
| 220 | t.* = try std.Thread.spawn(.{}, startGets, .{&context}); |
| 221 | } |
| 222 | |
| 223 | for (putters) |t| |
| 224 | t.join(); |
| 225 | @atomicStore(bool, &context.puts_done, true, .SeqCst); |
| 226 | for (getters) |t| |
| 227 | t.join(); |
| 228 | |
| 229 | try expect(context.queue.isEmpty()); |
| 230 | } |
| 231 | |
| 232 | if (context.put_sum != context.get_sum) { |
| 233 | std.debug.panic("failure\nput_sum:{} != get_sum:{}", .{ context.put_sum, context.get_sum }); |
| 234 | } |
| 235 | |
| 236 | if (context.get_count != puts_per_thread * put_thread_count) { |
| 237 | std.debug.panic("failure\nget_count:{} != puts_per_thread:{} * put_thread_count:{}", .{ |
| 238 | context.get_count, |
| 239 | @as(u32, puts_per_thread), |
| 240 | @as(u32, put_thread_count), |
| 241 | }); |
| 242 | } |
| 243 | } |
| 244 | |
| 245 | fn startPuts(ctx: *Context) u8 { |
| 246 | var put_count: usize = puts_per_thread; |
| 247 | var prng = std.rand.DefaultPrng.init(0xdeadbeef); |
| 248 | const random = prng.random(); |
| 249 | while (put_count != 0) : (put_count -= 1) { |
| 250 | std.time.sleep(1); // let the os scheduler be our fuzz |
| 251 | const x = @as(i32, @bitCast(random.int(u32))); |
| 252 | const node = ctx.allocator.create(Queue(i32).Node) catch unreachable; |
| 253 | node.* = .{ |
| 254 | .prev = undefined, |
| 255 | .next = undefined, |
| 256 | .data = x, |
| 257 | }; |
| 258 | ctx.queue.put(node); |
| 259 | _ = @atomicRmw(isize, &ctx.put_sum, .Add, x, .SeqCst); |
| 260 | } |
| 261 | return 0; |
| 262 | } |
| 263 | |
| 264 | fn startGets(ctx: *Context) u8 { |
| 265 | while (true) { |
| 266 | const last = @atomicLoad(bool, &ctx.puts_done, .SeqCst); |
| 267 | |
| 268 | while (ctx.queue.get()) |node| { |
| 269 | std.time.sleep(1); // let the os scheduler be our fuzz |
| 270 | _ = @atomicRmw(isize, &ctx.get_sum, .Add, node.data, .SeqCst); |
| 271 | _ = @atomicRmw(usize, &ctx.get_count, .Add, 1, .SeqCst); |
| 272 | } |
| 273 | |
| 274 | if (last) return 0; |
| 275 | } |
| 276 | } |
| 277 | |
| 278 | test "std.atomic.Queue single-threaded" { |
| 279 | var queue = Queue(i32).init(); |
| 280 | try expect(queue.isEmpty()); |
| 281 | |
| 282 | var node_0 = Queue(i32).Node{ |
| 283 | .data = 0, |
| 284 | .next = undefined, |
| 285 | .prev = undefined, |
| 286 | }; |
| 287 | queue.put(&node_0); |
| 288 | try expect(!queue.isEmpty()); |
| 289 | |
| 290 | var node_1 = Queue(i32).Node{ |
| 291 | .data = 1, |
| 292 | .next = undefined, |
| 293 | .prev = undefined, |
| 294 | }; |
| 295 | queue.put(&node_1); |
| 296 | try expect(!queue.isEmpty()); |
| 297 | |
| 298 | try expect(queue.get().?.data == 0); |
| 299 | try expect(!queue.isEmpty()); |
| 300 | |
| 301 | var node_2 = Queue(i32).Node{ |
| 302 | .data = 2, |
| 303 | .next = undefined, |
| 304 | .prev = undefined, |
| 305 | }; |
| 306 | queue.put(&node_2); |
| 307 | try expect(!queue.isEmpty()); |
| 308 | |
| 309 | var node_3 = Queue(i32).Node{ |
| 310 | .data = 3, |
| 311 | .next = undefined, |
| 312 | .prev = undefined, |
| 313 | }; |
| 314 | queue.put(&node_3); |
| 315 | try expect(!queue.isEmpty()); |
| 316 | |
| 317 | try expect(queue.get().?.data == 1); |
| 318 | try expect(!queue.isEmpty()); |
| 319 | |
| 320 | try expect(queue.get().?.data == 2); |
| 321 | try expect(!queue.isEmpty()); |
| 322 | |
| 323 | var node_4 = Queue(i32).Node{ |
| 324 | .data = 4, |
| 325 | .next = undefined, |
| 326 | .prev = undefined, |
| 327 | }; |
| 328 | queue.put(&node_4); |
| 329 | try expect(!queue.isEmpty()); |
| 330 | |
| 331 | try expect(queue.get().?.data == 3); |
| 332 | node_3.next = null; |
| 333 | try expect(!queue.isEmpty()); |
| 334 | |
| 335 | queue.unget(&node_3); |
| 336 | try expect(queue.get().?.data == 3); |
| 337 | try expect(!queue.isEmpty()); |
| 338 | |
| 339 | try expect(queue.get().?.data == 4); |
| 340 | try expect(queue.isEmpty()); |
| 341 | |
| 342 | try expect(queue.get() == null); |
| 343 | try expect(queue.isEmpty()); |
| 344 | |
| 345 | // unget an empty queue |
| 346 | queue.unget(&node_4); |
| 347 | try expect(queue.tail == &node_4); |
| 348 | try expect(queue.head == &node_4); |
| 349 | |
| 350 | try expect(queue.get().?.data == 4); |
| 351 | |
| 352 | try expect(queue.get() == null); |
| 353 | try expect(queue.isEmpty()); |
| 354 | } |
| 355 | |
| 356 | test "std.atomic.Queue dump" { |
| 357 | const mem = std.mem; |
| 358 | var buffer: [1024]u8 = undefined; |
| 359 | var expected_buffer: [1024]u8 = undefined; |
| 360 | var fbs = std.io.fixedBufferStream(&buffer); |
| 361 | |
| 362 | var queue = Queue(i32).init(); |
| 363 | |
| 364 | // Test empty stream |
| 365 | fbs.reset(); |
| 366 | try queue.dumpToStream(fbs.writer()); |
| 367 | try expect(mem.eql(u8, buffer[0..fbs.pos], |
| 368 | \\head: (null) |
| 369 | \\tail: (null) |
| 370 | \\ |
| 371 | )); |
| 372 | |
| 373 | // Test a stream with one element |
| 374 | var node_0 = Queue(i32).Node{ |
| 375 | .data = 1, |
| 376 | .next = undefined, |
| 377 | .prev = undefined, |
| 378 | }; |
| 379 | queue.put(&node_0); |
| 380 | |
| 381 | fbs.reset(); |
| 382 | try queue.dumpToStream(fbs.writer()); |
| 383 | |
| 384 | var expected = try std.fmt.bufPrint(expected_buffer[0..], |
| 385 | \\head: 0x{x}=1 |
| 386 | \\ (null) |
| 387 | \\tail: 0x{x}=1 |
| 388 | \\ (null) |
| 389 | \\ |
| 390 | , .{ @intFromPtr(queue.head), @intFromPtr(queue.tail) }); |
| 391 | try expect(mem.eql(u8, buffer[0..fbs.pos], expected)); |
| 392 | |
| 393 | // Test a stream with two elements |
| 394 | var node_1 = Queue(i32).Node{ |
| 395 | .data = 2, |
| 396 | .next = undefined, |
| 397 | .prev = undefined, |
| 398 | }; |
| 399 | queue.put(&node_1); |
| 400 | |
| 401 | fbs.reset(); |
| 402 | try queue.dumpToStream(fbs.writer()); |
| 403 | |
| 404 | expected = try std.fmt.bufPrint(expected_buffer[0..], |
| 405 | \\head: 0x{x}=1 |
| 406 | \\ 0x{x}=2 |
| 407 | \\ (null) |
| 408 | \\tail: 0x{x}=2 |
| 409 | \\ (null) |
| 410 | \\ |
| 411 | , .{ @intFromPtr(queue.head), @intFromPtr(queue.head.?.next), @intFromPtr(queue.tail) }); |
| 412 | try expect(mem.eql(u8, buffer[0..fbs.pos], expected)); |
| 413 | } |