Address review comments

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -696,26 +696,12 @@ extension PostgresConnection {
 
 // MARK: Copy from
 
-fileprivate extension EventLoop {
-    /// If we are on the given event loop, execute `task` immediately. Otherwise schedule it for execution.
-    func executeImmediatelyOrSchedule(_ task: @Sendable @escaping () -> Void) {
-        if inEventLoop {
-            return task()
-        }
-        return execute(task)
-    }
-}
-
 /// A handle to send 
 public struct PostgresCopyFromWriter: Sendable {
     private let channelHandler: NIOLoopBound<PostgresChannelHandler>
     private let context: NIOLoopBound<ChannelHandlerContext>
     private let eventLoop: any EventLoop
 
-    struct NotWritableError: Error, CustomStringConvertible {
-        var description = "No data must be written to `PostgresCopyFromWriter` after it has sent a CopyDone or CopyFail message, ie. after the closure producing the copy data has finished"
-    }
-
     init(handler: PostgresChannelHandler, context: ChannelHandlerContext, eventLoop: any EventLoop) {
         self.channelHandler = NIOLoopBound(handler, eventLoop: eventLoop)
         self.context = NIOLoopBound(context, eventLoop: eventLoop)
@@ -725,8 +711,12 @@ public struct PostgresCopyFromWriter: Sendable {
     /// Send data for a `COPY ... FROM STDIN` operation to the backend.
     public func write(_ byteBuffer: ByteBuffer) async throws {
         await withCheckedContinuation { (continuation: CheckedContinuation<Void, Never>) in 
-            eventLoop.executeImmediatelyOrSchedule {
+            if eventLoop.inEventLoop {
                 self.channelHandler.value.copyData(byteBuffer, context: self.context.value, readyForMoreWriteContinuation: continuation)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.copyData(byteBuffer, context: self.context.value, readyForMoreWriteContinuation: continuation)
+                }
             }
         }
     }
@@ -735,8 +725,12 @@ public struct PostgresCopyFromWriter: Sendable {
     /// the backend.
     func done() async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
-            eventLoop.executeImmediatelyOrSchedule {
+            if eventLoop.inEventLoop {
                 self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)
+                }
             }
         }
     }
@@ -745,8 +739,12 @@ public struct PostgresCopyFromWriter: Sendable {
     /// the backend.
     func failed(error: any Error) async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
-            eventLoop.executeImmediatelyOrSchedule {
+            if eventLoop.inEventLoop {
                 self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)
+            } else {
+                eventLoop.execute {
+                    self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)
+                }
             }
         }
     }

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

@@ -122,6 +122,11 @@ struct ConnectionStateMachine {
         case succeedClose(CloseCommandContext)
         case failClose(CloseCommandContext, with: PSQLError, cleanupContext: CleanUpContext?)
     }
+
+    enum ChannelWritabilityChangedAction {
+        case none
+        case resumeContinuation(CheckedContinuation<Void, Never>)
+    }
 
     private var state: State
     private let requireBackendKeyData: Bool
@@ -679,13 +684,14 @@ struct ConnectionStateMachine {
         }
     }
 
-    mutating func channelWritabilityChanged(isWritable: Bool) {
+    mutating func channelWritabilityChanged(isWritable: Bool) -> ChannelWritabilityChangedAction {
         guard case .extendedQuery(var queryState, let connectionContext) = state else {
-            return
+            return .none
         }
         self.state = .modifying // avoid CoW
-        queryState.channelWritabilityChanged(isWritable: isWritable)
+        let action = queryState.channelWritabilityChanged(isWritable: isWritable)
         self.state = .extendedQuery(queryState, connectionContext)
+        return action
     }
 
     // MARK: - Running Queries -

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

@@ -598,13 +598,13 @@ struct ExtendedQueryStateMachine {
         }
     }
 
-    mutating func channelWritabilityChanged(isWritable: Bool) {
+    mutating func channelWritabilityChanged(isWritable: Bool) -> ConnectionStateMachine.ChannelWritabilityChangedAction {
         guard case .copyingData(.pendingBackpressureRelieve(let continuation)) = state else {
-            return
+            return .none
         }
-        self.avoidingStateMachineCoW { state in
+        return self.avoidingStateMachineCoW { state in
             state = .copyingData(.readyToSend)
-            continuation.resume()
+            return .resumeContinuation(continuation)
         }
     }
 

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -203,7 +203,13 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
     }
 
     func channelWritabilityChanged(context: ChannelHandlerContext) {
-        self.state.channelWritabilityChanged(isWritable: context.channel.isWritable)
+        let action = self.state.channelWritabilityChanged(isWritable: context.channel.isWritable)
+        switch action {
+        case .none: 
+            break
+        case .resumeContinuation(let continuation): 
+            continuation.resume()
+        }
     }
 
     func userInboundEventTriggered(context: ChannelHandlerContext, event: Any) {