Fix race condition when websocket connection is upgraded on MacOS (maybe BSD) · poecoh.com/http.zig@5d160ba

+70 -1

src/httpz.zig

··· 800 800 { 801 801 default_server = try Server(void).init(ga, .{ 802 802 .address = .localhost(5992), 803 - .request = .{ .lazy_read_size = 4_096, .max_body_size = 1_048_576, }, 803 + .request = .{ 804 + .lazy_read_size = 4_096, 805 + .max_body_size = 1_048_576, 806 + }, 804 807 }, {}); 805 808 806 809 // only need to do this because we're using listenInNewThread instead ··· 1937 1940 try t.expectEqual(10, buf[1]); 1938 1941 try t.expectString("over 9000!", buf[2..12]); 1939 1942 try t.expectString(&.{ 136, 2, 3, 232 }, buf[12..16]); 1943 + } 1944 + 1945 + // Stress test: multiple concurrent websocket clients sending many messages each. 1946 + // Run repeatedly (e.g. zig build test -Dtest-filter="websocket: stress" or run 50x) 1947 + // to verify no race in reader.done() / allocator free when using thread pool. 1948 + test "websocket: stress" { 1949 + if (force_blocking) return; // non-blocking mode only (thread pool) 1950 + const num_clients = 8; 1951 + const messages_per_client = 150; 1952 + 1953 + // When run with -Dtest-filter="websocket: stress", tests:beforeAll may not run, 1954 + // so nothing is listening on 5998. Wait for port and start our own server if needed. 1955 + var stress_server: ?Server(TestWebsocketHandler) = null; 1956 + var stress_listen_thread: ?Thread = null; 1957 + testing.waitForPort(5998) catch { 1958 + stress_server = try Server(TestWebsocketHandler).init(t.allocator, .{ .address = .localhost(5998) }, TestWebsocketHandler{}); 1959 + var router = try stress_server.?.router(.{}); 1960 + router.get("/ws", TestWebsocketHandler.upgrade, .{}); 1961 + stress_listen_thread = try stress_server.?.listenInNewThread(); 1962 + try testing.waitForPort(5998); 1963 + }; 1964 + defer if (stress_server) |*srv| { 1965 + srv.stop(); 1966 + if (stress_listen_thread) |thrd| thrd.join(); 1967 + srv.deinit(); 1968 + }; 1969 + 1970 + var threads: [num_clients]Thread = undefined; 1971 + for (0..num_clients) |i| { 1972 + threads[i] = Thread.spawn(.{}, struct { 1973 + fn run(_: usize) void { 1974 + const stream = testStream(5998); 1975 + defer stream.close(); 1976 + 1977 + var writer = stream.writer(&.{}); 1978 + const w = &writer.interface; 1979 + w.writeAll("GET /ws HTTP/1.1\r\nContent-Length: 0\r\n") catch return; 1980 + w.writeAll("upgrade: WEBsocket\r\n") catch return; 1981 + w.writeAll("Sec-Websocket-verSIon: 13\r\n") catch return; 1982 + w.writeAll("ConnectioN: upgrade\r\n") catch return; 1983 + w.writeAll("SEC-WEBSOCKET-KeY: a-secret-key\r\n\r\n") catch return; 1984 + w.flush() catch return; 1985 + 1986 + var buf: [1024]u8 = undefined; 1987 + var pos: usize = 0; 1988 + var reader = stream.reader(&.{}); 1989 + const r = reader.interface(); 1990 + while (!std.mem.endsWith(u8, buf[0..pos], "\r\n\r\n")) { 1991 + if (pos >= buf.len) return; 1992 + var vecs: [1][]u8 = .{buf[pos..]}; 1993 + const n = r.readVec(&vecs) catch return; 1994 + if (n == 0) return; 1995 + pos += n; 1996 + } 1997 + if (pos < 12 or !std.mem.startsWith(u8, buf[0..12], "HTTP/1.1 101")) return; 1998 + 1999 + for (0..messages_per_client) |_| { 2000 + const frame = websocket.frameText("stress"); 2001 + w.writeAll(&frame) catch return; 2002 + } 2003 + w.writeAll(&websocket.frameText("close")) catch return; 2004 + w.flush() catch return; 2005 + } 2006 + }.run, .{i}) catch return; 2007 + } 2008 + for (&threads) |*th| th.join(); 1940 2009 } 1941 2010 1942 2011 test "ContentType: forX" {

+3 -3

src/key_value.zig

··· 25 25 const allocation = try allocator.alignedAlloc(u8, std.mem.Alignment.fromByteUnits(alignment), max * size); 26 26 return .{ 27 27 .len = 0, 28 - .keys = @as([*][]const u8, @alignCast(@ptrCast(if (kFirst) allocation.ptr else allocation[max * @sizeOf(V) ..].ptr)))[0..max], 29 - .values = @as([*]V, @alignCast(@ptrCast(if (kFirst) allocation[max * @sizeOf([]const u8) ..].ptr else allocation.ptr)))[0..max], 28 + .keys = @as([*][]const u8, @ptrCast(@alignCast(if (kFirst) allocation.ptr else allocation[max * @sizeOf(V) ..].ptr)))[0..max], 29 + .values = @as([*]V, @ptrCast(@alignCast(if (kFirst) allocation[max * @sizeOf([]const u8) ..].ptr else allocation.ptr)))[0..max], 30 30 .hashes = allocation[max * @sizeOf([]const u8) + max * @sizeOf(V) ..], 31 31 }; 32 32 } 33 33 34 34 pub fn deinit(self: *Self, allocator: Allocator) void { 35 - const allocation = @as([*]align(alignment) u8, @alignCast(@ptrCast(if (kFirst) self.keys.ptr else self.values.ptr))); 35 + const allocation = @as([*]align(alignment) u8, @ptrCast(@alignCast(if (kFirst) self.keys.ptr else self.values.ptr))); 36 36 allocator.free(allocation[0 .. self.keys.len * size]); 37 37 } 38 38

