+5 -2

readme.md

View file Reviewed

··· 21 21 fork-specific. 22 22 23 23 For intentional `runIo` shutdown, an optional server-handler `serverClose` hook 24 24 - can send the close frame appropriate to the application. The two drain clocks 25 25 - begin together, so their waits use the larger duration, not their sum: 24 24 + can send the close frame appropriate to the application. Hooks fan out 25 25 + concurrently, so a backpressured peer cannot prevent other clients from being 26 26 + notified; a peer still blocked at its WebSocket deadline has both socket 27 27 + directions interrupted. The two drain clocks begin together, so their waits use 28 28 + the larger duration, not their sum: 26 29 27 30 ```zig 28 31 var server = try ws.Server(Handler).init(allocator, io, .{

+197 -5

src/server/server.zig

View file Reviewed

··· 720 720 fn shutdown(self: *Self) void { 721 721 const conn_manager = &self.conn_manager; 722 722 const started = Io.Timestamp.now(conn_manager.io, .awake).toNanoseconds(); 723 723 - conn_manager.beginShutdown(); 723 723 + var notifications: Io.Group = .init; 724 724 + conn_manager.beginShutdownConcurrent(&notifications); 724 725 725 726 self.waitForConnections(started, self.websocket_shutdown_grace, true); 727 727 + conn_manager.interruptWebSockets(self); 726 728 conn_manager.forceWebSockets(self); 727 729 self.waitForConnections(started, self.http_shutdown_grace, false); 728 730 conn_manager.forceAll(self); 731 731 + notifications.await(conn_manager.io) catch {}; 732 732 + } 733 733 + 734 734 + fn interruptShutdown(_: *Self, hc: *HandlerConn(H)) void { 735 735 + const io = hc.conn.io; 736 736 + io.vtable.netShutdown(io.userdata, hc.socket, .both) catch {}; 729 737 } 730 738 731 739 fn waitForConnections(self: *Self, started_ns: i96, grace: Io.Duration, websocket_only: bool) void { ··· 1462 1470 cleanup: Io.Mutex = .init, 1463 1471 compression: ?Compression = null, 1464 1472 shutdown_notified: bool = false, 1473 1473 + // Shutdown must still be able to count and interrupt this socket while 1474 1474 + // its notification task is blocked inside the handler's write path. 1475 1475 + upgraded: std.atomic.Value(bool) = .init(false), 1476 1476 + // A queued notification borrows this HandlerConn. Cleanup waits for 1477 1477 + // the borrow to end before returning the allocation to the pool. 1478 1478 + shutdown_notification_refs: std.atomic.Value(usize) = .init(0), 1465 1479 next: ?*HandlerConn(H) = null, 1466 1480 prev: ?*HandlerConn(H) = null, 1467 1481 ··· 1535 1549 .reader = null, 1536 1550 .compression = null, 1537 1551 .shutdown_notified = false, 1552 1552 + .upgraded = .init(false), 1553 1553 + .shutdown_notification_refs = .init(0), 1538 1554 .conn = .{ 1539 1555 ._closed = false, 1540 1556 .accept_compression = true, ··· 1604 1620 h.close(); 1605 1621 } 1606 1622 hc.handler = null; 1623 1623 + hc.upgraded.store(false, .release); 1607 1624 } 1608 1625 1609 1626 if (hc.conn.compression) |c| { ··· 1623 1640 self.compression_pool.destroy(c); 1624 1641 } 1625 1642 1643 1643 + while (hc.shutdown_notification_refs.load(.acquire) != 0) 1644 1644 + self.io.sleep(.fromMilliseconds(1), .awake) catch {}; 1645 1645 + 1626 1646 self.pool.destroy(hc); 1627 1647 self.lock.unlock(self.io); 1628 1648 } ··· 1665 1685 var total: usize = 0; 1666 1686 var next_node = head; 1667 1687 while (next_node) |hc| : (next_node = hc.next) { 1688 1688 + if (hc.upgraded.load(.acquire)) total += 1; 1689 1689 + } 1690 1690 + return total; 1691 1691 + } 1692 1692 + 1693 1693 + pub fn beginShutdownConcurrent(self: *Self, notifications: *Io.Group) void { 1694 1694 + self.lock.lockUncancelable(self.io); 1695 1695 + defer self.lock.unlock(self.io); 1696 1696 + 1697 1697 + notifyListConcurrent(self.active.head, notifications, self.io); 1698 1698 + notifyListConcurrent(self.pending.head, notifications, self.io); 1699 1699 + } 1700 1700 + 1701 1701 + fn notifyListConcurrent(head: ?*HandlerConn(H), notifications: *Io.Group, io: Io) void { 1702 1702 + var next_node = head; 1703 1703 + while (next_node) |hc| : (next_node = hc.next) { 1668 1704 hc.cleanup.lockUncancelable(hc.conn.io); 1669 1669 - if (hc.handler != null) total += 1; 1705 1705 + const should_notify = hc.handler != null and !hc.shutdown_notified; 1706 1706 + if (should_notify) { 1707 1707 + hc.shutdown_notified = true; 1708 1708 + _ = hc.shutdown_notification_refs.fetchAdd(1, .release); 1709 1709 + } 1670 1710 hc.cleanup.unlock(hc.conn.io); 1711 1711 + 1712 1712 + if (!should_notify) continue; 1713 1713 + notifications.concurrent(io, notifyOne, .{hc}) catch { 1714 1714 + hc.cleanup.lockUncancelable(hc.conn.io); 1715 1715 + hc.shutdown_notified = false; 1716 1716 + hc.cleanup.unlock(hc.conn.io); 1717 1717 + _ = hc.shutdown_notification_refs.fetchSub(1, .release); 1718 1718 + }; 1671 1719 } 1672 1672 - return total; 1720 1720 + } 1721 1721 + 1722 1722 + fn notifyOne(hc: *HandlerConn(H)) void { 1723 1723 + hc.cleanup.lockUncancelable(hc.conn.io); 1724 1724 + if (hc.handler) |*handler| { 1725 1725 + if (comptime std.meta.hasFn(H, "serverClose")) handler.serverClose(); 1726 1726 + } 1727 1727 + hc.cleanup.unlock(hc.conn.io); 1728 1728 + _ = hc.shutdown_notification_refs.fetchSub(1, .release); 1673 1729 } 1674 1730 1675 1731 pub fn beginShutdown(self: *Self) void { ··· 1702 1758 forceList(self.pending.head, worker, true); 1703 1759 } 1704 1760 1761 1761 + pub fn interruptWebSockets(self: *Self, worker: anytype) void { 1762 1762 + self.lock.lockUncancelable(self.io); 1763 1763 + defer self.lock.unlock(self.io); 1764 1764 + 1765 1765 + interruptList(self.active.head, worker); 1766 1766 + interruptList(self.pending.head, worker); 1767 1767 + } 1768 1768 + 1769 1769 + fn interruptList(head: ?*HandlerConn(H), worker: anytype) void { 1770 1770 + var next_node = head; 1771 1771 + while (next_node) |hc| : (next_node = hc.next) { 1772 1772 + if (hc.upgraded.load(.acquire)) worker.interruptShutdown(hc); 1773 1773 + } 1774 1774 + } 1775 1775 + 1705 1776 pub fn forceAll(self: *Self, worker: anytype) void { 1706 1777 self.lock.lockUncancelable(self.io); 1707 1778 defer self.lock.unlock(self.io); ··· 1731 1802 } 1732 1803 h.close(); 1733 1804 hc.handler = null; 1805 1805 + hc.upgraded.store(false, .release); 1734 1806 } 1735 1807 } 1736 1808 hc.cleanup.unlock(hc.conn.io); ··· 2070 2142 }; 2071 2143 2072 2144 hc.handler = handler; 2145 2145 + hc.upgraded.store(true, .release); 2073 2146 2074 2147 var negotiated: ?Handshake.Compression = null; 2075 2148 if (conn.accept_compression) if (handshake.compression) |offer| if (worker.compression) |configured| { ··· 2609 2682 try t.expectEqual(@as(u8, 2), lifecycle.load(.acquire)); 2610 2683 } 2611 2684 2685 2685 + test "Server: shutdown fan-out is concurrent when an earlier writer is wedged" { 2686 2686 + var threaded: std.Io.Threaded = .init(t.allocator, .{}); 2687 2687 + defer threaded.deinit(); 2688 2688 + const io = threaded.io(); 2689 2689 + 2690 2690 + var context: FanoutShutdownContext = .{}; 2691 2691 + var server = try Server(FanoutShutdownHandler).init(t.allocator, io, .{ 2692 2692 + .port = 9299, 2693 2693 + .address = "127.0.0.1", 2694 2694 + .websocket_shutdown_grace = .fromMilliseconds(150), 2695 2695 + .http_shutdown_grace = .fromMilliseconds(150), 2696 2696 + }); 2697 2697 + defer server.deinit(); 2698 2698 + 2699 2699 + var addr = net.IpAddress.parse("127.0.0.1", 9299) catch unreachable; 2700 2700 + var listener = try addr.listen(io, .{ .reuse_address = true }); 2701 2701 + const Run = struct { 2702 2702 + fn go(srv: *Server(FanoutShutdownHandler), l: *net.Server, ctx: *FanoutShutdownContext) void { 2703 2703 + srv.runIo(l, ctx); 2704 2704 + } 2705 2705 + }; 2706 2706 + var future = try io.concurrent(Run.go, .{ &server, &listener, &context }); 2707 2707 + 2708 2708 + // List traversal is connection order. The blocked writer comes first so 2709 2709 + // this test would time out under the old serial notification loop. 2710 2710 + var blocked = try testStreamPortPath(9299, "/blocked"); 2711 2711 + defer blocked.close(); 2712 2712 + var healthy = try testStreamPortPath(9299, "/healthy"); 2713 2713 + defer healthy.close(); 2714 2714 + const receive_timeout = std.mem.toBytes(posix.timeval{ .sec = 1, .usec = 0 }); 2715 2715 + try posix.setsockopt(healthy.socket, posix.SOL.SOCKET, posix.SO.RCVTIMEO, &receive_timeout); 2716 2716 + 2717 2717 + var attempts: usize = 0; 2718 2718 + while ((!context.blocked_ready.load(.acquire) or !context.healthy_ready.load(.acquire)) and attempts < 1000) : (attempts += 1) 2719 2719 + _ = try io.sleep(.fromMilliseconds(1), .awake); 2720 2720 + try std.testing.expect(context.blocked_ready.load(.acquire)); 2721 2721 + try std.testing.expect(context.healthy_ready.load(.acquire)); 2722 2722 + 2723 2723 + const Stop = struct { 2724 2724 + fn go(server_future: *Io.Future(void), task_io: Io) void { 2725 2725 + _ = server_future.cancel(task_io); 2726 2726 + } 2727 2727 + }; 2728 2728 + const started = Io.Timestamp.now(io, .awake).toNanoseconds(); 2729 2729 + var stopping = try io.concurrent(Stop.go, .{ &future, io }); 2730 2730 + 2731 2731 + var close_frame: [4]u8 = undefined; 2732 2732 + _ = try healthy.readAtLeast(&close_frame, close_frame.len); 2733 2733 + const healthy_elapsed = Io.Timestamp.now(io, .awake).toNanoseconds() - started; 2734 2734 + try t.expectSlice(u8, &.{ 0x88, 0x02, 0x03, 0xe9 }, &close_frame); 2735 2735 + try std.testing.expect(healthy_elapsed < 100 * std.time.ns_per_ms); 2736 2736 + try healthy.writeAll(&proto.frame(.close, "\x03\xe9")); 2737 2737 + 2738 2738 + stopping.await(io); 2739 2739 + listener.deinit(io); 2740 2740 + const total_elapsed = Io.Timestamp.now(io, .awake).toNanoseconds() - started; 2741 2741 + try std.testing.expect(context.blocked_close_started.load(.acquire)); 2742 2742 + try std.testing.expect(context.blocked_close_unblocked.load(.acquire)); 2743 2743 + try std.testing.expect(context.healthy_notified.load(.acquire)); 2744 2744 + try std.testing.expect(total_elapsed >= 110 * std.time.ns_per_ms); 2745 2745 + try std.testing.expect(total_elapsed < 400 * std.time.ns_per_ms); 2746 2746 + } 2747 2747 + 2612 2748 test "Server: HTTP shutdown grace forces an active fallback at its own deadline" { 2613 2749 var threaded: std.Io.Threaded = .init(t.allocator, .{}); 2614 2750 defer threaded.deinit(); ··· 2971 3107 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.RCVTIMEO, &timeout); 2972 3108 try posix.setsockopt(socket, posix.SOL.SOCKET, posix.SO.SNDTIMEO, &timeout); 2973 3109 2974 2974 - var result = TestStream{ .socket = socket, .io = io }; 3110 3110 + const result = TestStream{ .socket = socket, .io = io }; 2975 3111 2976 3112 if (handshake == false) { 2977 3113 return result; 2978 3114 } 2979 3115 2980 2980 - try result.writeAll("GET / HTTP/1.1\r\ncontent-length: 0\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nconnection: upgrade\r\nsec-websocket-key: my-key\r\n\r\n"); 3116 3116 + return finishTestHandshake(result, "/"); 3117 3117 + } 3118 3118 + 3119 3119 + fn testStreamPortPath(port: u16, path: []const u8) !TestStream { 3120 3120 + const result = try testStreamPort(false, port); 3121 3121 + return finishTestHandshake(result, path); 3122 3122 + } 3123 3123 + 3124 3124 + fn finishTestHandshake(result: TestStream, path: []const u8) !TestStream { 3125 3125 + try result.writeAll("GET "); 3126 3126 + try result.writeAll(path); 3127 3127 + try result.writeAll(" HTTP/1.1\r\ncontent-length: 0\r\nupgrade: websocket\r\nsec-websocket-version: 13\r\nconnection: upgrade\r\nsec-websocket-key: my-key\r\n\r\n"); 2981 3128 var buf: [1024]u8 = undefined; 2982 3129 var pos: usize = 0; 2983 3130 while (pos < buf.len) { ··· 3107 3254 } 3108 3255 3109 3256 pub fn close(_: *LifecycleHandler) void {} 3257 3257 + }; 3258 3258 + 3259 3259 + const FanoutShutdownContext = struct { 3260 3260 + blocked_ready: std.atomic.Value(bool) = .init(false), 3261 3261 + healthy_ready: std.atomic.Value(bool) = .init(false), 3262 3262 + blocked_close_started: std.atomic.Value(bool) = .init(false), 3263 3263 + blocked_close_unblocked: std.atomic.Value(bool) = .init(false), 3264 3264 + healthy_notified: std.atomic.Value(bool) = .init(false), 3265 3265 + }; 3266 3266 + 3267 3267 + const FanoutShutdownHandler = struct { 3268 3268 + context: *FanoutShutdownContext, 3269 3269 + conn: *Conn, 3270 3270 + blocked: bool, 3271 3271 + 3272 3272 + pub fn init(handshake: *const Handshake, conn: *Conn, context: *FanoutShutdownContext) !FanoutShutdownHandler { 3273 3273 + const blocked = std.mem.eql(u8, handshake.url, "/blocked"); 3274 3274 + if (blocked) { 3275 3275 + const send_buffer = std.mem.toBytes(@as(c_int, 4096)); 3276 3276 + setSockOptBestEffort(conn.stream.socket.handle, posix.SOL.SOCKET, posix.SO.SNDBUF, &send_buffer); 3277 3277 + try conn.writeTimeout(5000); 3278 3278 + context.blocked_ready.store(true, .release); 3279 3279 + } else { 3280 3280 + context.healthy_ready.store(true, .release); 3281 3281 + } 3282 3282 + return .{ .context = context, .conn = conn, .blocked = blocked }; 3283 3283 + } 3284 3284 + 3285 3285 + pub fn clientMessage(_: *FanoutShutdownHandler, _: []const u8) !void {} 3286 3286 + 3287 3287 + pub fn serverClose(self: *FanoutShutdownHandler) void { 3288 3288 + if (self.blocked) { 3289 3289 + self.context.blocked_close_started.store(true, .release); 3290 3290 + const payload = [_]u8{'x'} ** (16 * 1024); 3291 3291 + while (true) self.conn.writeBin(&payload) catch { 3292 3292 + self.context.blocked_close_unblocked.store(true, .release); 3293 3293 + return; 3294 3294 + }; 3295 3295 + } 3296 3296 + 3297 3297 + self.conn.writeClose(.{ .code = 1001 }) catch return; 3298 3298 + self.context.healthy_notified.store(true, .release); 3299 3299 + } 3300 3300 + 3301 3301 + pub fn close(_: *FanoutShutdownHandler) void {} 3110 3302 }; 3111 3303 3112 3304 const SlowHttpContext = struct {