Avoid storing ChannelHandlerContext in PostgresCopyFromWriter

Author: ahoppen

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -710,12 +710,10 @@ public struct PostgresCopyFromWriter: Sendable {
     }
 
     private let channelHandler: NIOLoopBound<PostgresChannelHandler>
-    private let context: NIOLoopBound<ChannelHandlerContext>
     private let eventLoop: any EventLoop
 
-    init(handler: PostgresChannelHandler, context: ChannelHandlerContext, eventLoop: any EventLoop) {
+    init(handler: PostgresChannelHandler, eventLoop: any EventLoop) {
         self.channelHandler = NIOLoopBound(handler, eventLoop: eventLoop)
-        self.context = NIOLoopBound(context, eventLoop: eventLoop)
         self.eventLoop = eventLoop
     }
 
@@ -728,20 +726,20 @@ public struct PostgresCopyFromWriter: Sendable {
         // error during the data transfer and thus cannot process any more data.
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
             if eventLoop.inEventLoop {
-                self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)
+                self.channelHandler.value.waitForWritableBuffer(continuation)
             } else {
                 eventLoop.execute {
-                    self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)
+                    self.channelHandler.value.waitForWritableBuffer(continuation)
                 }
             }
         }
 
         // Run the actual data transfer
         if eventLoop.inEventLoop {
-            self.channelHandler.value.copyData(byteBuffer, context: self.context.value)
+            self.channelHandler.value.copyData(byteBuffer)
         } else {
             eventLoop.execute {
-                self.channelHandler.value.copyData(byteBuffer, context: self.context.value)
+                self.channelHandler.value.copyData(byteBuffer)
             }
         }
     }
@@ -751,10 +749,10 @@ public struct PostgresCopyFromWriter: Sendable {
     func done() async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
             if eventLoop.inEventLoop {
-                self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)
+                self.channelHandler.value.sendCopyDone(continuation: continuation)
             } else {
                 eventLoop.execute {
-                    self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)
+                    self.channelHandler.value.sendCopyDone(continuation: continuation)
                 }
             }
         }
@@ -765,10 +763,10 @@ public struct PostgresCopyFromWriter: Sendable {
     func failed(error: any Error) async throws {
         try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in 
             if eventLoop.inEventLoop {
-                self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)
+                self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation)
             } else {
                 eventLoop.execute {
-                    self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)
+                    self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation)
                 }
             }
         }
@@ -777,10 +775,10 @@ public struct PostgresCopyFromWriter: Sendable {
     /// Send a `Sync` message to the backend.
     func sync() {
         if eventLoop.inEventLoop {
-            self.channelHandler.value.sendSync(context: self.context.value)
+            self.channelHandler.value.sendSync()
         } else {
             eventLoop.execute {
-                self.channelHandler.value.sendSync(context: self.context.value)
+                self.channelHandler.value.sendSync()
             }
         }
     }

Sources/PostgresNIO/New/PostgresChannelHandler.swift

@@ -176,11 +176,11 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
     /// 
     /// This fails the continuation with a `PostgresCopyFromWriter.CopyCancellationError` when the server has cancelled
     /// the data transfer to indicate that the frontend should not send any more data.
-    func waitForWritableBuffer(context: ChannelHandlerContext, _ continuation: CheckedContinuation<Void, any Error>) {
-        let action = self.state.waitForWritableBuffer(channel: context.channel, continuation: continuation)
+    func waitForWritableBuffer(_ continuation: CheckedContinuation<Void, any Error>) {
+        let action = self.state.waitForWritableBuffer(channel: handlerContext!.channel, continuation: continuation)
         switch action {
         case .waitForBackpressureRelieve:
-            context.channel.flush()
+            self.handlerContext!.channel.flush()
         case .resumeContinuation(let continuation):
             continuation.resume()
         case .failContinuation(_, error: let error):
@@ -189,27 +189,27 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
     }
 
     /// Send a `CopyData` message to the backend using the given data.
-    func copyData(_ data: ByteBuffer, context: ChannelHandlerContext) {
+    func copyData(_ data: ByteBuffer) {
         self.encoder.copyData(data: data)
-        context.write(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
+        self.handlerContext!.write(self.wrapOutboundOut(self.encoder.flushBuffer()), 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>, context: ChannelHandlerContext) {
+    func sendCopyDone(continuation: CheckedContinuation<Void, any Error>) {
         let action = self.state.sendCopyDone(continuation: continuation)
-        self.run(action, with: context)
+        self.run(action, with: self.handlerContext!)
     }
 
     /// Put the state machine out of the copying mode and send a `CopyFail` message to the backend.
-    func sendCopyFailed(message: String, continuation: CheckedContinuation<Void, any Error>, context: ChannelHandlerContext) {
+    func sendCopyFailed(message: String, continuation: CheckedContinuation<Void, any Error>) {
         let action = self.state.sendCopyFail(message: message, continuation: continuation)
-        self.run(action, with: context)
+        self.run(action, with: self.handlerContext!)
     }
 
     /// Send a `Sync` message to the backend.
-    func sendSync(context: ChannelHandlerContext) {
+    func sendSync() {
         self.encoder.sync()
-        context.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
+        self.handlerContext!.writeAndFlush(self.wrapOutboundOut(self.encoder.flushBuffer()), promise: nil)
     }
 
     func channelReadComplete(context: ChannelHandlerContext) {
@@ -423,7 +423,7 @@ final class PostgresChannelHandler: ChannelDuplexHandler {
             }
             continuation.resume(throwing: error)
         case .triggerCopyData(let triggerCopy):
-            let writer = PostgresCopyFromWriter(handler: self, context: context, eventLoop: eventLoop)
+            let writer = PostgresCopyFromWriter(handler: self, eventLoop: eventLoop)
             triggerCopy.resume(returning: writer)
         case .sendCopyDone:
             self.encoder.copyDone()