+2 -2

src/response.zig

··· 161 161 buf[len] = '\r'; 162 162 buf[len + 1] = '\n'; 163 163 164 - var vec = [2][]const u8{buf[0..len+2], data}; 164 + var vec = [2][]const u8{ buf[0 .. len + 2], data }; 165 165 try conn.writeAllIOVec(&vec); 166 166 } 167 167 ··· 197 197 const buffered = self.buffer.writer.buffered(); 198 198 const body = if (buffered.len > 0) buffered else self.body; 199 199 200 - var vec = [2][]const u8{header_buf, body}; 200 + var vec = [2][]const u8{ header_buf, body }; 201 201 return conn.writeAllIOVec(&vec); 202 202 } 203 203

+16

src/testing.zig

··· 287 287 }; 288 288 } 289 289 290 + /// Waits until a TCP port is accepting connections (e.g. after starting a server in another thread). 291 + /// Tries up to 100 times with 20ms sleep between attempts. 292 + pub fn waitForPort(port: u16) !void { 293 + const address = std.net.Address.parseIp("127.0.0.1", port) catch unreachable; 294 + for (0..100) |_| { 295 + if (std.net.tcpConnectToAddress(address)) |stream| { 296 + stream.close(); 297 + return; 298 + } else |err| { 299 + if (err != error.ConnectionRefused) return err; 300 + std.Thread.sleep(20 * std.time.ns_per_ms); 301 + } 302 + } 303 + return error.ConnectionRefused; 304 + } 305 + 290 306 fn decodeChunkedEncoding(full_dest: []u8, full_src: []u8) usize { 291 307 var src = full_src; 292 308 var dest = full_dest;

+40 -4

src/worker.zig

