Implement cancellation during copy operations

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection+CopyFrom.swift

@@ -23,13 +23,17 @@ public struct PostgresCopyFromWriter: Sendable {
         precondition(eventLoop.inEventLoop)
         let promise = eventLoop.makePromise(of: Void.self)
         self.channelHandler.value.checkBackendCanReceiveCopyData(promise: promise)
-        promise.futureResult.map {
+        promise.futureResult.flatMap {
             if eventLoop.inEventLoop {
-                self.channelHandler.value.sendCopyData(byteBuffer)
+                return eventLoop.makeCompletedFuture(withResultOf: {
+                    try self.channelHandler.value.sendCopyData(byteBuffer)
+                })
             } else {
+                let promise = eventLoop.makePromise(of: Void.self)
                 eventLoop.execute {
-                    self.channelHandler.value.sendCopyData(byteBuffer)
+                    promise.completeWith(Result(catching: { try self.channelHandler.value.sendCopyData(byteBuffer) }))
                 }
+                return promise.futureResult
             }
         }.whenComplete { result in
             continuation.resume(with: result)
@@ -45,13 +49,32 @@ public struct PostgresCopyFromWriter: Sendable {
         // `writeData` closure. It is likely that the user would forget to do so.
         try Task.checkCancellation()
 
-        // TODO: Listen for task cancellation while we are waiting for backpressure to clear.
-        try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
+        try await withTaskCancellationHandler {
+            do {
+                try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
+                    if eventLoop.inEventLoop {
+                        writeAssumingInEventLoop(byteBuffer, continuation)
+                    } else {
+                        eventLoop.execute {
+                            writeAssumingInEventLoop(byteBuffer, continuation)
+                        }
+                    }
+                }
+            } catch {
+                if Task.isCancelled {
+                    // If the task was cancelled, we might receive a postgres error which is an artifact about how we
+                    // communicate the cancellation to the state machine. Throw a `CancellationError` to the user
+                    // instead, which looks more like native Swift Concurrency code.
+                    throw CancellationError()
+                }
+                throw error
+            }
+        } onCancel: {
             if eventLoop.inEventLoop {
-                writeAssumingInEventLoop(byteBuffer, continuation)
+                self.channelHandler.value.cancel()
             } else {
                 eventLoop.execute {
-                    writeAssumingInEventLoop(byteBuffer, continuation)
+                    self.channelHandler.value.cancel()
                 }
             }
         }

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

@@ -110,6 +110,12 @@ struct ConnectionStateMachine {
         /// Send a `CopyFail` message to the backend with the given error message.
         case sendCopyFail(message: String)
 
+        /// Fail the promise with the given error and close the connection.
+        ///
+        /// This is used when we want to cancel a COPY operation while waiting for backpressure relieve. In that case we
+        /// can't recover the connection because we can't send any messages to the backend, so we need to close it.
+        case failPromiseAndCloseConnection(EventLoopPromise<Void>, error: PSQLError, cleanupContext: CleanUpContext)
+
         // --- streaming actions
         // actions if query has requested next row but we are waiting for backend
         case forwardRows([DataRow])
@@ -867,9 +873,10 @@ struct ConnectionStateMachine {
 
     // MARK: Consumer
 
-    mutating func cancelQueryStream() -> ConnectionAction {
+    mutating func cancel() -> ConnectionAction {
         guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
-            preconditionFailure("Tried to cancel stream without active query")
+            // We are not in a state in which we can cancel. Do nothing.
+            return .wait
         }
 
         self.state = .modifying // avoid CoW
@@ -955,7 +962,8 @@ struct ConnectionStateMachine {
                  .triggerCopyData,
                  .sendCopyDoneAndSync,
                  .sendCopyFail,
-                 .succeedQueryContinuation:
+                 .succeedQueryContinuation,
+                 .failPromiseAndCloseConnection:
                 preconditionFailure("Invalid query state machine action in state: \(self.state), action: \(action)")
 
             case .evaluateErrorAtConnectionLevel:
@@ -1149,6 +1157,9 @@ extension ConnectionStateMachine {
             return .sendCopyDoneAndSync
         case .sendCopyFail(message: let message):
             return .sendCopyFail(message: message)
+        case .failPromiseAndCloseConnection(let promise, error: let error):
+            let cleanupContext = self.setErrorAndCreateCleanupContext(error)
+            return .failPromiseAndCloseConnection(promise, error: error, cleanupContext: cleanupContext)
         case .forwardRows(let buffer):
             return .forwardRows(buffer)
         case .forwardStreamComplete(let buffer, let commandTag):

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

@@ -79,6 +79,12 @@ struct ExtendedQueryStateMachine {
         /// Send a `CopyFail` message to the backend with the given error message.
         case sendCopyFail(message: String)
 
+        /// Fail the promise with the given error and close the connection.
+        ///
+        /// This is used when we want to cancel a COPY operation while waiting for backpressure relieve. In that case we
+        /// can't recover the connection because we can't send any messages to the backend, so we need to close it.
+        case failPromiseAndCloseConnection(EventLoopPromise<Void>, error: PSQLError)
+
         // --- streaming actions
         // actions if query has requested next row but we are waiting for backend
         case forwardRows([DataRow])
@@ -155,8 +161,16 @@ struct ExtendedQueryStateMachine {
                 return .failPreparedStatementCreation(eventLoopPromise, with: .queryCancelled)
             }
 
-        case .copyingData:
-            return .sendCopyFail(message: "Copy cancelled")
+        case .copyingData(.readyToSend):
+            // We can't initiate an exit from the copy state here because `copyingFinished`, which is the state that is
+            // reached after sending a `CopyFail` requires a continuation that waits for the `CommandComplete` or
+            // `ErrorResponse`. Instead, we assume that the next call to `CopyFromWriter.write` checks cancellation and
+            // initiates the `CopyFail` with the cancellation.
+            return .wait
+
+        case .copyingData(.pendingBackpressureRelieve(let promise)):
+            self.state = .error(.queryCancelled)
+            return .failPromiseAndCloseConnection(promise, error: .queryCancelled)
 
         case .copyingFinished:
             // We already finished the copy and are awaiting the `CommandComplete` or `ErrorResponse` from it. There's

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -177,26 +177,50 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
     /// `ErrorResponse`. This is mostly the case when malformed data is sent to it. In that case, the data transfer
     /// should be aborted to avoid unnecessary work.
     func checkBackendCanReceiveCopyData(promise: EventLoopPromise<Void>) {
-        self.state.checkBackendCanReceiveCopyData(channelIsWritable: handlerContext!.channel.isWritable, promise: promise)
+        guard let handlerContext else {
+            promise.fail(PostgresError.connectionClosed)
+            return
+        }
+        self.state.checkBackendCanReceiveCopyData(channelIsWritable: handlerContext.channel.isWritable, promise: promise)
+    }
+
+    /// Cancel the currently executing operation, if it is cancellable.
+    func cancel() {
+        guard let handlerContext else {
+            return
+        }
+        let action = self.state.cancel()
+        self.run(action, with: handlerContext)
     }
 
     /// Send a `CopyData` message to the backend using the given data.
-    func sendCopyData(_ data: ByteBuffer) {
+    func sendCopyData(_ data: ByteBuffer) throws {
+        guard let handlerContext else {
+            throw PostgresError.connectionClosed
+        }
         self.encoder.copyDataHeader(dataLength: UInt32(data.readableBytes))
-        self.handlerContext!.write(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
-        self.handlerContext!.writeAndFlush(self.wrapOutboundOut(data), promise: nil)
+        handlerContext.write(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
+        handlerContext.writeAndFlush(self.wrapOutboundOut(data), promise: nil)
     }
 
     /// Put the state machine out of the copying mode and send a `CopyDone` message to the backend.
     func sendCopyDone(continuation: CheckedContinuation<Void, any Error>) {
+        guard let handlerContext else {
+            continuation.resume(throwing: PostgresError.connectionClosed)
+            return
+        }
         let action = self.state.sendCopyDone(continuation: continuation)
-        self.run(action, with: self.handlerContext!)
+        self.run(action, with: handlerContext)
     }
 
     /// Put the state machine out of the copying mode and send a `CopyFail` message to the backend.
     func sendCopyFail(message: String, continuation: CheckedContinuation<Void, any Error>) {
+        guard let handlerContext else {
+            continuation.resume(throwing: PostgresError.connectionClosed)
+            return
+        }
         let action = self.state.sendCopyFail(message: message, continuation: continuation)
-        self.run(action, with: self.handlerContext!)
+        self.run(action, with: handlerContext)
     }
 
     func channelReadComplete(context: ChannelHandlerContext) {
@@ -489,6 +513,9 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             }
         case .forwardNotificationToListeners(let notification):
             self.forwardNotificationToListeners(notification, context: context)
+        case .failPromiseAndCloseConnection(let promise, let error, let cleanupContext):
+            promise.fail(error)
+            self.closeConnectionAndCleanup(cleanupContext, context: context)
         }
     }
 
@@ -860,11 +887,10 @@ extension PostgresChannelHandler: PSQLRowsDataSource {
     }
 
     func cancel(for stream: PSQLRowStream) {
-        guard self.rowStream === stream, let handlerContext = self.handlerContext else {
+        guard self.rowStream === stream else {
             return
         }
-        let action = self.state.cancelQueryStream()
-        self.run(action, with: handlerContext)
+        self.cancel()
     }
 }
 

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

@@ -140,7 +140,7 @@ class ExtendedQueryStateMachineTests: XCTestCase {
 
         XCTAssertEqual(state.rowDescriptionReceived(.init(columns: input)), .wait)
         XCTAssertEqual(state.bindCompleteReceived(), .succeedQuery(promise, with: .init(value: .rowDescription(expected), logger: logger)))
-        XCTAssertEqual(state.cancelQueryStream(), .forwardStreamError(.queryCancelled, read: false, cleanupContext: nil))
+        XCTAssertEqual(state.cancel(), .forwardStreamError(.queryCancelled, read: false, cleanupContext: nil))
         XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test1")]), .wait)
         XCTAssertEqual(state.channelReadComplete(), .wait)
         XCTAssertEqual(state.readEventCaught(), .read)
@@ -188,7 +188,7 @@ class ExtendedQueryStateMachineTests: XCTestCase {
         XCTAssertEqual(state.dataRowReceived(row1), .wait)
         XCTAssertEqual(state.channelReadComplete(), .forwardRows([row1]))
         XCTAssertEqual(state.readEventCaught(), .wait)
-        XCTAssertEqual(state.cancelQueryStream(), .forwardStreamError(.queryCancelled, read: true, cleanupContext: nil))
+        XCTAssertEqual(state.cancel(), .forwardStreamError(.queryCancelled, read: true, cleanupContext: nil))
 
         XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test2")]), .wait)
         XCTAssertEqual(state.dataRowReceived([ByteBuffer(string: "test3")]), .wait)
@@ -287,7 +287,7 @@ class ExtendedQueryStateMachineTests: XCTestCase {
         XCTAssertEqual(state.enqueue(task: .extendedQuery(queryContext)), .sendParseDescribeBindExecuteSync(query))
         XCTAssertEqual(state.parseCompleteReceived(), .wait)
         XCTAssertEqual(state.parameterDescriptionReceived(.init(dataTypes: [.int8])), .wait)
-        XCTAssertEqual(state.cancelQueryStream(), .failQuery(promise, with: .queryCancelled, cleanupContext: .none))
+        XCTAssertEqual(state.cancel(), .failQuery(promise, with: .queryCancelled, cleanupContext: .none))
 
         let serverError = PostgresBackendMessage.ErrorResponse(fields: [.severity: "Error", .sqlState: "123"])
         XCTAssertEqual(state.errorReceived(serverError), .wait)