Drain buffered inbound connections on cancel to avoid writer trap

Sources/NIOHTTPServer/NIOHTTPServer+HTTP1_1.swift

@@ -13,12 +13,19 @@
 //===----------------------------------------------------------------------===//
 
 import Logging
+
 import NIOCore
+
 import NIOExtras
+
 import NIOHTTP1
+
 import NIOHTTPTypes
+
 import NIOHTTPTypesHTTP1
+
 import NIOPosix
+
 import NIOSSL
 
 @available(anyAppleOS 27.0, *)
@@ -37,31 +44,12 @@ extension NIOHTTPServer {
         connectionStream: NIOAsyncChannelInboundStream<HTTPRequestChannelAndCancellationSignal>,
         connectionHandler: Handler
     ) async throws {
-        // We don't use a `withThrowingDiscardingTaskGroup` here because an error thrown from the body or a child
-        // task would immediately propagate upwards, cancelling all child tasks and bringing down the entire server.
-        // We instead use a non-throwing discarding task group so that errors in the body (e.g. from iterating
-        // `inbound`) must be caught and handled directly.
-        let inboundConnectionIterationError = await withDiscardingTaskGroup { group -> (any Error)? in
-            do {
-                for try await connectionChannel in connectionStream {
-                    group.addTask {
-                        await self.dispatchPlaintextHTTP1_1Connection(
-                            requestChannel: connectionChannel.channel,
-                            clientClosed: connectionChannel.clientClosed,
-                            connectionHandler: connectionHandler
-                        )
-                    }
-                }
-
-                return nil
-            } catch {
-                return error
-            }
-        }
-
-        if let inboundConnectionIterationError {
-            // The error occurred while iterating the inbound connection stream
-            throw inboundConnectionIterationError
+        try await self.serveInboundConnections(connectionStream: connectionStream) { connectionChannel in
+            await self.dispatchPlaintextHTTP1_1Connection(
+                requestChannel: connectionChannel.channel,
+                clientClosed: connectionChannel.clientClosed,
+                connectionHandler: connectionHandler
+            )
         }
     }
 

Sources/NIOHTTPServer/NIOHTTPServer+SecureUpgrade.swift

@@ -57,30 +57,11 @@ extension NIOHTTPServer {
         connectionStream: NIOAsyncChannelInboundStream<EventLoopFuture<NegotiationResult>>,
         connectionHandler: Handler
     ) async throws {
-        // We don't use a `withThrowingDiscardingTaskGroup` here because an error thrown from the body or a child
-        // task would immediately propagate upwards, cancelling all child tasks and bringing down the entire server.
-        // We instead use a non-throwing discarding task group so that errors in the body (e.g. from iterating
-        // `inbound`) must be caught and handled directly.
-        let inboundConnectionIterationError = await withDiscardingTaskGroup { connectionGroup -> (any Error)? in
-            do {
-                for try await upgradeResult in connectionStream {
-                    connectionGroup.addTask {
-                        await self.dispatchSecureConnection(
-                            upgradeResult: upgradeResult,
-                            connectionHandler: connectionHandler
-                        )
-                    }
-                }
-
-                return nil
-            } catch {
-                return error
-            }
-        }
-
-        if let inboundConnectionIterationError {
-            // The error occurred while iterating the inbound connection stream
-            throw inboundConnectionIterationError
+        try await self.serveInboundConnections(connectionStream: connectionStream) { upgradeResult in
+            await self.dispatchSecureConnection(
+                upgradeResult: upgradeResult,
+                connectionHandler: connectionHandler
+            )
         }
     }
 

Sources/NIOHTTPServer/NIOHTTPServer.swift

@@ -664,4 +664,103 @@ extension NIOHTTPServer {
             serverChannel.channel.close(promise: nil)
         }
     }
+
+    /// Coordinates the inbound connection accept loop and cancellation draining.
+    private final class AcceptLoopCoordinator: Sendable {
+        private enum Status {
+            case running
+            case streamFinished
+            case cancelled
+        }
+
+        private struct State {
+            var status: Status = .running
+            var task: Task<Void, any Error>?
+        }
+
+        private let state = NIOLockedValueBox(State())
+
+        func setTask(_ task: Task<Void, any Error>) {
+            self.state.withLockedValue { state in
+                state.task = task
+            }
+        }
+
+        func onStreamFinished() -> Bool {
+            self.state.withLockedValue { state -> Bool in
+                switch state.status {
+                case .running:
+                    state.status = .streamFinished
+                    return false
+                case .cancelled:
+                    return true
+                case .streamFinished:
+                    return false
+                }
+            }
+        }
+
+        func onCancel() {
+            let taskToCancel = self.state.withLockedValue { state -> Task<Void, any Error>? in
+                switch state.status {
+                case .running:
+                    state.status = .cancelled
+                    // Do not cancel the accept task while running; allow it to drain buffered connections.
+                    return nil
+                case .streamFinished:
+                    state.status = .cancelled
+                    // Stream is already finished, safe to cancel the task to abort running connection handlers.
+                    return state.task
+                case .cancelled:
+                    return nil
+                }
+            }
+            taskToCancel?.cancel()
+        }
+    }
+
+    /// Serves inbound connections from `connectionStream` by dispatching each connection to `dispatchConnection`
+    /// in a child task.
+    ///
+    /// When cancelled, the accept loop continues to drain buffered connections that were accepted before the listening
+    /// socket closed, preventing `NIOThrowingAsyncSequenceProducer` from dropping them with un-finished writers.
+    /// Once the stream finishes, all active child tasks are cancelled.
+    func serveInboundConnections<Item: Sendable>(
+        connectionStream: NIOAsyncChannelInboundStream<Item>,
+        dispatchConnection: @escaping @Sendable (Item) async -> Void
+    ) async throws {
+        let coordinator = AcceptLoopCoordinator()
+
+        try await withTaskCancellationHandler {
+            let task = Task {
+                var iterationError: (any Error)?
+                await withDiscardingTaskGroup { group in
+                    do {
+                        for try await item in connectionStream {
+                            group.addTask {
+                                await dispatchConnection(item)
+                            }
+                        }
+                    } catch {
+                        iterationError = error
+                    }
+
+                    let wasCancelled = coordinator.onStreamFinished()
+                    if wasCancelled {
+                        group.cancelAll()
+                    }
+                }
+
+                if let iterationError {
+                    throw iterationError
+                }
+            }
+
+            coordinator.setTask(task)
+            try await task.value
+            try Task.checkCancellation()
+        } onCancel: {
+            coordinator.onCancel()
+        }
+    }
 }

