Avoid storing ChannelHandlerContext in PostgresCopyFromWriter

ahoppen · ahoppen · commit 8138f431569a · 2025-07-04T08:12:15.000+02:00
diff --git a/Sources/PostgresNIO/Connection/PostgresConnection.swift b/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()
             }
         }
     }
diff --git a/Sources/PostgresNIO/New/PostgresChannelHandler.swift b/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()

Original file line number	Diff line number	Diff line change
`@@ -710,12 +710,10 @@ public struct PostgresCopyFromWriter: Sendable {`
`710`	`710`	`}`
`711`	`711`
`712`	`712`	`private let channelHandler: NIOLoopBound<PostgresChannelHandler>`
`713`		`- private let context: NIOLoopBound<ChannelHandlerContext>`
`714`	`713`	`private let eventLoop: any EventLoop`
`715`	`714`
`716`		`- init(handler: PostgresChannelHandler, context: ChannelHandlerContext, eventLoop: any EventLoop) {`
	`715`	`+ init(handler: PostgresChannelHandler, eventLoop: any EventLoop) {`
`717`	`716`	`self.channelHandler = NIOLoopBound(handler, eventLoop: eventLoop)`
`718`		`- self.context = NIOLoopBound(context, eventLoop: eventLoop)`
`719`	`717`	`self.eventLoop = eventLoop`
`720`	`718`	`}`
`721`	`719`
`@@ -728,20 +726,20 @@ public struct PostgresCopyFromWriter: Sendable {`
`728`	`726`	`// error during the data transfer and thus cannot process any more data.`
`729`	`727`	`try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in`
`730`	`728`	`if eventLoop.inEventLoop {`
`731`		`- self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)`
	`729`	`+ self.channelHandler.value.waitForWritableBuffer(continuation)`
`732`	`730`	`} else {`
`733`	`731`	`eventLoop.execute {`
`734`		`- self.channelHandler.value.waitForWritableBuffer(context: self.context.value, continuation)`
	`732`	`+ self.channelHandler.value.waitForWritableBuffer(continuation)`
`735`	`733`	`}`
`736`	`734`	`}`
`737`	`735`	`}`
`738`	`736`
`739`	`737`	`// Run the actual data transfer`
`740`	`738`	`if eventLoop.inEventLoop {`
`741`		`- self.channelHandler.value.copyData(byteBuffer, context: self.context.value)`
	`739`	`+ self.channelHandler.value.copyData(byteBuffer)`
`742`	`740`	`} else {`
`743`	`741`	`eventLoop.execute {`
`744`		`- self.channelHandler.value.copyData(byteBuffer, context: self.context.value)`
	`742`	`+ self.channelHandler.value.copyData(byteBuffer)`
`745`	`743`	`}`
`746`	`744`	`}`
`747`	`745`	`}`
`@@ -751,10 +749,10 @@ public struct PostgresCopyFromWriter: Sendable {`
`751`	`749`	`func done() async throws {`
`752`	`750`	`try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in`
`753`	`751`	`if eventLoop.inEventLoop {`
`754`		`- self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)`
	`752`	`+ self.channelHandler.value.sendCopyDone(continuation: continuation)`
`755`	`753`	`} else {`
`756`	`754`	`eventLoop.execute {`
`757`		`- self.channelHandler.value.sendCopyDone(continuation: continuation, context: self.context.value)`
	`755`	`+ self.channelHandler.value.sendCopyDone(continuation: continuation)`
`758`	`756`	`}`
`759`	`757`	`}`
`760`	`758`	`}`
`@@ -765,10 +763,10 @@ public struct PostgresCopyFromWriter: Sendable {`
`765`	`763`	`func failed(error: any Error) async throws {`
`766`	`764`	`try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in`
`767`	`765`	`if eventLoop.inEventLoop {`
`768`		`- self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)`
	`766`	`+ self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation)`
`769`	`767`	`} else {`
`770`	`768`	`eventLoop.execute {`
`771`		`- self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation, context: self.context.value)`
	`769`	`+ self.channelHandler.value.sendCopyFailed(message: "\(error)", continuation: continuation)`
`772`	`770`	`}`
`773`	`771`	`}`
`774`	`772`	`}`
`@@ -777,10 +775,10 @@ public struct PostgresCopyFromWriter: Sendable {`
`777`	`775`	/// Send a `Sync` message to the backend.
`778`	`776`	`func sync() {`
`779`	`777`	`if eventLoop.inEventLoop {`
`780`		`- self.channelHandler.value.sendSync(context: self.context.value)`
	`778`	`+ self.channelHandler.value.sendSync()`
`781`	`779`	`} else {`
`782`	`780`	`eventLoop.execute {`
`783`		`- self.channelHandler.value.sendSync(context: self.context.value)`
	`781`	`+ self.channelHandler.value.sendSync()`
`784`	`782`	`}`
`785`	`783`	`}`
`786`	`784`	`}`