diff --git a/Sources/NIOHTTPServer/NIOHTTPServer+HTTP1_1.swift b/Sources/NIOHTTPServer/NIOHTTPServer+HTTP1_1.swift index 192c831..5596d69 100644 --- a/Sources/NIOHTTPServer/NIOHTTPServer+HTTP1_1.swift +++ b/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, 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 + ) } } diff --git a/Sources/NIOHTTPServer/NIOHTTPServer+SecureUpgrade.swift b/Sources/NIOHTTPServer/NIOHTTPServer+SecureUpgrade.swift index 95fbca7..5eabecc 100644 --- a/Sources/NIOHTTPServer/NIOHTTPServer+SecureUpgrade.swift +++ b/Sources/NIOHTTPServer/NIOHTTPServer+SecureUpgrade.swift @@ -57,30 +57,11 @@ extension NIOHTTPServer { connectionStream: NIOAsyncChannelInboundStream>, 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 + ) } } diff --git a/Sources/NIOHTTPServer/NIOHTTPServer.swift b/Sources/NIOHTTPServer/NIOHTTPServer.swift index 40eb530..97de94f 100644 --- a/Sources/NIOHTTPServer/NIOHTTPServer.swift +++ b/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? + } + + private let state = NIOLockedValueBox(State()) + + func setTask(_ task: Task) { + 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? 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( + connectionStream: NIOAsyncChannelInboundStream, + 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() + } + } } diff --git a/Tests/NIOHTTPServerTests/ConnectionLifecycleTests.swift b/Tests/NIOHTTPServerTests/ConnectionLifecycleTests.swift index a671754..3cccbb0 100644 --- a/Tests/NIOHTTPServerTests/ConnectionLifecycleTests.swift +++ b/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(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..