From a9f341ea45d89e50190cb9da38040ff491b027ff Mon Sep 17 00:00:00 2001 From: Kazuki Nakashima <65545348+lynnswap@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:37:25 +0900 Subject: [PATCH] Own MCP connection request shutdown --- .../CodexReviewMCPHTTPServer.swift | 106 ++- .../MCPHTTPNetworkResourceOwner.swift | 630 ++++++++++++++++++ .../CodexReviewMCPHTTPServerTests.swift | 35 + .../MCPHTTPNetworkResourceOwnerTests.swift | 288 ++++++++ 4 files changed, 1036 insertions(+), 23 deletions(-) create mode 100644 Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift create mode 100644 Tests/CodexReviewMCPServerTests/MCPHTTPNetworkResourceOwnerTests.swift diff --git a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift index 2109606..cadfbdf 100644 --- a/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift +++ b/Sources/CodexReviewMCPServer/CodexReviewMCPHTTPServer.swift @@ -208,11 +208,17 @@ package actor CodexReviewMCPHTTPServer { private final class StartingGeneration: @unchecked Sendable { let id: UInt64 + let networkResources: MCPHTTPNetworkResourceOwner let task: Task var admissionClosed = false - init(id: UInt64, task: Task) { + init( + id: UInt64, + networkResources: MCPHTTPNetworkResourceOwner, + task: Task + ) { self.id = id + self.networkResources = networkResources self.task = task } } @@ -223,6 +229,7 @@ package actor CodexReviewMCPHTTPServer { let eventLoopGroup: MultiThreadedEventLoopGroup let cleanupTask: Task let boundURL: URL + let networkResources: MCPHTTPNetworkResourceOwner var admissionClosed = false var listenerCloseTask: Task? @@ -231,13 +238,15 @@ package actor CodexReviewMCPHTTPServer { listener: any Channel, eventLoopGroup: MultiThreadedEventLoopGroup, cleanupTask: Task, - boundURL: URL + boundURL: URL, + networkResources: MCPHTTPNetworkResourceOwner ) { self.id = id self.listener = listener self.eventLoopGroup = eventLoopGroup self.cleanupTask = cleanupTask self.boundURL = boundURL + self.networkResources = networkResources } } @@ -347,10 +356,18 @@ package actor CodexReviewMCPHTTPServer { try result.get() nextGenerationID &+= 1 let id = nextGenerationID + let networkResources = MCPHTTPNetworkResourceOwner(generationID: id) let task = Task { [self] in - await performStartGeneration(id: id) + await performStartGeneration( + id: id, + networkResources: networkResources + ) } - let operation = StartingGeneration(id: id, task: task) + let operation = StartingGeneration( + id: id, + networkResources: networkResources, + task: task + ) lifecycleState = .starting(operation) let result = await task.value try publishStartResult(result, operation: operation) @@ -378,14 +395,23 @@ package actor CodexReviewMCPHTTPServer { } } - private func performStartGeneration(id: UInt64) async -> StartingGenerationResult { + private func performStartGeneration( + id: UInt64, + networkResources: MCPHTTPNetworkResourceOwner + ) async -> StartingGenerationResult { let group = MultiThreadedEventLoopGroup(numberOfThreads: 1) let bootstrap = ServerBootstrap(group: group) .serverChannelOption(ChannelOptions.backlog, value: 128) .serverChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) .childChannelInitializer { channel in - channel.pipeline.configureHTTPServerPipeline().flatMap { - channel.pipeline.addHandler(CodexReviewMCPHTTPHandler(server: self)) + guard let connection = networkResources.admitConnection(channel) else { + return channel.close(mode: .all) + } + return channel.pipeline.configureHTTPServerPipeline().flatMap { + channel.pipeline.addHandler(CodexReviewMCPHTTPHandler( + server: self, + connection: connection + )) } } .childChannelOption(ChannelOptions.socketOption(.so_reuseaddr), value: 1) @@ -411,10 +437,13 @@ package actor CodexReviewMCPHTTPServer { listener: channel, eventLoopGroup: group, cleanupTask: cleanupTask, - boundURL: configuration.url(boundPort: actualPort) + boundURL: configuration.url(boundPort: actualPort), + networkResources: networkResources )) } catch { var failures: [LifecycleError.Failure] = [] + networkResources.closeAdmission() + let closingNetworkResources = networkResources.beginClosing(.serverStop) if let listener { do { try await listener.close() @@ -428,6 +457,7 @@ package actor CodexReviewMCPHTTPServer { failures.append(listenerFailure) } } + await closingNetworkResources.waitUntilClosed() do { try await group.shutdownGracefully() } catch { @@ -496,11 +526,14 @@ package actor CodexReviewMCPHTTPServer { case .running(let resources): resources.admissionClosed = true + resources.networkResources.closeAdmission() + let closingNetworkResources = resources.networkResources.beginClosing(.serverStop) id = resources.id let groupFailure = consumeEventLoopGroupShutdownFailureForTesting() let newTask = Task { [self] in await performStopGeneration( resources, + closingNetworkResources: closingNetworkResources, injectedGroupFailure: groupFailure ) } @@ -512,7 +545,7 @@ package actor CodexReviewMCPHTTPServer { task = newTask case .starting(let operation): - closeStartingAdmission(operation) + let closingNetworkResources = closeStartingAdmission(operation) id = operation.id let newTask = Task { [self] in switch await operation.task.value { @@ -520,6 +553,7 @@ package actor CodexReviewMCPHTTPServer { resources.admissionClosed = true return await performStopGeneration( resources, + closingNetworkResources: closingNetworkResources, injectedGroupFailure: consumeEventLoopGroupShutdownFailureForTesting() ) case .failure(let failure): @@ -540,7 +574,7 @@ package actor CodexReviewMCPHTTPServer { case .stopped: return case .starting(let operation): - closeStartingAdmission(operation) + let closingNetworkResources = closeStartingAdmission(operation) let id = operation.id let task = Task { [self] in switch await operation.task.value { @@ -548,6 +582,7 @@ package actor CodexReviewMCPHTTPServer { resources.admissionClosed = true return await performStopGeneration( resources, + closingNetworkResources: closingNetworkResources, injectedGroupFailure: consumeEventLoopGroupShutdownFailureForTesting() ) case .failure(let failure): @@ -559,6 +594,7 @@ package actor CodexReviewMCPHTTPServer { finishStopIfCurrent(id: id, result: result) case .running(let resources): resources.admissionClosed = true + resources.networkResources.closeAdmission() _ = await listenerCloseTask(for: resources).value case .stopping(let id, _, let task): let result = await task.value @@ -566,22 +602,26 @@ package actor CodexReviewMCPHTTPServer { } } - private func closeStartingAdmission(_ operation: StartingGeneration) { - guard operation.admissionClosed == false else { - return - } - operation.admissionClosed = true - operation.task.cancel() - lastStartingAdmissionClosedGenerationID = operation.id - let waiters = startingAdmissionCloseWaiters - startingAdmissionCloseWaiters.removeAll(keepingCapacity: false) - for waiter in waiters { - waiter.resume() + private func closeStartingAdmission( + _ operation: StartingGeneration + ) -> MCPHTTPNetworkResourceOwner.ClosingGeneration { + if operation.admissionClosed == false { + operation.admissionClosed = true + operation.networkResources.closeAdmission() + operation.task.cancel() + lastStartingAdmissionClosedGenerationID = operation.id + let waiters = startingAdmissionCloseWaiters + startingAdmissionCloseWaiters.removeAll(keepingCapacity: false) + for waiter in waiters { + waiter.resume() + } } + return operation.networkResources.beginClosing(.serverStop) } private func performStopGeneration( _ resources: RunningGeneration, + closingNetworkResources: MCPHTTPNetworkResourceOwner.ClosingGeneration, injectedGroupFailure: LifecycleError.Failure? ) async -> StoppingGenerationResult { let listenerCloseTask = listenerCloseTask(for: resources) @@ -589,6 +629,7 @@ package actor CodexReviewMCPHTTPServer { await closeAllSessions() let listenerResult = await listenerCloseTask.value await resources.cleanupTask.value + await closingNetworkResources.waitUntilClosed() await stopCompletionGate.waitIfNeeded() var failures: [LifecycleError.Failure] = [] @@ -1171,13 +1212,18 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } private let server: CodexReviewMCPHTTPServer + private let connection: MCPHTTPNetworkResourceOwner.Connection private var requestState: RequestState? private var activeStreamTask: Task? private var activeStreamID: UUID? private var activeStreamCompletion: ActiveRequestCompletion? - init(server: CodexReviewMCPHTTPServer) { + init( + server: CodexReviewMCPHTTPServer, + connection: MCPHTTPNetworkResourceOwner.Connection + ) { self.server = server + self.connection = connection } func channelRead(context: ChannelHandlerContext, data: NIOAny) { @@ -1195,10 +1241,21 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked return } requestState = nil + guard let admittedRequest = connection.admitRequest() else { + context.close(promise: nil) + return + } nonisolated(unsafe) let context = context - Task { + let task = Task { [self] in + defer { + admittedRequest.lease.acknowledgeCompletion() + } + guard await admittedRequest.lease.waitUntilStartIsAllowed() else { + return + } await handleRequest(state: state, context: context) } + admittedRequest.lease.install(task) } } @@ -1208,12 +1265,14 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func channelInactive(context: ChannelHandlerContext) { + connection.peerClosed() finishActiveStream() context.fireChannelInactive() } func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) { if case ChannelEvent.inputClosed = event { + connection.peerClosed() finishActiveStream() context.close(promise: nil) return @@ -1222,6 +1281,7 @@ private final class CodexReviewMCPHTTPHandler: ChannelInboundHandler, @unchecked } func errorCaught(context: ChannelHandlerContext, error: any Error) { + connection.transportFailed(error.localizedDescription) finishActiveStream() context.close(promise: nil) } diff --git a/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift new file mode 100644 index 0000000..25a9e92 --- /dev/null +++ b/Sources/CodexReviewMCPServer/MCPHTTPNetworkResourceOwner.swift @@ -0,0 +1,630 @@ +import Foundation +@preconcurrency import NIOCore + +package protocol MCPHTTPConnectionResource: AnyObject, Sendable { + func signalClose() + func installCloseAcknowledgement(_ acknowledgement: @escaping @Sendable () -> Void) +} + +private final class NIOHTTPConnectionResource: MCPHTTPConnectionResource, @unchecked Sendable { + private let channel: any Channel + + init(channel: any Channel) { + self.channel = channel + } + + func signalClose() { + channel.close(mode: .all, promise: nil) + } + + func installCloseAcknowledgement(_ acknowledgement: @escaping @Sendable () -> Void) { + channel.closeFuture.whenComplete { _ in + acknowledgement() + } + } +} + +package final class MCPHTTPNetworkResourceOwner: @unchecked Sendable { + package enum TerminalCause: Equatable, Sendable { + case serverStop + case peerClosed + case transportFailure(String) + } + + package enum GenerationPhase: Equatable, Sendable { + case accepting + case admissionClosed + case closing(TerminalCause) + case closed + } + + package enum ConnectionPhase: Equatable, Sendable { + case accepting + case admissionClosed + case closing(TerminalCause) + case closed + } + + package enum RequestWorkPhase: Equatable, Sendable { + case reserved + case installed + case running + case closing(TerminalCause) + case closed(TerminalCause?) + } + + package struct RequestSnapshot: Equatable, Sendable { + package let id: UUID + package let admissionOrdinal: UInt64 + package let phase: RequestWorkPhase + } + + package struct ConnectionSnapshot: Equatable, Sendable { + package let id: UUID + package let admissionOrdinal: UInt64 + package let phase: ConnectionPhase + package let closeAcknowledged: Bool + package let requests: [RequestSnapshot] + } + + package struct Snapshot: Equatable, Sendable { + package let generationID: UInt64 + package let phase: GenerationPhase + package let connections: [ConnectionSnapshot] + + package var isClosed: Bool { + phase == .closed && connections.isEmpty + } + } + + package final class WorkLease: @unchecked Sendable { + fileprivate let id: UUID + private weak var operation: RequestOperation? + + fileprivate init(operation: RequestOperation, id: UUID) { + self.id = id + self.operation = operation + } + + package func install(_ task: Task) { + operation?.install(task: task, leaseID: id) + } + + package func waitUntilStartIsAllowed() async -> Bool { + guard let operation else { + return false + } + return await operation.waitUntilStartIsAllowed(leaseID: id) + } + + package func acknowledgeCompletion() { + operation?.acknowledgeCompletion(leaseID: id) + } + } + + package final class RequestOperation: @unchecked Sendable { + private enum WorkState { + case reserved + case installed(@Sendable () -> Void) + case running(@Sendable () -> Void) + case acknowledged + } + + package let id = UUID() + package let admissionOrdinal: UInt64 + private weak var connection: Connection? + private let lock = NSLock() + private let leaseID: UUID + private var workState: WorkState = .reserved + private var terminalCause: TerminalCause? + private var startWasRequested = false + private var startWaiter: CheckedContinuation? + private var closeWaiters: [CheckedContinuation] = [] + + fileprivate init(admissionOrdinal: UInt64, connection: Connection) { + self.admissionOrdinal = admissionOrdinal + self.connection = connection + let leaseID = UUID() + self.leaseID = leaseID + } + + fileprivate func makeLease() -> WorkLease { + WorkLease(operation: self, id: leaseID) + } + + fileprivate func beginClosing(_ cause: TerminalCause) { + var cancellation: (@Sendable () -> Void)? + var waiter: CheckedContinuation? + lock.lock() + if terminalCause == nil { + terminalCause = cause + } + switch workState { + case .installed(let cancel), .running(let cancel): + cancellation = cancel + waiter = startWaiter + startWaiter = nil + case .reserved, .acknowledged: + break + } + lock.unlock() + cancellation?() + waiter?.resume(returning: false) + } + + package func waitUntilClosed() async -> TerminalCause? { + await withCheckedContinuation { continuation in + lock.lock() + if case .acknowledged = workState { + let cause = terminalCause + lock.unlock() + continuation.resume(returning: cause) + } else { + closeWaiters.append(continuation) + lock.unlock() + } + } + } + + fileprivate func install(task: Task, leaseID: UUID) { + var shouldCancel = false + var waiter: CheckedContinuation? + lock.lock() + precondition(self.leaseID == leaseID, "A request WorkLease belongs to exactly one admitted operation.") + guard case .reserved = workState else { + lock.unlock() + preconditionFailure("A request WorkLease can be installed exactly once.") + } + let cancellation: @Sendable () -> Void = { task.cancel() } + if startWasRequested { + waiter = startWaiter + startWaiter = nil + if terminalCause == nil { + workState = .running(cancellation) + } else { + workState = .installed(cancellation) + shouldCancel = true + } + } else { + workState = .installed(cancellation) + shouldCancel = terminalCause != nil + } + lock.unlock() + if shouldCancel { + task.cancel() + } + waiter?.resume(returning: shouldCancel == false) + } + + fileprivate func waitUntilStartIsAllowed(leaseID: UUID) async -> Bool { + await withCheckedContinuation { continuation in + var cancellation: (@Sendable () -> Void)? + lock.lock() + precondition(self.leaseID == leaseID, "A request WorkLease belongs to exactly one admitted operation.") + precondition(startWasRequested == false, "A request WorkLease can start exactly once.") + startWasRequested = true + switch workState { + case .reserved: + startWaiter = continuation + lock.unlock() + case .installed(let cancel): + if terminalCause == nil { + workState = .running(cancel) + lock.unlock() + continuation.resume(returning: true) + } else { + cancellation = cancel + lock.unlock() + cancellation?() + continuation.resume(returning: false) + } + case .running, .acknowledged: + lock.unlock() + preconditionFailure("A request WorkLease can start exactly once.") + } + } + } + + fileprivate func acknowledgeCompletion(leaseID: UUID) { + let waiters: [CheckedContinuation] + let cause: TerminalCause? + lock.lock() + precondition(self.leaseID == leaseID, "A request WorkLease belongs to exactly one admitted operation.") + guard case .acknowledged = workState else { + workState = .acknowledged + cause = terminalCause + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in waiters { + waiter.resume(returning: cause) + } + connection?.requestDidClose(self) + return + } + lock.unlock() + } + + fileprivate func snapshot() -> RequestSnapshot { + lock.lock() + let phase: RequestWorkPhase + switch workState { + case .reserved: + phase = terminalCause.map(RequestWorkPhase.closing) ?? .reserved + case .installed: + phase = terminalCause.map(RequestWorkPhase.closing) ?? .installed + case .running: + phase = terminalCause.map(RequestWorkPhase.closing) ?? .running + case .acknowledged: + phase = .closed(terminalCause) + } + lock.unlock() + return .init(id: id, admissionOrdinal: admissionOrdinal, phase: phase) + } + } + + package final class Connection: @unchecked Sendable { + package struct AdmittedRequest: Sendable { + package let operation: RequestOperation + package let lease: WorkLease + } + + package let id = UUID() + package let admissionOrdinal: UInt64 + private weak var owner: MCPHTTPNetworkResourceOwner? + private let resource: any MCPHTTPConnectionResource + private let lock = NSLock() + private var phase: ConnectionPhase = .accepting + private var nextRequestOrdinal: UInt64 = 0 + private var requests: [UUID: RequestOperation] = [:] + private var closeAcknowledged = false + private var closeWaiters: [CheckedContinuation] = [] + + fileprivate init( + admissionOrdinal: UInt64, + resource: any MCPHTTPConnectionResource, + owner: MCPHTTPNetworkResourceOwner + ) { + self.admissionOrdinal = admissionOrdinal + self.resource = resource + self.owner = owner + } + + fileprivate func installCloseAcknowledgement() { + resource.installCloseAcknowledgement { [weak self] in + self?.acknowledgePeerClose() + } + } + + package func admitRequest() -> AdmittedRequest? { + lock.lock() + guard phase == .accepting else { + lock.unlock() + return nil + } + nextRequestOrdinal &+= 1 + let operation = RequestOperation( + admissionOrdinal: nextRequestOrdinal, + connection: self + ) + let lease = operation.makeLease() + requests[operation.id] = operation + lock.unlock() + return .init(operation: operation, lease: lease) + } + + package func peerClosed() { + beginClosing(.peerClosed, signalResourceClose: false) + } + + package func transportFailed(_ message: String) { + beginClosing(.transportFailure(message), signalResourceClose: true) + } + + package func waitUntilClosed() async { + await withCheckedContinuation { continuation in + lock.lock() + if phase == .closed { + lock.unlock() + continuation.resume() + } else { + closeWaiters.append(continuation) + lock.unlock() + } + } + } + + fileprivate func closeAdmission() { + lock.lock() + if phase == .accepting { + phase = .admissionClosed + } + lock.unlock() + } + + fileprivate func beginClosing(_ cause: TerminalCause) { + beginClosing(cause, signalResourceClose: true) + } + + fileprivate func requestDidClose(_ operation: RequestOperation) { + var didClose = false + var waiters: [CheckedContinuation] = [] + lock.lock() + requests.removeValue(forKey: operation.id) + if case .closing = phase, closeAcknowledged, requests.isEmpty { + phase = .closed + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + didClose = true + } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + if didClose { + owner?.connectionDidClose(self) + } + } + + fileprivate func snapshot() -> ConnectionSnapshot { + lock.lock() + let phase = phase + let closeAcknowledged = closeAcknowledged + let requests = requests.values.sorted { + $0.admissionOrdinal < $1.admissionOrdinal + } + lock.unlock() + return .init( + id: id, + admissionOrdinal: admissionOrdinal, + phase: phase, + closeAcknowledged: closeAcknowledged, + requests: requests.map { $0.snapshot() } + ) + } + + private func beginClosing( + _ cause: TerminalCause, + signalResourceClose: Bool + ) { + let requests: [RequestOperation] + var shouldSignalClose = false + lock.lock() + switch phase { + case .accepting, .admissionClosed: + phase = .closing(cause) + requests = Array(self.requests.values) + shouldSignalClose = signalResourceClose + case .closing, .closed: + requests = [] + } + lock.unlock() + for request in requests { + request.beginClosing(cause) + } + if shouldSignalClose { + resource.signalClose() + } + finishIfQuiescent() + } + + private func acknowledgePeerClose() { + let requests: [RequestOperation] + lock.lock() + closeAcknowledged = true + switch phase { + case .accepting, .admissionClosed: + phase = .closing(.peerClosed) + requests = Array(self.requests.values) + case .closing, .closed: + requests = [] + } + lock.unlock() + for request in requests { + request.beginClosing(.peerClosed) + } + finishIfQuiescent() + } + + private func finishIfQuiescent() { + var didClose = false + var waiters: [CheckedContinuation] = [] + lock.lock() + if case .closing = phase, closeAcknowledged, requests.isEmpty { + phase = .closed + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + didClose = true + } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + if didClose { + owner?.connectionDidClose(self) + } + } + } + + package final class ClosingGeneration: @unchecked Sendable { + private weak var owner: MCPHTTPNetworkResourceOwner? + + fileprivate init(owner: MCPHTTPNetworkResourceOwner) { + self.owner = owner + } + + package func waitUntilClosed() async { + await owner?.waitUntilClosed() + } + } + + private struct GenerationState { + var connections: [UUID: Connection] + } + + private enum State { + case accepting(GenerationState) + case admissionClosed(GenerationState) + case closing(GenerationState, TerminalCause) + case closed + } + + package let generationID: UInt64 + private let lock = NSLock() + private var state: State = .accepting(.init(connections: [:])) + private var nextConnectionOrdinal: UInt64 = 0 + private var closeWaiters: [CheckedContinuation] = [] + + package init(generationID: UInt64) { + self.generationID = generationID + } + + func admitConnection(_ channel: any Channel) -> Connection? { + admitConnection(NIOHTTPConnectionResource(channel: channel)) + } + + package func admitConnection(_ resource: any MCPHTTPConnectionResource) -> Connection? { + lock.lock() + guard case .accepting(var accepting) = state else { + lock.unlock() + resource.signalClose() + return nil + } + nextConnectionOrdinal &+= 1 + let connection = Connection( + admissionOrdinal: nextConnectionOrdinal, + resource: resource, + owner: self + ) + accepting.connections[connection.id] = connection + state = .accepting(accepting) + lock.unlock() + connection.installCloseAcknowledgement() + return connection + } + + package func closeAdmission() { + let connections: [Connection] + lock.lock() + switch state { + case .accepting(let accepting): + state = .admissionClosed(accepting) + connections = Array(accepting.connections.values) + case .admissionClosed(let current), .closing(let current, _): + connections = Array(current.connections.values) + case .closed: + connections = [] + } + lock.unlock() + for connection in connections { + connection.closeAdmission() + } + } + + package func beginClosing(_ cause: TerminalCause) -> ClosingGeneration { + let connections: [Connection] + let effectiveCause: TerminalCause + var waiters: [CheckedContinuation] = [] + lock.lock() + switch state { + case .accepting(let current), .admissionClosed(let current): + effectiveCause = cause + if current.connections.isEmpty { + state = .closed + connections = [] + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + } else { + state = .closing(current, cause) + connections = current.connections.values.sorted { + $0.admissionOrdinal < $1.admissionOrdinal + } + } + case .closing(let current, let firstCause): + effectiveCause = firstCause + connections = current.connections.values.sorted { + $0.admissionOrdinal < $1.admissionOrdinal + } + case .closed: + effectiveCause = cause + connections = [] + } + lock.unlock() + for connection in connections { + connection.beginClosing(effectiveCause) + } + for waiter in waiters { + waiter.resume() + } + return ClosingGeneration(owner: self) + } + + package func snapshot() -> Snapshot { + lock.lock() + let phase: GenerationPhase + let connections: [Connection] + switch state { + case .accepting(let current): + phase = .accepting + connections = Array(current.connections.values) + case .admissionClosed(let current): + phase = .admissionClosed + connections = Array(current.connections.values) + case .closing(let current, let cause): + phase = .closing(cause) + connections = Array(current.connections.values) + case .closed: + phase = .closed + connections = [] + } + lock.unlock() + return .init( + generationID: generationID, + phase: phase, + connections: connections.sorted { + $0.admissionOrdinal < $1.admissionOrdinal + }.map { $0.snapshot() } + ) + } + + fileprivate func connectionDidClose(_ connection: Connection) { + var waiters: [CheckedContinuation] = [] + lock.lock() + switch state { + case .accepting(var current): + current.connections.removeValue(forKey: connection.id) + state = .accepting(current) + case .admissionClosed(var current): + current.connections.removeValue(forKey: connection.id) + state = .admissionClosed(current) + case .closing(var current, let cause): + current.connections.removeValue(forKey: connection.id) + if current.connections.isEmpty { + state = .closed + waiters = closeWaiters + closeWaiters.removeAll(keepingCapacity: false) + } else { + state = .closing(current, cause) + } + case .closed: + break + } + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + private func waitUntilClosed() async { + await withCheckedContinuation { continuation in + lock.lock() + if case .closed = state { + lock.unlock() + continuation.resume() + } else { + closeWaiters.append(continuation) + lock.unlock() + } + } + } +} diff --git a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift index 146a316..d949238 100644 --- a/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift +++ b/Tests/CodexReviewMCPServerTests/CodexReviewMCPHTTPServerTests.swift @@ -284,6 +284,41 @@ struct CodexReviewMCPHTTPServerTests { } } + @Test func concurrentStopsJoinTheSameListenerFailureResult() async throws { + let backend = FakeCodexReviewBackend() + let store = CodexReviewStore.makeTestingStore( + backend: TestingCodexReviewStoreBackend(reviewBackend: backend) + ) + let server = CodexReviewMCPHTTPServer( + adapter: CodexReviewMCPServer(store: store), + configuration: .init(host: "127.0.0.1", port: 0) + ) + try await server.start() + await server.holdNextStopCompletionForTesting() + await server.injectNextListenerCloseFailureForTesting("joined listener failure") + + let firstStop = Task { + await lifecycleError { + try await server.stop() + } + } + await server.waitUntilStopCompletionIsHeldForTesting() + let joinedStop = Task { + await lifecycleError { + try await server.stop() + } + } + await server.waitUntilStopJoinsStoppingGenerationForTesting() + await server.releaseStopCompletionForTesting() + + let expected = CodexReviewMCPHTTPServer.LifecycleError( + first: .init(resource: .listener, message: "joined listener failure") + ) + #expect(await firstStop.value == expected) + #expect(await joinedStop.value == expected) + #expect(await server.eventLoopGroupShutdownCountForTesting() == 1) + } + @Test func runningGenerationAdmissionCloseIsRetainedByStop() async throws { let backend = FakeCodexReviewBackend() let store = CodexReviewStore.makeTestingStore( diff --git a/Tests/CodexReviewMCPServerTests/MCPHTTPNetworkResourceOwnerTests.swift b/Tests/CodexReviewMCPServerTests/MCPHTTPNetworkResourceOwnerTests.swift new file mode 100644 index 0000000..e41d3a1 --- /dev/null +++ b/Tests/CodexReviewMCPServerTests/MCPHTTPNetworkResourceOwnerTests.swift @@ -0,0 +1,288 @@ +import Foundation +import Testing +import CodexReviewMCPServer +import CodexReviewTesting + +@Suite("MCP HTTP network resource owner") +struct MCPHTTPNetworkResourceOwnerTests { + @Test func acceptedConnectionClosesOnlyAfterItsChannelAcknowledges() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 1) + let resource = TestingConnectionResource() + _ = try #require(owner.admitConnection(resource)) + + let closing = owner.beginClosing(.serverStop) + await resource.waitUntilCloseIsSignalled() + + let closingSnapshot = owner.snapshot() + #expect(closingSnapshot.phase == .closing(.serverStop)) + #expect(closingSnapshot.connections.count == 1) + #expect(closingSnapshot.connections[0].closeAcknowledged == false) + + resource.acknowledgeClose() + await closing.waitUntilClosed() + + #expect(owner.snapshot().isClosed) + } + + @Test func requestAndWorkLeaseAreRegisteredBeforeTaskCreation() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 2) + let resource = TestingConnectionResource() + let connection = try #require(owner.admitConnection(resource)) + let admitted = try #require(connection.admitRequest()) + + let beforeTask = try #require(owner.snapshot().connections.first?.requests.first) + #expect(beforeTask.id == admitted.operation.id) + #expect(beforeTask.phase == .reserved) + + let didRun = CompletionFlag() + let task = Task { + defer { + admitted.lease.acknowledgeCompletion() + } + guard await admitted.lease.waitUntilStartIsAllowed() else { + return + } + await didRun.complete() + } + admitted.lease.install(task) + await task.value + + #expect(await didRun.isCompleted()) + #expect(owner.snapshot().connections.first?.requests.isEmpty == true) + + let closing = owner.beginClosing(.serverStop) + await resource.waitUntilCloseIsSignalled() + resource.acknowledgeClose() + await closing.waitUntilClosed() + } + + @Test func stopCancellationIsInstalledBeforeReservedWorkCanStart() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 8) + let resource = TestingConnectionResource() + let connection = try #require(owner.admitConnection(resource)) + let admitted = try #require(connection.admitRequest()) + let observation = StartObservation() + let task = Task { + defer { + admitted.lease.acknowledgeCompletion() + } + let wasAllowed = await admitted.lease.waitUntilStartIsAllowed() + await observation.record( + wasAllowed: wasAllowed, + wasCancelled: Task.isCancelled + ) + } + + owner.closeAdmission() + let closing = owner.beginClosing(.serverStop) + await resource.waitUntilCloseIsSignalled() + admitted.lease.install(task) + resource.acknowledgeClose() + + await task.value + await closing.waitUntilClosed() + #expect(await observation.value() == .init( + wasAllowed: false, + wasCancelled: true + )) + } + + @Test func admissionCloseRejectsLateConnectionsAndRequests() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 3) + let acceptedResource = TestingConnectionResource() + let connection = try #require(owner.admitConnection(acceptedResource)) + + owner.closeAdmission() + + #expect(owner.snapshot().phase == .admissionClosed) + #expect(connection.admitRequest() == nil) + let lateResource = TestingConnectionResource() + #expect(owner.admitConnection(lateResource) == nil) + await lateResource.waitUntilCloseIsSignalled() + + let closing = owner.beginClosing(.serverStop) + await acceptedResource.waitUntilCloseIsSignalled() + acceptedResource.acknowledgeClose() + await closing.waitUntilClosed() + } + + @Test func transportFailureWinsPeerCloseAndDrainsItsRequest() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 4) + let resource = TestingConnectionResource() + let connection = try #require(owner.admitConnection(resource)) + let admitted = try #require(connection.admitRequest()) + let started = AsyncGate() + let completionGate = AsyncGate() + let task = Task { + defer { + admitted.lease.acknowledgeCompletion() + } + guard await admitted.lease.waitUntilStartIsAllowed() else { + return + } + await started.open() + await completionGate.waitIgnoringCancellation() + } + admitted.lease.install(task) + await started.wait() + + connection.transportFailed("broken pipe") + await resource.waitUntilCloseIsSignalled() + resource.acknowledgeClose() + + let closingRequest = try #require(owner.snapshot().connections.first?.requests.first) + #expect(closingRequest.phase == .closing(.transportFailure("broken pipe"))) + #expect(owner.snapshot().connections.first?.closeAcknowledged == true) + + await completionGate.open() + await task.value + await connection.waitUntilClosed() + + #expect(owner.snapshot().phase == .accepting) + #expect(owner.snapshot().connections.isEmpty) + } + + @Test func peerCloseCancelsAndDrainsItsRequest() async throws { + let owner = MCPHTTPNetworkResourceOwner(generationID: 5) + let resource = TestingConnectionResource() + let connection = try #require(owner.admitConnection(resource)) + let admitted = try #require(connection.admitRequest()) + let started = AsyncGate() + let completionGate = AsyncGate() + let task = Task { + defer { + admitted.lease.acknowledgeCompletion() + } + guard await admitted.lease.waitUntilStartIsAllowed() else { + return + } + await started.open() + await completionGate.waitIgnoringCancellation() + } + admitted.lease.install(task) + await started.wait() + + resource.acknowledgeClose() + + let closingRequest = try #require(owner.snapshot().connections.first?.requests.first) + #expect(closingRequest.phase == .closing(.peerClosed)) + + await completionGate.open() + await task.value + await connection.waitUntilClosed() + + #expect(owner.snapshot().phase == .accepting) + #expect(owner.snapshot().connections.isEmpty) + } + + @Test func aStoppedGenerationRejectsLateRegistrationWhileRestartUsesANewOwner() async throws { + let first = MCPHTTPNetworkResourceOwner(generationID: 6) + first.closeAdmission() + let firstClosing = first.beginClosing(.serverStop) + await firstClosing.waitUntilClosed() + + let lateFirstResource = TestingConnectionResource() + #expect(first.admitConnection(lateFirstResource) == nil) + await lateFirstResource.waitUntilCloseIsSignalled() + + let second = MCPHTTPNetworkResourceOwner(generationID: 7) + let secondResource = TestingConnectionResource() + _ = try #require(second.admitConnection(secondResource)) + + #expect(first.snapshot().generationID == 6) + #expect(first.snapshot().isClosed) + #expect(second.snapshot().generationID == 7) + #expect(second.snapshot().connections.count == 1) + + let secondClosing = second.beginClosing(.serverStop) + await secondResource.waitUntilCloseIsSignalled() + secondResource.acknowledgeClose() + await secondClosing.waitUntilClosed() + } +} + +private final class TestingConnectionResource: MCPHTTPConnectionResource, @unchecked Sendable { + private let lock = NSLock() + private var closeAcknowledgement: (@Sendable () -> Void)? + private var closeWasAcknowledged = false + private var closeWasSignalled = false + private var closeSignalWaiters: [CheckedContinuation] = [] + + func signalClose() { + let waiters: [CheckedContinuation] + lock.lock() + closeWasSignalled = true + waiters = closeSignalWaiters + closeSignalWaiters.removeAll(keepingCapacity: false) + lock.unlock() + for waiter in waiters { + waiter.resume() + } + } + + func installCloseAcknowledgement(_ acknowledgement: @escaping @Sendable () -> Void) { + var shouldAcknowledge = false + lock.lock() + precondition(closeAcknowledgement == nil) + closeAcknowledgement = acknowledgement + shouldAcknowledge = closeWasAcknowledged + lock.unlock() + if shouldAcknowledge { + acknowledgement() + } + } + + func acknowledgeClose() { + let acknowledgement: (@Sendable () -> Void)? + lock.lock() + closeWasAcknowledged = true + acknowledgement = closeAcknowledgement + lock.unlock() + acknowledgement?() + } + + func waitUntilCloseIsSignalled() async { + await withCheckedContinuation { continuation in + lock.lock() + if closeWasSignalled { + lock.unlock() + continuation.resume() + } else { + closeSignalWaiters.append(continuation) + lock.unlock() + } + } + } +} + +private actor CompletionFlag { + private var completed = false + + func complete() { + completed = true + } + + func isCompleted() -> Bool { + completed + } +} + +private actor StartObservation { + struct Value: Equatable, Sendable { + let wasAllowed: Bool + let wasCancelled: Bool + } + + private var recordedValue: Value? + + func record(wasAllowed: Bool, wasCancelled: Bool) { + recordedValue = .init( + wasAllowed: wasAllowed, + wasCancelled: wasCancelled + ) + } + + func value() -> Value? { + recordedValue + } +}