Tests/NIOHTTPServerTests/ConnectionLifecycleTests.swift

@@ -491,6 +491,55 @@ struct ConnectionLifecycleTests {
             #expect(sawOK, "Closure-overloaded handleRequests should serve the request.")
         }
     }
+
+    /// Verifies that cancelling the task running `serve` while accepted connections are still buffered
+    /// in the inbound connection stream does not drop child channels without finishing their writers,
+    /// which would cause a fatal error ("Deinited NIOAsyncWriter without calling finish()").
+    @available(anyAppleOS 27.0, *)
+    @Test(
+        "Cancelling serve does not trap when accepted connections are buffered",
+        arguments: [NIOHTTPServer.HTTPVersion.plaintextHTTP1_1, .http1_1]
+    )
+    func testCancellationWithBufferedAcceptedConnections(httpVersion: NIOHTTPServer.HTTPVersion) async throws {
+        for _ in 0..<5 {
+            let (server, _) = try TestHelpers.makeServerAndClientConfiguration(
+                for: httpVersion,
+                clientLogger: Self.clientLogger,
+                serverLogger: Self.serverLogger
+            )
+
+            try await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    try await server.serve { request, requestContext, reader, responseSender in
+                        var buffer = UniqueArray<UInt8>(copying: [])
+                        try await responseSender.sendAndFinish(.init(status: .ok), buffer: &buffer)
+                    }
+                }
+
+                let address = try #require(try await server.listeningAddresses.first)
+                let port = address.port
+
+                let burst = 32
+                let clientBootstrap = ClientBootstrap(
+                    group: MultiThreadedEventLoopGroup.singletonMultiThreadedEventLoopGroup
+                )
+                var clients = [any Channel]()
+                clients.reserveCapacity(burst)
+
+                for _ in 0..<burst {
+                    if let channel = try? await clientBootstrap.connect(host: address.host, port: port).get() {
+                        clients.append(channel)
+                    }
+                }
+
+                group.cancelAll()
+                for client in clients {
+                    try? await client.close()
+                }
+                _ = try? await group.next()
+            }
+        }
+    }
 }
 
 /// Connection handler that throws on the first invocation and runs normally thereafter, so the test can verify a thrown