fix(pg-cloudflare): batch extended query protocol writes

packages/pg-cloudflare/src/index.ts

@@ -1,6 +1,8 @@
 import { SocketOptions, Socket, TlsOptions } from 'cloudflare:sockets'
 import { EventEmitter } from 'events'
 
+type WriteCallback = (error?: unknown) => void
+
 /**
  * Wrapper around the Cloudflare built-in socket that can be used by the `Connection`.
  */
@@ -13,6 +15,8 @@ export class CloudflareSocket extends EventEmitter {
   private _cfSocket: Socket | null = null
   private _cfWriter: WritableStreamDefaultWriter | null = null
   private _cfReader: ReadableStreamDefaultReader | null = null
+  private _corked = 0
+  private _corkedWrites: { data: Uint8Array; callback: WriteCallback }[] = []
 
   constructor(readonly ssl: boolean) {
     super()
@@ -31,6 +35,26 @@ export class CloudflareSocket extends EventEmitter {
     return this
   }
 
+  cork() {
+    this._corked++
+  }
+
+  uncork() {
+    if (this._corked === 0) return
+    this._corked--
+    if (this._corked > 0) return
+
+    const writes = this._corkedWrites
+    this._corkedWrites = []
+    if (writes.length === 0) return
+
+    const data = writes.length === 1 ? writes[0].data : Buffer.concat(writes.map((write) => write.data))
+    const sent = data.length === 0 ? Promise.resolve() : this._write(data)
+    for (const { callback } of writes) {
+      sent.then(() => callback(), callback)
+    }
+  }
+
   async connect(port: number, host: string, connectListener?: (...args: unknown[]) => void) {
     try {
       log('connecting')
@@ -82,32 +106,38 @@ export class CloudflareSocket extends EventEmitter {
     this.emit('data', Buffer.from(value))
   }
 
-  write(data: Uint8Array | string, callback?: (error?: unknown) => void): true | void
-  write(data: Uint8Array | string, encoding?: BufferEncoding, callback?: (error?: unknown) => void): true | void
+  write(data: Uint8Array | string, callback?: WriteCallback): true | void
+  write(data: Uint8Array | string, encoding?: BufferEncoding, callback?: WriteCallback): true | void
   write(
     data: Uint8Array | string,
-    encodingOrCallback: BufferEncoding | ((error?: unknown) => void) = 'utf8',
-    callback: (error?: unknown) => void = () => {}
+    encodingOrCallback: BufferEncoding | WriteCallback = 'utf8',
+    callback: WriteCallback = () => {}
   ): true | void {
     const encoding = typeof encodingOrCallback === 'function' ? 'utf8' : encodingOrCallback
     if (typeof encodingOrCallback === 'function') callback = encodingOrCallback
-    if (data.length === 0) return callback()
     if (typeof data === 'string') data = Buffer.from(data, encoding)
+    if (this._corked > 0) {
+      this._corkedWrites.push({ data, callback })
+      return true
+    }
+    if (data.length === 0) return callback()
 
-    log('sending data direct:', data)
-    this._cfWriter!.write(data).then(
-      () => {
-        log('data sent')
-        callback()
-      },
-      (err) => {
-        log('send error', err)
-        callback(err)
-      }
-    )
+    this._write(data).then(() => callback(), callback)
     return true
   }
 
+  private _write(data: Uint8Array) {
+    log('sending data direct:', data)
+    const sent = this._cfWriter!.write(data)
+    if (debug) {
+      sent.then(
+        () => log('data sent'),
+        (err) => log('send error', err)
+      )
+    }
+    return sent
+  }
+
   end(data = Buffer.alloc(0), encoding: BufferEncoding = 'utf8', callback: (...args: unknown[]) => void = () => {}) {
     log('ending CF socket')
     this.write(data, encoding, (err) => {
@@ -123,12 +153,24 @@ export class CloudflareSocket extends EventEmitter {
         .catch((e) => this.emit('error', e))
       if (callback) callback(err)
     })
+    if (this._corked > 0) {
+      // Ending flushes even when cork() was called more than once.
+      this._corked = 1
+      this.uncork()
+    }
     return this
   }
 
   destroy(reason: string) {
     log('destroying CF socket', reason)
     this.destroyed = true
+    this._corked = 0
+    const writes = this._corkedWrites
+    this._corkedWrites = []
+    const error = reason ?? new Error('Socket destroyed before buffered writes could be sent')
+    for (const { callback } of writes) {
+      Promise.resolve().then(() => callback(error))
+    }
     return this.end()
   }
 

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

@@ -1,5 +1,7 @@
 import assert from 'assert'
+import { once } from 'events'
 import { describe, it } from 'node:test'
+import pg from 'pg'
 import { CloudflareSocket } from 'pg-cloudflare'
 
 describe('pg-cloudflare', () => {
@@ -38,6 +40,120 @@ describe('pg-cloudflare', () => {
     await promise
   })
 
+  it('writes the extended query protocol messages together without changing the bytes', () => {
+    const socket = new CloudflareSocket()
+    const writes = []
+    socket.writable = true
+    socket._cfWriter = {
+      write(data) {
+        writes.push(data)
+        return Promise.resolve()
+      },
+    }
+
+    const separateWrites = []
+    const plainStream = { writable: true, write: (data) => separateWrites.push(data) }
+    const query = { text: 'select $1::int', values: [42] }
+    assert.equal(new pg.Query(query).submit(new pg.Connection({ stream: plainStream })), null)
+    assert.equal(new pg.Query(query).submit(new pg.Connection({ stream: socket })), null)
+
+    assert.equal(separateWrites.length, 5)
+    assert.deepStrictEqual(writes, [Buffer.concat(separateWrites)])
+  })
+
+  it('flushes on the last uncork and completes buffered callbacks in write order', async () => {
+    const socket = new CloudflareSocket()
+    const writes = []
+    const completed = []
+    socket._cfWriter = {
+      write(data) {
+        writes.push(data)
+        return Promise.resolve()
+      },
+    }
+
+    socket.cork()
+    socket.cork()
+    const finished = new Promise((resolve) => {
+      socket.write(Uint8Array.of(0xaa), (err) => completed.push([1, err]))
+      socket.write('ff', 'hex', (err) => {
+        completed.push([2, err])
+        resolve()
+      })
+    })
+    socket.uncork()
+    assert.deepStrictEqual(writes, [])
+    assert.deepStrictEqual(completed, [])
+
+    socket.uncork()
+    socket.uncork()
+    await finished
+    assert.deepStrictEqual(writes, [Buffer.from([0xaa, 0xff])])
+    assert.deepStrictEqual(completed, [
+      [1, undefined],
+      [2, undefined],
+    ])
+  })
+
+  it('reports a failed batch to each buffered write callback', async () => {
+    const socket = new CloudflareSocket()
+    const error = new Error('write failed')
+    socket._cfWriter = { write: () => Promise.reject(error) }
+
+    socket.cork()
+    const first = new Promise((resolve) => socket.write('first', resolve))
+    const second = new Promise((resolve) => socket.write('second', resolve))
+    socket.uncork()
+    assert.deepStrictEqual(await Promise.all([first, second]), [error, error])
+  })
+
+  it('flushes buffered data before closing, even without an end payload', async () => {
+    const socket = new CloudflareSocket()
+    let finishWrite
+    const writes = []
+    let closed = false
+    socket._cfWriter = {
+      write(data) {
+        writes.push(data)
+        return new Promise((resolve) => {
+          finishWrite = resolve
+        })
+      },
+    }
+    socket._cfSocket = {
+      close: async () => {
+        closed = true
+      },
+    }
+
+    socket.cork()
+    socket.cork()
+    socket.write('pending')
+    const close = once(socket, 'close')
+    socket.end()
+    assert.deepStrictEqual(writes, [Buffer.from('pending')])
+    assert.equal(closed, false)
+    finishWrite()
+    await close
+    assert.equal(closed, true)
+  })
+
+  it('discards buffered data on destroy and completes its callbacks with the reason', async () => {
+    const socket = new CloudflareSocket()
+    socket._cfWriter = { write: () => assert.fail('destroy must not send buffered data') }
+    socket._cfSocket = { close: async () => {} }
+
+    const reason = new Error('socket destroyed')
+    socket.cork()
+    const result = new Promise((resolve) => socket.write('pending', resolve))
+    const close = once(socket, 'close')
+    socket.destroy(reason)
+    socket.uncork()
+
+    assert.equal(await result, reason)
+    await close
+  })
+
   it('should emit close when ending a socket whose closed promise never settles', async () => {
     const socket = new CloudflareSocket()
     socket._cfSocket = {