··· 563 563 std.Thread.sleep(std.time.ns_per_s); 564 564 continue; 565 565 }; 566 + 566 567 now = timestamp(now); 567 568 var closed_conn = false; 568 569 ··· 588 589 589 590 // At this point, the connection is either in 590 591 // keepalive or request. Either way, we know no 591 - // other thread is access the connection, so we 592 + // other thread is accessing the connection, so we 592 593 // can access _state directly. 593 594 594 595 const stream = http_conn.stream; ··· 613 614 self.swapList(conn, .active); 614 615 thread_pool.spawn(.{ self, now, conn }); 615 616 }, 616 - .websocket => thread_pool.spawn(.{ self, now, conn }), 617 + .websocket => { 618 + if (conn.acquireProcessing() == false) { 619 + // Connection is already being processed. We need 620 + // to wait for the current processing to complete. 621 + // See the processing field in Conn 622 + continue; 623 + } 624 + thread_pool.spawn(.{ self, now, conn }); 625 + }, 617 626 }, 618 627 .shutdown => return, 619 628 } ··· 806 815 var handover = http_conn.handover; 807 816 http_conn.requestDone(self.retain_allocated_bytes, handover == .keepalive or handover == .websocket) catch { 808 817 // This means we failed to put the connection into 809 - // nonblocking mode. Rare, but safer to clos the connection 818 + // nonblocking mode. Rare, but safer to close the connection 810 819 // at this point. 811 820 handover = .close; 812 821 }; ··· 835 844 } 836 845 837 846 pub fn processWebsocketData(self: *Self, conn: *Conn(WSH), thread_buf: []u8, hc: *ws.HandlerConn(WSH)) void { 847 + defer conn.releaseProcessing(); 848 + 838 849 var ws_conn = &hc.conn; 839 850 const success = self.websocket.worker.dataAvailable(hc, thread_buf); 840 851 if (success == false) { ··· 1090 1101 _ = try posix.kevent(self.fd, &.{.{ 1091 1102 .ident = @intCast(conn.getSocket()), 1092 1103 .filter = posix.system.EVFILT.READ, 1093 - .flags = posix.system.EV.ENABLE, 1104 + .flags = posix.system.EV.ADD, 1094 1105 .fflags = 0, 1095 1106 .data = 0, 1096 1107 .udata = @intFromPtr(conn), ··· 1458 1469 next: ?*Conn(WSH), 1459 1470 prev: ?*Conn(WSH), 1460 1471 1472 + // Used exclusively in the websocket path. In the HTTP Path, the 1473 + // HTTPConn's state is used to prevent 2 requests on the same socket 1474 + // from being processed at the same time. This is 100% necessary because 1475 + // our event loop will keep emitting events for new data. For HTTP, this 1476 + // is OK, because a "health" flow is REQ->RES - there shouldn't be 1477 + // multiple inflight requests. NOT using ONESHOT/EV_DISPATCH is more 1478 + // efficient because it avoids a bunch of system calls. 1479 + // But for WebSockets, it's normal to have multiple incoming messages. 1480 + // So we switch the event notifier to be ONESHOT/DISPATCH. And, that 1481 + // SHOULD be enough to stop 2 concurrent messages from being processed 1482 + // at the same time. And on Linux, it seems to work fine. But on MacOS, 1483 + // there's a window where we might get 2 messages across 2 different 1484 + // wait calls as we're switching to DISPATCH. So this guard is there just 1485 + // for thar narrow window. 1486 + processing: bool = false, 1487 + 1461 1488 const Self = @This(); 1462 1489 1463 1490 fn close(self: *Self) void { ··· 1472 1499 .http => |hc| hc.stream.handle, 1473 1500 .websocket => |hc| hc.socket, 1474 1501 }; 1502 + } 1503 + 1504 + pub fn acquireProcessing(self: *Self) bool { 1505 + // returns true if it was previously false 1506 + return @atomicRmw(bool, &self.processing, .Xchg, true, .monotonic) == false; 1507 + } 1508 + 1509 + pub fn releaseProcessing(self: *Self) void { 1510 + return @atomicStore(bool, &self.processing, false, .release); 1475 1511 } 1476 1512 }; 1477 1513 }

Configure Feed

Configure Feed