[Fix] Query Hangs if Connection is Closed

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -222,7 +222,9 @@ public final class PostgresConnection: @unchecked Sendable {
             promise: promise
         )
 
-        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(HandlerTask.extendedQuery(context), promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
 
         return promise.futureResult
     }
@@ -239,7 +241,10 @@ public final class PostgresConnection: @unchecked Sendable {
             promise: promise
         )
 
-        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(HandlerTask.extendedQuery(context), promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
+
         return promise.futureResult.map { rowDescription in
             PSQLPreparedStatement(name: name, query: query, connection: self, rowDescription: rowDescription)
         }
@@ -255,15 +260,21 @@ public final class PostgresConnection: @unchecked Sendable {
             logger: logger,
             promise: promise)
 
-        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(HandlerTask.extendedQuery(context), promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
+
         return promise.futureResult
     }
 
     func close(_ target: CloseTarget, logger: Logger) -> EventLoopFuture<Void> {
         let promise = self.channel.eventLoop.makePromise(of: Void.self)
         let context = CloseCommandContext(target: target, logger: logger, promise: promise)
 
-        self.channel.write(HandlerTask.closeCommand(context), promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(HandlerTask.closeCommand(context), promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
+
         return promise.futureResult
     }
 
@@ -426,7 +437,9 @@ extension PostgresConnection {
             promise: promise
         )
 
-        self.channel.write(HandlerTask.extendedQuery(context), promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(HandlerTask.extendedQuery(context), promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
 
         do {
             return try await promise.futureResult.map({ $0.asyncSequence() }).get()
@@ -455,7 +468,11 @@ extension PostgresConnection {
 
                 let task = HandlerTask.startListening(listener)
 
-                self.channel.write(task, promise: nil)
+                let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+                self.channel.write(task, promise: writePromise)
+                writePromise.futureResult.whenFailure { error in
+                    continuation.resume(throwing: error)
+                }
             }
         } onCancel: {
             let task = HandlerTask.cancelListening(channel, id)
@@ -480,7 +497,11 @@ extension PostgresConnection {
             logger: logger,
             promise: promise
         ))
-        self.channel.write(task, promise: nil)
+
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(task, promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
+
         do {
             return try await promise.futureResult
                 .map { $0.asyncSequence() }
@@ -515,7 +536,11 @@ extension PostgresConnection {
             logger: logger,
             promise: promise
         ))
-        self.channel.write(task, promise: nil)
+
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(task, promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: promise)
+
         do {
             return try await promise.futureResult
                 .map { $0.commandTag }
@@ -674,7 +699,7 @@ internal enum PostgresCommands: PostgresRequest {
 
 /// Context for receiving NotificationResponse messages on a connection, used for PostgreSQL's `LISTEN`/`NOTIFY` support.
 public final class PostgresListenContext: Sendable {
-    private let promise: EventLoopPromise<Void>
+    let promise: EventLoopPromise<Void>
 
     var future: EventLoopFuture<Void> {
         self.promise.futureResult
@@ -714,7 +739,9 @@ extension PostgresConnection {
         )
 
         let task = HandlerTask.startListening(listener)
-        self.channel.write(task, promise: nil)
+        let writePromise = self.channel.eventLoop.makePromise(of: Void.self)
+        self.channel.write(task, promise: writePromise)
+        writePromise.futureResult.cascadeFailure(to: listenContext.promise)
 
         listenContext.future.whenComplete { _ in
             let task = HandlerTask.cancelListening(channel, id)

Tests/IntegrationTests/PSQLIntegrationTests.swift

@@ -360,4 +360,22 @@ final class IntegrationTests: XCTestCase {
         }
     }
 
+    func testConnectionClosureMidQueryDoesNotHang() async throws {
+        _ = await withThrowingTaskGroup(of: Void.self) { taskGroup in
+            for _ in (0 ..< 1_000) {
+                taskGroup.addTask {
+                    let conn = try await PostgresConnection.test(
+                        on: NIOSingletons.posixEventLoopGroup.next()
+                    ).get()
+
+                    async let close: () = conn.closeGracefully()
+                    async let query = conn.query("SELECT 1", logger: .psqlTest)
+
+                    _ = try await (close, query)
+                }
+            }
+
+            // Ignore failures
+        }
+    }
 }