Cancellation works when talking to a silent server (#647)

Author: fabianfett

Sources/PostgresNIO/Connection/ConnectCancelHandler.swift

@@ -0,0 +1,79 @@
+import NIOConcurrencyHelpers
+import NIOCore
+
+final class ConnectCancelHandler: Sendable {
+    private enum State {
+        case waiting
+        case cancelledBeforeConnect
+        case tcpConnected(any Channel)
+        case done
+    }
+
+    private let state: NIOLockedValueBox<State>
+
+    init() {
+        self.state = NIOLockedValueBox(.waiting)
+    }
+
+    enum ChannelConnectedAction {
+        case none
+        case close(any Channel)
+    }
+
+    func channelConnected(_ channel: any Channel) -> EventLoopFuture<Void>? {
+        let action = self.state.withLockedValue { state -> ChannelConnectedAction in
+            switch state {
+            case .waiting:
+                state = .tcpConnected(channel)
+                return .none
+            case .cancelledBeforeConnect:
+                state = .done
+                return .close(channel)
+            case .tcpConnected, .done:
+                preconditionFailure("channelConnected called in invalid state")
+            }
+        }
+
+        switch action {
+        case .none:
+            return nil
+        case .close(let channel):
+            channel.close(mode: .all, promise: nil)
+            return channel.closeFuture
+        }
+    }
+
+    enum CancelAction {
+        case none
+        case close(any Channel)
+    }
+
+    func cancel() -> EventLoopFuture<Void>? {
+        let action = self.state.withLockedValue { state -> CancelAction in
+            switch state {
+            case .waiting:
+                state = .cancelledBeforeConnect
+                return .none
+            case .tcpConnected(let channel):
+                state = .done
+                return .close(channel)
+            case .cancelledBeforeConnect, .done:
+                return .none
+            }
+        }
+
+        switch action {
+        case .none:
+            return nil
+        case .close(let channel):
+            channel.close(mode: .all, promise: nil)
+            return channel.closeFuture
+        }
+    }
+
+    func postgresHandshakeDone() {
+        self.state.withLockedValue { state in
+            state = .done
+        }
+    }
+}

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -1,4 +1,5 @@
 import Atomics
+import NIOConcurrencyHelpers
 import NIOCore
 import NIOPosix
 #if canImport(Network)
@@ -129,33 +130,37 @@ public final class PostgresConnection: @unchecked Sendable {
         id connectionID: ID,
         logger: Logger
     ) -> EventLoopFuture<PostgresConnection> {
-        self.connect(
+        let (future, _) = self.connect(
             connectionID: connectionID,
             configuration: .init(configuration),
             logger: logger,
             on: eventLoop
         )
+        return future
     }
 
     static func connect(
         connectionID: ID,
         configuration: PostgresConnection.InternalConfiguration,
         logger: Logger,
         on eventLoop: any EventLoop
-    ) -> EventLoopFuture<PostgresConnection> {
+    ) -> (EventLoopFuture<PostgresConnection>, ConnectCancelHandler) {
 
         var mlogger = logger
         mlogger[postgresMetadataKey: .connectionID] = "\(connectionID)"
         let logger = mlogger
 
+        let cancelHandler = ConnectCancelHandler()
+        let deadline = NIODeadline.now() + configuration.options.connectTimeout
+
         // Here we dispatch to the `eventLoop` first before we setup the EventLoopFuture chain, to
         // ensure all `flatMap`s are executed on the EventLoop (this means the enqueuing of the
         // callbacks).
         //
         // This saves us a number of context switches between the thread the Connection is created
         // on and the EventLoop. In addition, it eliminates all potential races between the creating
         // thread and the EventLoop.
-        return eventLoop.flatSubmit { () -> EventLoopFuture<PostgresConnection> in
+        let future = eventLoop.flatSubmit { () -> EventLoopFuture<PostgresConnection> in
             let connectFuture: EventLoopFuture<any Channel>
 
             switch configuration.connection {
@@ -176,17 +181,48 @@ public final class PostgresConnection: @unchecked Sendable {
             }
 
             return connectFuture.flatMap { channel -> EventLoopFuture<PostgresConnection> in
+                // 1. check if the connection request was cancelled in the mean time.
+                if let closeFuture = cancelHandler.channelConnected(channel) {
+                    return closeFuture.flatMapThrowing { throw CancellationError() }
+                }
+
+                // 2. check if the deadline has elapsed
+                let remaining = deadline - .now()
+                guard remaining > .nanoseconds(0) else {
+                    channel.close(mode: .all, promise: nil)
+                    return channel.closeFuture.flatMapThrowing {
+                        throw PSQLError.connectionError(underlying: ChannelError.connectTimeout(configuration.options.connectTimeout))
+                    }
+                }
+
+                // 3. setup time to enforce connect deadline 
+                let timeoutTask = eventLoop.scheduleTask(deadline: deadline) {
+                    channel.pipeline.fireErrorCaught(
+                        ChannelError.connectTimeout(configuration.options.connectTimeout)
+                    )
+                }
+
                 let connection = PostgresConnection(channel: channel, connectionID: connectionID, logger: logger)
-                return connection.start(configuration: configuration).map { _ in connection }
+                return connection.start(configuration: configuration).map { _ in
+                    timeoutTask.cancel()
+                    return connection
+                }.flatMapError { error in
+                    timeoutTask.cancel()
+                    return eventLoop.makeFailedFuture(error)
+                }
             }.flatMapErrorThrowing { error -> PostgresConnection in
                 switch error {
-                case is PSQLError:
+                case is PSQLError, is CancellationError:
                     throw error
                 default:
                     throw PSQLError.connectionError(underlying: error)
                 }
             }
         }
+
+        future.whenComplete { _ in cancelHandler.postgresHandshakeDone() }
+
+        return (future, cancelHandler)
     }
 
     static func makeBootstrap(
@@ -319,12 +355,13 @@ extension PostgresConnection {
                 options: options
             )
 
-            return PostgresConnection.connect(
+            let (future, _) = PostgresConnection.connect(
                 connectionID: self.idGenerator.wrappingIncrementThenLoad(ordering: .relaxed),
                 configuration: configuration,
                 logger: logger,
                 on: eventLoop
             )
+            return future
         }.flatMapErrorThrowing { error in
             throw error.asAppropriatePostgresError
         }
@@ -373,12 +410,17 @@ extension PostgresConnection {
         id connectionID: ID,
         logger: Logger
     ) async throws -> PostgresConnection {
-        try await self.connect(
+        let (future, cancelHandler) = self.connect(
             connectionID: connectionID,
             configuration: .init(configuration),
             logger: logger,
             on: eventLoop
-        ).get()
+        )
+        return try await withTaskCancellationHandler {
+            try await future.get()
+        } onCancel: {
+            cancelHandler.cancel()
+        }
     }
 
     /// Closes the connection to the server.

Sources/PostgresNIO/New/PSQLEventsHandler.swift

@@ -38,6 +38,7 @@ final class PSQLEventsHandler: ChannelInboundHandler {
         case connected
         case readyForStartup
         case authenticated
+        case closed
     }
 
     private var readyForStartupPromise: EventLoopPromise<Void>!
@@ -65,7 +66,7 @@ final class PSQLEventsHandler: ChannelInboundHandler {
                 // successful
                 self.state = .authenticated
                 self.authenticatePromise.succeed(Void())
-            case .authenticated:
+            case .authenticated, .closed:
                 break
             }
         default:
@@ -89,16 +90,37 @@ final class PSQLEventsHandler: ChannelInboundHandler {
         context.fireChannelActive()
     }
 
+    func channelInactive(context: ChannelHandlerContext) {
+        switch self.state {
+        case .initialized:
+            self.state = .closed
+            self.readyForStartupPromise?.fail(PSQLError.clientClosedConnection(underlying: nil))
+            self.authenticatePromise?.fail(PSQLError.clientClosedConnection(underlying: nil))
+        case .connected:
+            self.state = .closed
+            self.readyForStartupPromise.fail(PSQLError.clientClosedConnection(underlying: nil))
+            self.authenticatePromise.fail(PSQLError.clientClosedConnection(underlying: nil))
+        case .readyForStartup:
+            self.state = .closed
+            self.authenticatePromise.fail(PSQLError.clientClosedConnection(underlying: nil))
+        case .authenticated, .closed:
+            break
+        }
+        context.fireChannelInactive()
+    }
+
     func errorCaught(context: ChannelHandlerContext, error: any Error) {
         switch self.state {
         case .initialized:
             preconditionFailure("Unexpected message for state")
         case .connected:
+            self.state = .closed
             self.readyForStartupPromise.fail(error)
             self.authenticatePromise.fail(error)
         case .readyForStartup:
+            self.state = .closed
             self.authenticatePromise.fail(error)
-        case .authenticated:
+        case .authenticated, .closed:
             break
         }
 

Sources/PostgresNIO/Pool/ConnectionFactory.swift

@@ -46,7 +46,7 @@ final class ConnectionFactory: Sendable {
             configuration: config,
             id: connectionID,
             logger: connectionLogger
-        ).get()
+        )
     }
 
     func makeConnectionConfig() async throws -> PostgresConnection.Configuration {

Tests/PostgresNIOTests/New/PostgresClientTests.swift

@@ -0,0 +1,41 @@
+import Logging
+import NIOCore
+import PostgresNIO
+import Testing
+
+@Suite struct PostgresClientTests {
+
+    @Test(.timeLimit(.minutes(1)))
+    func clientRunExitsPromptlyOnCancellationWhileConnecting() async throws {
+        try await withSilentServer { port in
+            var config = PostgresClient.Configuration(
+                host: "127.0.0.1",
+                port: port,
+                username: "postgres",
+                password: "irrelevant",
+                database: "test",
+                tls: .disable
+            )
+            // Long connect timeout so the hang is obvious if the fix regresses.
+            config.options.connectTimeout = .seconds(30)
+            config.options.minimumConnections = 1
+
+            let client = PostgresClient(
+                configuration: config,
+                eventLoopGroup: NIOSingletons.posixEventLoopGroup,
+                backgroundLogger: Logger(label: "test")
+            )
+
+            await withTaskGroup(of: Void.self) { group in
+                group.addTask { await client.run() }
+
+                // Give the pool enough time to start the connection attempt (enter .starting state).
+                try? await Task.sleep(for: .milliseconds(200))
+
+                // Cancelling the group must make pool.run() return quickly — not in 30 seconds.
+                group.cancelAll()
+            }
+            // If we reach here the test passed (the .timeLimit above enforces the deadline).
+        }
+    }
+}

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -1042,6 +1042,34 @@ import Synchronization
         }
     }
 
+    @Test(.timeLimit(.minutes(1))) func connectEnforcesDeadlineWithSilentServer() async throws {
+        try await withSilentServer { port in
+            var config = PostgresConnection.Configuration(
+                host: "127.0.0.1",
+                port: port,
+                username: "postgres",
+                password: "irrelevant",
+                database: "test",
+                tls: .disable
+            )
+            config.options.connectTimeout = .milliseconds(500)
+
+            let start = ContinuousClock.now
+
+            await #expect(throws: PSQLError.self) {
+                _ = try await PostgresConnection.connect(
+                    configuration: config,
+                    id: 1,
+                    logger: Logger(label: "test")
+                )
+            }
+
+            let elapsed = ContinuousClock.now - start
+            #expect(elapsed < .seconds(5))
+            #expect(elapsed >= .milliseconds(400))
+        }
+    }
+
     func withAsyncTestingChannel(_ body: (PostgresConnection, NIOAsyncTestingChannel) async throws -> ()) async throws {
         let eventLoop = NIOAsyncTestingEventLoop()
         let channel = try await NIOAsyncTestingChannel(loop: eventLoop) { channel in

Tests/PostgresNIOTests/Utilities.swift

@@ -1,3 +1,5 @@
+import NIOCore
+import NIOPosix
 import Logging
 
 extension Logger {
@@ -7,3 +9,16 @@ extension Logger {
         return logger
     }
 }
+
+func withSilentServer<Success: ~Copyable>(_ body: (_ port: Int) async throws -> Success) async throws -> Success{
+    let server = try await ServerBootstrap(group: NIOSingletons.posixEventLoopGroup)
+        .bind(to: .init(ipAddress: "127.0.0.1", port: 0)).get()
+    let result: Result<Success, any Error>
+    do {
+        result = .success(try await body(server.localAddress!.port!))
+    } catch {
+        result = .failure(error)
+    }
+    try? await server.close().get()
+    return try result.get()
+}