Communicate side effects of checkBackendCanReceiveCopyData via an action

Author: ahoppen

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

@@ -141,6 +141,17 @@ struct ConnectionStateMachine {
         case succeedPromise(EventLoopPromise<Void>)
     }
 
+    enum CheckBackendCanReceiveCopyDataAction {
+        /// Don't perform any action.
+        case none
+
+        /// Succeed the promise with a Void result.
+        case succeedPromise(EventLoopPromise<Void>)
+
+        /// Fail the promise with the given error.
+        case failPromise(EventLoopPromise<Void>, error: any Error)
+    }
+
     private var state: State
     private let requireBackendKeyData: Bool
     private var taskQueue = CircularBuffer<PSQLTask>()
@@ -815,14 +826,15 @@ struct ConnectionStateMachine {
     /// The promise may be failed if the backend indicated that it can't handle any more data by sending an
     /// `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.
-    mutating func checkBackendCanReceiveCopyData(channelIsWritable: Bool, promise: EventLoopPromise<Void>) {
+    mutating func checkBackendCanReceiveCopyData(channelIsWritable: Bool, promise: EventLoopPromise<Void>) -> CheckBackendCanReceiveCopyDataAction {
         guard case .extendedQuery(var queryState, let connectionContext) = self.state else {
             preconditionFailure("Copy mode is only supported for extended queries")
         }
 
         self.state = .modifying // avoid CoW
-        queryState.checkBackendCanReceiveCopyData(channelIsWritable: channelIsWritable, promise: promise)
+        let action = queryState.checkBackendCanReceiveCopyData(channelIsWritable: channelIsWritable, promise: promise)
         self.state = .extendedQuery(queryState, connectionContext)
+        return action
     }
 
     /// Put the state machine out of the copying mode and send a `CopyDone` message to the backend.

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

@@ -423,22 +423,22 @@ struct ExtendedQueryStateMachine {
     /// The promise may be failed if the backend indicated that it can't handle any more data by sending an
     /// `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.
-    mutating func checkBackendCanReceiveCopyData(channelIsWritable: Bool, promise: EventLoopPromise<Void>) {
+    mutating func checkBackendCanReceiveCopyData(channelIsWritable: Bool, promise: EventLoopPromise<Void>) -> ConnectionStateMachine.CheckBackendCanReceiveCopyDataAction {
         if case .error(let error) = self.state {
             // The backend sent us an ErrorResponse during the copy operation. Indicate to the client that it should
             // abort the data transfer.
             promise.fail(error)
-            return
+            return . failPromise(promise, error: error)
         }
         guard case .copyingData(.readyToSend) = self.state else {
             preconditionFailure("Not ready to send data")
         }
         if channelIsWritable {
-            promise.succeed()
-            return
+            return .succeedPromise(promise)
         }
         return avoidingStateMachineCoW { state in
             state = .copyingData(.pendingBackpressureRelieve(promise))
+            return .none
         }
     }
 

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -181,7 +181,15 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             promise.fail(PostgresError.connectionClosed)
             return
         }
-        self.state.checkBackendCanReceiveCopyData(channelIsWritable: handlerContext.channel.isWritable, promise: promise)
+        let action = self.state.checkBackendCanReceiveCopyData(channelIsWritable: handlerContext.channel.isWritable, promise: promise)
+        switch action {
+        case .none:
+            break
+        case .succeedPromise(let promise):
+            promise.succeed()
+        case .failPromise(let promise, error: let error):
+            promise.fail(error)
+}
     }
 
     /// Cancel the currently executing operation, if it is cancellable.