Fix row stream cancel/error behavior (vapor#353)

Author: fabianfett

Diff for: Package.swift

@@ -14,7 +14,7 @@ let package = Package(
     ],
     dependencies: [
         .package(url: "https://github.com/apple/swift-atomics.git", from: "1.1.0"),
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.50.0"),
+        .package(url: "https://github.com/apple/swift-nio.git", from: "2.51.1"),
         .package(url: "https://github.com/apple/swift-nio-transport-services.git", from: "1.16.0"),
         .package(url: "https://github.com/apple/swift-nio-ssl.git", from: "2.23.1"),
         .package(url: "https://github.com/apple/swift-crypto.git", "1.0.0" ..< "3.0.0"),

Diff for: Sources/PostgresNIO/New/Connection State Machine/ConnectionStateMachine.swift

@@ -841,7 +841,7 @@ struct ConnectionStateMachine {
     // MARK: Consumer
 
     mutating func cancelQueryStream() -> ConnectionAction {
-        guard case .extendedQuery(var queryState, let connectionContext) = self.state, !queryState.isComplete else {
+        guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
             preconditionFailure("Tried to cancel stream without active query")
         }
 
@@ -926,6 +926,8 @@ struct ConnectionStateMachine {
                  .wait,
                  .read:
                 preconditionFailure("Expecting only failure actions if an error happened")
+            case .evaluateErrorAtConnectionLevel:
+                return .closeConnectionAndCleanup(cleanupContext)
             case .failQuery(let queryContext, with: let error):
                 return .failQuery(queryContext, with: error, cleanupContext: cleanupContext)
             case .forwardStreamError(let error, let read):
@@ -1169,6 +1171,12 @@ extension ConnectionStateMachine {
         case .forwardStreamError(let error, let read):
             let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error)
             return .forwardStreamError(error, read: read, cleanupContext: cleanupContext)
+
+        case .evaluateErrorAtConnectionLevel(let error):
+            if let cleanupContext = self.setErrorAndCreateCleanupContextIfNeeded(error) {
+                return .closeConnectionAndCleanup(cleanupContext)
+            }
+            return .wait
         case .read:
             return .read
         case .wait:

Diff for: Sources/PostgresNIO/New/Connection State Machine/ExtendedQueryStateMachine.swift

@@ -32,7 +32,9 @@ struct ExtendedQueryStateMachine {
         case failQuery(ExtendedQueryContext, with: PSQLError)
         case succeedQuery(ExtendedQueryContext, columns: [RowDescription.Column])
         case succeedQueryNoRowsComming(ExtendedQueryContext, commandTag: String)
-        
+
+        case evaluateErrorAtConnectionLevel(PSQLError)
+
         // --- streaming actions
         // actions if query has requested next row but we are waiting for backend
         case forwardRows([DataRow])
@@ -422,11 +424,15 @@ struct ExtendedQueryStateMachine {
              .noDataMessageReceived(let context),
              .bindCompleteReceived(let context):
             self.state = .error(error)
-            return .failQuery(context, with: error)
+            if self.isCancelled {
+                return .evaluateErrorAtConnectionLevel(error)
+            } else {
+                return .failQuery(context, with: error)
+            }
 
         case .drain:
             self.state = .error(error)
-            return .forwardStreamError(error, read: false)
+            return .evaluateErrorAtConnectionLevel(error)
 
         case .streaming(_, var streamStateMachine):
             self.state = .error(error)

Diff for: Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -273,9 +273,8 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
 
 
         case .forwardStreamError(let error, let read, let cleanupContext):
-            let rowStream = self.rowStream!
+            self.rowStream!.receive(completion: .failure(error))
             self.rowStream = nil
-            rowStream.receive(completion: .failure(error))
             if let cleanupContext = cleanupContext {
                 self.closeConnectionAndCleanup(cleanupContext, context: context)
             } else if read {
@@ -512,7 +511,6 @@ extension PostgresChannelHandler: PSQLRowsDataSource {
         guard self.rowStream === stream, let handlerContext = self.handlerContext else {
             return
         }
-        // we ignore this right now :)
         let action = self.state.cancelQueryStream()
         self.run(action, with: handlerContext)
     }

Diff for: Tests/IntegrationTests/AsyncTests.swift

@@ -75,6 +75,43 @@ final class AsyncPostgresConnectionTests: XCTestCase {
             }
 
             XCTAssertFalse(connection.isClosed, "Connection should survive!")
+
+            for num in 0..<10 {
+                for try await decoded in try await connection.query("SELECT \(num);", logger: .psqlTest).decode(Int.self) {
+                    XCTAssertEqual(decoded, num)
+                }
+            }
+        }
+    }
+
+    func testConnectionSurvives1kQueriesWithATypo() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        let start = 1
+        let end = 10000
+
+        try await withTestConnection(on: eventLoop) { connection -> () in
+            for _ in 0..<1000 {
+                do {
+                    try await connection.query("SELECT generte_series(\(start), \(end));", logger: .psqlTest)
+                    XCTFail("Expected to throw from the request")
+                } catch {
+                    guard let error = error as? PSQLError else { return XCTFail("Unexpected error type: \(error)") }
+
+                    XCTAssertEqual(error.code, .server)
+                    XCTAssertEqual(error.serverInfo?[.severity], "ERROR")
+                }
+            }
+
+            // the connection survived all of this, we can still run normal queries:
+
+            for num in 0..<10 {
+                for try await decoded in try await connection.query("SELECT \(num);", logger: .psqlTest).decode(Int.self) {
+                    XCTAssertEqual(decoded, num)
+                }
+            }
         }
     }
 
@@ -172,6 +209,53 @@ final class AsyncPostgresConnectionTests: XCTestCase {
         }
     }
     #endif
+
+    func testCancelTaskThatIsVeryLongRunningWhichAlsoFailsWhileInStreamingMode() async throws {
+        let eventLoopGroup = MultiThreadedEventLoopGroup(numberOfThreads: 1)
+        defer { XCTAssertNoThrow(try eventLoopGroup.syncShutdownGracefully()) }
+        let eventLoop = eventLoopGroup.next()
+
+        // we cancel the query after 400ms.
+        // the server times out the query after 1sec.
+
+        try await withTestConnection(on: eventLoop) { connection -> () in
+            try await connection.query("SET statement_timeout=1000;", logger: .psqlTest) // 1000 milliseconds
+
+            try await withThrowingTaskGroup(of: Void.self) { group in
+                group.addTask {
+                    let start = 1
+                    let end = 100_000_000
+
+                    let rows = try await connection.query("SELECT generate_series(\(start), \(end));", logger: .psqlTest)
+                    var counter = 0
+                    do {
+                        for try await element in rows.decode(Int.self, context: .default) {
+                            XCTAssertEqual(element, counter + 1)
+                            counter += 1
+                        }
+                        XCTFail("Expected to get cancelled while reading the query")
+                        XCTAssertEqual(counter, end)
+                    } catch let error as CancellationError {
+                        XCTAssertGreaterThanOrEqual(counter, 1)
+                        // Expected
+                        print("\(error)")
+                    } catch {
+                        XCTFail("Unexpected error: \(error)")
+                    }
+
+                    XCTAssertTrue(Task.isCancelled)
+                    XCTAssertFalse(connection.isClosed, "Connection should survive!")
+                }
+
+                let delay: UInt64 = 400_000_000 // 400 milliseconds
+                try await Task.sleep(nanoseconds: delay)
+
+                group.cancelAll()
+            }
+
+            try await connection.query("SELECT 1;", logger: .psqlTest)
+        }
+    }
 }
 
 extension XCTestCase {

Diff for: Tests/PostgresNIOTests/New/Connection State Machine/ExtendedQueryStateMachineTests.swift

@@ -181,4 +181,100 @@ class ExtendedQueryStateMachineTests: XCTestCase {
         XCTAssertEqual(state.commandCompletedReceived("SELECT 4"), .wait)
         XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
     }
+
+    func testCancelQueryAfterServerError() {
+        var state = ConnectionStateMachine.readyForQuery()
+
+        let logger = Logger.psqlTest
+        let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
+        promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
+        let query: PostgresQuery = "SELECT version()"
+        let queryContext = ExtendedQueryContext(query: query, logger: logger, promise: promise)
+
+        XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query))
+        XCTAssertEqual(state.parseCompleteReceived(), .wait)
+        XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait)
+
+        // We need to ensure that even though the row description from the wire says that we
+        // will receive data in `.text` format, we will actually receive it in binary format,
+        // since we requested it in binary with our bind message.
+        let input: [RowDescription.Column] = [
+            .init(name: "version", tableOID: 0, columnAttributeNumber: 0, dataType: .text, dataTypeSize: -1, dataTypeModifier: -1, format: .text)
+        ]
+        let expected: [RowDescription.Column] = input.map {
+            .init(name: $0.name, tableOID: $0.tableOID, columnAttributeNumber: $0.columnAttributeNumber, dataType: $0.dataType,
+                  dataTypeSize: $0.dataTypeSize, dataTypeModifier: $0.dataTypeModifier, format: .binary)
+        }
+
+        XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
+        XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(queryContext, columns: expected))
+        let dataRows1: [DataRow] = [
+            [ByteBuffer(string: "test1")],
+            [ByteBuffer(string: "test2")],
+            [ByteBuffer(string: "test3")]
+        ]
+        for row in dataRows1 {
+            XCTAssertEqual(state.dataRowReceived(row), .wait)
+        }
+        XCTAssertEqual(state.channelReadComplete(), .forwardRows(dataRows1))
+        XCTAssertEqual(state.readEventCaught(), .wait)
+        XCTAssertEqual(state.requestQueryRows(), .read)
+        let dataRows2: [DataRow] = [
+            [ByteBuffer(string: "test4")],
+            [ByteBuffer(string: "test5")],
+            [ByteBuffer(string: "test6")]
+        ]
+        for row in dataRows2 {
+            XCTAssertEqual(state.dataRowReceived(row), .wait)
+        }
+        let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"])
+        XCTAssertEqual(state.errorReceived(serverError), .forwardStreamError(.server(serverError), read: false, cleanupContext: .none))
+
+        XCTAssertEqual(state.channelReadComplete(), .wait)
+        XCTAssertEqual(state.readEventCaught(), .read)
+
+        XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
+    }
+
+    func testQueryErrorDoesNotKillConnection() {
+        var state = ConnectionStateMachine.readyForQuery()
+
+        let logger = Logger.psqlTest
+        let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
+        promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
+        let query: PostgresQuery = "SELECT version()"
+        let queryContext = ExtendedQueryContext(query: query, logger: logger, promise: promise)
+
+        XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query))
+        XCTAssertEqual(state.parseCompleteReceived(), .wait)
+        XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait)
+
+        let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"])
+        XCTAssertEqual(
+            state.errorReceived(serverError), .failQuery(queryContext, with: .server(serverError), cleanupContext: .none)
+        )
+
+        XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
+    }
+
+    func testQueryErrorAfterCancelDoesNotKillConnection() {
+        var state = ConnectionStateMachine.readyForQuery()
+
+        let logger = Logger.psqlTest
+        let promise = EmbeddedEventLoop().makePromise(of: PSQLRowStream.self)
+        promise.fail(PSQLError.uncleanShutdown) // we don't care about the error at all.
+        let query: PostgresQuery = "SELECT version()"
+        let queryContext = ExtendedQueryContext(query: query, logger: logger, promise: promise)
+
+        XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query))
+        XCTAssertEqual(state.parseCompleteReceived(), .wait)
+        XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait)
+        XCTAssertEqual(state.cancelQueryStream(), .failQuery(queryContext, with: .queryCancelled, cleanupContext: .none))
+
+        let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"])
+        XCTAssertEqual(state.errorReceived(serverError), .wait)
+
+        XCTAssertEqual(state.readyForQueryReceived(.idle), .fireEventReadyForQuery)
+    }
+
 }