close() closes immediately; Add new closeGracefully() (#383)

Fixes #370.

This patch changes the behavior of `PostgresConnection.close()`. Currently `close()` terminates the connection only after all queued queries have been successfully processed by the server. This however leads to an unwanted dependency on the Postgres server to close a connection. If a server stops responding, the client is currently unable to close its connection. Because of this, this patch changes the behavior of `close()`. `close()` now terminates a connection immediately and fails all running or queued queries. To allow users to continue to use the existing behavior we introduce a `closeGracefully()` that now has the same behavior as close had previously. 

Since we never documented the old close behavior and we consider it dangerous in certain situations we are fine with changing the behavior without tagging a major version.

Author: fabianfett

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -384,6 +384,17 @@ extension PostgresConnection {
         try await self.close().get()
     }
 
+    /// Closes the connection to the server, _after all queries_ that have been created on this connection have been run.
+    public func closeGracefully() async throws {
+        try await withTaskCancellationHandler { () async throws -> () in
+            let promise = self.eventLoop.makePromise(of: Void.self)
+            self.channel.triggerUserOutboundEvent(PSQLOutgoingEvent.gracefulShutdown, promise: promise)
+            return try await promise.futureResult.get()
+        } onCancel: {
+            _ = self.close()
+        }
+    }
+
     /// Run a query on the Postgres server the connection is connected to.
     ///
     /// - Parameters:

Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift

@@ -36,7 +36,14 @@ struct ListenStateMachine {
     }
 
     mutating func stopListeningSucceeded(channel: String) -> StopListeningSuccessAction {
-        return self.channels[channel, default: .init()].stopListeningSucceeded()
+        switch self.channels[channel]!.stopListeningSucceeded() {
+        case .none:
+            self.channels.removeValue(forKey: channel)
+            return .none
+
+        case .startListening:
+            return .startListening
+        }
     }
 
     enum CancelAction {
@@ -46,7 +53,7 @@ struct ListenStateMachine {
     }
 
     mutating func cancelNotificationListener(channel: String, id: Int) -> CancelAction {
-        return self.channels[channel, default: .init()].cancelListening(id: id)
+        return self.channels[channel]?.cancelListening(id: id) ?? .none
     }
 
     mutating func fail(_ error: Error) -> [NotificationListener] {

Sources/PostgresNIO/New/Connection State Machine/ListenStateMachine.swift

@@ -18,8 +18,9 @@ public struct PSQLError: Error {
 
             case queryCancelled
             case tooManyParameters
-            case connectionQuiescing
-            case connectionClosed
+            case clientClosesConnection
+            case clientClosedConnection
+            case serverClosedConnection
             case connectionError
             case uncleanShutdown
 
@@ -45,13 +46,20 @@ public struct PSQLError: Error {
         public static let invalidCommandTag = Self(.invalidCommandTag)
         public static let queryCancelled = Self(.queryCancelled)
         public static let tooManyParameters = Self(.tooManyParameters)
-        public static let connectionQuiescing = Self(.connectionQuiescing)
-        public static let connectionClosed = Self(.connectionClosed)
+        public static let clientClosesConnection = Self(.clientClosesConnection)
+        public static let clientClosedConnection = Self(.clientClosedConnection)
+        public static let serverClosedConnection = Self(.serverClosedConnection)
         public static let connectionError = Self(.connectionError)
         public static let uncleanShutdown = Self.init(.uncleanShutdown)
         public static let listenFailed = Self.init(.listenFailed)
         public static let unlistenFailed = Self.init(.unlistenFailed)
 
+        @available(*, deprecated, renamed: "clientClosesConnection")
+        public static let connectionQuiescing = Self.clientClosesConnection
+
+        @available(*, deprecated, message: "Use the more specific `serverClosedConnection` or `clientClosedConnection` instead")
+        public static let connectionClosed = Self.serverClosedConnection
+
         public var description: String {
             switch self.base {
             case .sslUnsupported:
@@ -78,10 +86,12 @@ public struct PSQLError: Error {
                 return "queryCancelled"
             case .tooManyParameters:
                 return "tooManyParameters"
-            case .connectionQuiescing:
-                return "connectionQuiescing"
-            case .connectionClosed:
-                return "connectionClosed"
+            case .clientClosesConnection:
+                return "clientClosesConnection"
+            case .clientClosedConnection:
+                return "clientClosedConnection"
+            case .serverClosedConnection:
+                return "serverClosedConnection"
             case .connectionError:
                 return "connectionError"
             case .uncleanShutdown:
@@ -377,19 +387,33 @@ public struct PSQLError: Error {
         return new
     }
 
-    static var connectionQuiescing: PSQLError { PSQLError(code: .connectionQuiescing) }
+    static func clientClosesConnection(underlying: Error?) -> PSQLError {
+        var error = PSQLError(code: .clientClosesConnection)
+        error.underlying = underlying
+        return error
+    }
+
+    static func clientClosedConnection(underlying: Error?) -> PSQLError {
+        var error = PSQLError(code: .clientClosedConnection)
+        error.underlying = underlying
+        return error
+    }
 
-    static var connectionClosed: PSQLError { PSQLError(code: .connectionClosed) }
+    static func serverClosedConnection(underlying: Error?) -> PSQLError {
+        var error = PSQLError(code: .serverClosedConnection)
+        error.underlying = underlying
+        return error
+    }
 
-    static var authMechanismRequiresPassword: PSQLError { PSQLError(code: .authMechanismRequiresPassword) }
+    static let authMechanismRequiresPassword = PSQLError(code: .authMechanismRequiresPassword)
 
-    static var sslUnsupported: PSQLError { PSQLError(code: .sslUnsupported) }
+    static let sslUnsupported = PSQLError(code: .sslUnsupported)
 
-    static var queryCancelled: PSQLError { PSQLError(code: .queryCancelled) }
+    static let queryCancelled = PSQLError(code: .queryCancelled)
 
-    static var uncleanShutdown: PSQLError { PSQLError(code: .uncleanShutdown) }
+    static let uncleanShutdown = PSQLError(code: .uncleanShutdown)
 
-    static var receivedUnencryptedDataAfterSSLRequest: PSQLError { PSQLError(code: .receivedUnencryptedDataAfterSSLRequest) }
+    static let receivedUnencryptedDataAfterSSLRequest = PSQLError(code: .receivedUnencryptedDataAfterSSLRequest)
 
     static func server(_ response: PostgresBackendMessage.ErrorResponse) -> PSQLError {
         var error = PSQLError(code: .server)

Sources/PostgresNIO/New/PSQLError.swift

@@ -7,6 +7,8 @@ enum PSQLOutgoingEvent {
     ///
     /// this shall be removed with the next breaking change and always supplied with `PSQLConnection.Configuration`
     case authenticate(AuthContext)
+
+    case gracefulShutdown
 }
 
 enum PSQLEvent {

Sources/PostgresNIO/New/PSQLEventsHandler.swift

@@ -247,7 +247,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             return
         }
 
-        let action = self.state.close(promise)
+        let action = self.state.close(promise: promise)
         self.run(action, with: context)
     }
 
@@ -258,6 +258,11 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
         case PSQLOutgoingEvent.authenticate(let authContext):
             let action = self.state.provideAuthenticationContext(authContext)
             self.run(action, with: context)
+
+        case PSQLOutgoingEvent.gracefulShutdown:
+            let action = self.state.gracefulClose(promise)
+            self.run(action, with: context)
+
         default:
             context.triggerUserOutboundEvent(event, promise: promise)
         }

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -37,9 +37,9 @@ extension PSQLError {
             return self.underlying ?? self
         case .tooManyParameters, .invalidCommandTag:
             return self
-        case .connectionQuiescing:
-            return PostgresError.connectionClosed
-        case .connectionClosed:
+        case .clientClosesConnection, 
+             .clientClosedConnection,
+             .serverClosedConnection:
             return PostgresError.connectionClosed
         case .connectionError:
             return self.underlying ?? self

Sources/PostgresNIO/Postgres+PSQLCompat.swift

@@ -137,14 +137,14 @@ class ConnectionStateMachineTests: XCTestCase {
 
     func testErrorIsIgnoredWhenClosingConnection() {
         // test ignore unclean shutdown when closing connection
-        var stateIgnoreChannelError = ConnectionStateMachine(.closing)
-        
+        var stateIgnoreChannelError = ConnectionStateMachine(.closing(nil))
+
         XCTAssertEqual(stateIgnoreChannelError.errorHappened(.connectionError(underlying: NIOSSLError.uncleanShutdown)), .wait)
         XCTAssertEqual(stateIgnoreChannelError.closed(), .fireChannelInactive)
 
         // test ignore any other error when closing connection
 
-        var stateIgnoreErrorMessage = ConnectionStateMachine(.closing)
+        var stateIgnoreErrorMessage = ConnectionStateMachine(.closing(nil))
         XCTAssertEqual(stateIgnoreErrorMessage.errorReceived(.init(fields: [:])), .wait)
         XCTAssertEqual(stateIgnoreErrorMessage.closed(), .fireChannelInactive)
     }

Tests/PostgresNIOTests/New/Connection State Machine/ConnectionStateMachineTests.swift

@@ -22,13 +22,13 @@ final class PSQLRowStreamTests: XCTestCase {
 
     func testFailedStream() {
         let stream = PSQLRowStream(
-            source: .noRows(.failure(PSQLError.connectionClosed)),
+            source: .noRows(.failure(PSQLError.serverClosedConnection(underlying: nil))),
             eventLoop: self.eventLoop,
             logger: self.logger
         )
 
         XCTAssertThrowsError(try stream.all().wait()) {
-            XCTAssertEqual($0 as? PSQLError, .connectionClosed)
+            XCTAssertEqual($0 as? PSQLError, .serverClosedConnection(underlying: nil))
         }
     }
 

Tests/PostgresNIOTests/New/PSQLRowStreamTests.swift

@@ -24,8 +24,11 @@ class PostgresChannelHandlerTests: XCTestCase {
             ReverseMessageToByteHandler(PSQLBackendMessageEncoder()),
             handler
         ], loop: self.eventLoop)
-        defer { XCTAssertNoThrow(try embedded.finish()) }
-        
+        defer {
+            do { try embedded.finish() }
+            catch { print("\(String(reflecting: error))") }
+        }
+
         var maybeMessage: PostgresFrontendMessage?
         XCTAssertNoThrow(embedded.connect(to: try .init(ipAddress: "0.0.0.0", port: 5432), promise: nil))
         XCTAssertNoThrow(maybeMessage = try embedded.readOutbound(as: PostgresFrontendMessage.self))

Tests/PostgresNIOTests/New/PostgresChannelHandlerTests.swift

@@ -182,6 +182,98 @@ class PostgresConnectionTests: XCTestCase {
         }
     }
 
+    func testCloseGracefullyClosesWhenInternalQueueIsEmpty() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
+            for _ in 1...2 {
+                taskGroup.addTask {
+                    let rows = try await connection.query("SELECT 1;", logger: self.logger)
+                    var iterator = rows.decode(Int.self).makeAsyncIterator()
+                    let first = try await iterator.next()
+                    XCTAssertEqual(first, 1)
+                    let second = try await iterator.next()
+                    XCTAssertNil(second)
+                }
+            }
+
+            for i in 0...1 {
+                let listenMessage = try await channel.waitForUnpreparedRequest()
+                XCTAssertEqual(listenMessage.parse.query, "SELECT 1;")
+
+                if i == 0 {
+                    taskGroup.addTask {
+                        try await connection.closeGracefully()
+                    }
+                }
+
+                try await channel.writeInbound(PostgresBackendMessage.parseComplete)
+                try await channel.writeInbound(PostgresBackendMessage.parameterDescription(.init(dataTypes: [])))
+                let intDescription = RowDescription.Column(
+                    name: "",
+                    tableOID: 0,
+                    columnAttributeNumber: 0,
+                    dataType: .int8, dataTypeSize: 8, dataTypeModifier: 0, format: .binary
+                )
+                try await channel.writeInbound(PostgresBackendMessage.rowDescription(.init(columns: [intDescription])))
+                try await channel.testingEventLoop.executeInContext { channel.read() }
+                try await channel.writeInbound(PostgresBackendMessage.bindComplete)
+                try await channel.testingEventLoop.executeInContext { channel.read() }
+                try await channel.writeInbound(PostgresBackendMessage.dataRow([Int(1)]))
+                try await channel.testingEventLoop.executeInContext { channel.read() }
+                try await channel.writeInbound(PostgresBackendMessage.commandComplete("SELECT 1 1"))
+                try await channel.testingEventLoop.executeInContext { channel.read() }
+                try await channel.writeInbound(PostgresBackendMessage.readyForQuery(.idle))
+            }
+
+            let terminate = try await channel.waitForOutboundWrite(as: PostgresFrontendMessage.self)
+            XCTAssertEqual(terminate, .terminate)
+            try await channel.closeFuture.get()
+            XCTAssertEqual(channel.isActive, false)
+
+            while let taskResult = await taskGroup.nextResult() {
+                switch taskResult {
+                case .success:
+                    break
+                case .failure(let failure):
+                    XCTFail("Unexpected error: \(failure)")
+                }
+            }
+        }
+    }
+
+    func testCloseClosesImmediatly() async throws {
+        let (connection, channel) = try await self.makeTestConnectionWithAsyncTestingChannel()
+
+        try await withThrowingTaskGroup(of: Void.self) { taskGroup async throws -> () in
+            for _ in 1...2 {
+                taskGroup.addTask {
+                    try await connection.query("SELECT 1;", logger: self.logger)
+                }
+            }
+
+            let listenMessage = try await channel.waitForUnpreparedRequest()
+            XCTAssertEqual(listenMessage.parse.query, "SELECT 1;")
+
+            async let close: () = connection.close()
+
+            try await channel.closeFuture.get()
+            XCTAssertEqual(channel.isActive, false)
+
+            try await close
+
+            while let taskResult = await taskGroup.nextResult() {
+                switch taskResult {
+                case .success:
+                    XCTFail("Expected queries to fail")
+                case .failure(let failure):
+                    guard let error = failure as? PSQLError else {
+                        return XCTFail("Unexpected error type: \(failure)")
+                    }
+                    XCTAssertEqual(error.code, .clientClosedConnection)
+                }
+            }
+        }
+    }
 
     func makeTestConnectionWithAsyncTestingChannel() async throws -> (PostgresConnection, NIOAsyncTestingChannel) {
         let eventLoop = NIOAsyncTestingEventLoop()

Tests/PostgresNIOTests/New/PostgresConnectionTests.swift

@@ -183,7 +183,7 @@ final class PostgresRowSequenceTests: XCTestCase {
             logger: self.logger
         )
 
-        stream.receive(completion: .failure(PSQLError.connectionClosed))
+        stream.receive(completion: .failure(PSQLError.serverClosedConnection(underlying: nil)))
 
         let rowSequence = stream.asyncSequence()
 
@@ -194,7 +194,7 @@ final class PostgresRowSequenceTests: XCTestCase {
             }
             XCTFail("Expected that an error was thrown before.")
         } catch {
-            XCTAssertEqual(error as? PSQLError, .connectionClosed)
+            XCTAssertEqual(error as? PSQLError, .serverClosedConnection(underlying: nil))
         }
     }
 
@@ -255,14 +255,14 @@ final class PostgresRowSequenceTests: XCTestCase {
         XCTAssertEqual(try row1?.decode(Int.self, context: .default), 0)
 
         DispatchQueue.main.asyncAfter(deadline: .now() + .seconds(1)) {
-            stream.receive(completion: .failure(PSQLError.connectionClosed))
+            stream.receive(completion: .failure(PSQLError.serverClosedConnection(underlying: nil)))
         }
 
         do {
             _ = try await rowIterator.next()
             XCTFail("Expected that an error was thrown before.")
         } catch {
-            XCTAssertEqual(error as? PSQLError, .connectionClosed)
+            XCTAssertEqual(error as? PSQLError, .serverClosedConnection(underlying: nil))
         }
     }