fix(pg-cloudflare): prevent Client.end() from hanging (#3752)

* fix(pg-cloudflare): safely end closed sockets

* fix(pg-cloudflare): complete socket shutdown handling

* test(pg-cloudflare): retain write callback coverage

Author: pratik-desgn

packages/pg-cloudflare/src/index.ts

@@ -111,7 +111,16 @@ export class CloudflareSocket extends EventEmitter {
   end(data = Buffer.alloc(0), encoding: BufferEncoding = 'utf8', callback: (...args: unknown[]) => void = () => {}) {
     log('ending CF socket')
     this.write(data, encoding, (err) => {
-      this._cfSocket?.close()
+      const socket = this._cfSocket
+      const closePromise = socket?.close()
+      closePromise
+        ?.then(() => {
+          if (this._cfSocket === socket) {
+            this._cfSocket = null
+            this.emit('close')
+          }
+        })
+        .catch((e) => this.emit('error', e))
       if (callback) callback(err)
     })
     return this
@@ -140,16 +149,21 @@ export class CloudflareSocket extends EventEmitter {
   }
 
   _addClosedHandler() {
-    this._cfSocket!.closed.then(() => {
-      if (!this._upgrading) {
-        log('CF socket closed')
-        this._cfSocket = null
-        this.emit('close')
-      } else {
-        this._upgrading = false
-        this._upgraded = true
-      }
-    }).catch((e) => this.emit('error', e))
+    const socket = this._cfSocket!
+    socket.closed
+      .then(() => {
+        if (!this._upgrading) {
+          log('CF socket closed')
+          if (this._cfSocket === socket) {
+            this._cfSocket = null
+            this.emit('close')
+          }
+        } else {
+          this._upgrading = false
+          this._upgraded = true
+        }
+      })
+      .catch((e) => this.emit('error', e))
   }
 }
 

packages/pg-esm-test/pg-cloudflare.test.js

@@ -37,4 +37,22 @@ describe('pg-cloudflare', () => {
 
     await promise
   })
+
+  it('should emit close when ending a socket whose closed promise never settles', async () => {
+    const socket = new CloudflareSocket()
+    socket._cfSocket = {
+      closed: new Promise(() => {}),
+      close: async () => {},
+    }
+    socket._addClosedHandler()
+
+    await new Promise((resolve, reject) => {
+      const timer = setTimeout(() => reject(new Error('close event was not emitted')), 100)
+      socket.once('close', () => {
+        clearTimeout(timer)
+        resolve()
+      })
+      socket.end()
+    })
+  })
 })