[Fix] Query Hangs if Connection is Closed

Sources/PostgresNIO/Connection/PostgresConnection.swift

@@ -239,7 +239,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)
         }
@@ -426,7 +429,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()

Tests/IntegrationTests/PSQLIntegrationTests.swift

@@ -360,4 +360,28 @@ final class IntegrationTests: XCTestCase {
         }
     }
 
+    func testConnectionClosureMidQueryDoesNotHang() async throws {
+        let badQuery: PostgresQuery = """
+        SELECT * FROM non_existent_table
+        """
+
+        _ = try 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(badQuery, logger: .psqlTest)
+
+                    _ = try await (close, query)
+                }
+            }
+
+            /// Ignore failures
+        }
+    }
 }