pg-protocol: defer DataRow cell decoding to a lazy fields getter

packages/pg-protocol/src/b-data-row.ts

@@ -0,0 +1,157 @@
+// file for microbenchmarking DataRow parsing
+//
+//   node dist/b-data-row.js
+//
+// Synthesizes a large text result set, hands it to the parser in socket-sized chunks, and times
+// the three ways a consumer can treat a row: never reading a cell, reading the `fields` array
+// once (what pg's `Result.parseRow` does), and touching every cell.
+
+import { PerformanceObserver } from 'perf_hooks'
+import { BackendMessage } from './messages'
+import { Parser } from './parser'
+
+const ROWS = 100_000
+const COLUMNS = 40
+const CHUNK_SIZE = 64 * 1024
+const RUNS = 5
+// width of the synthesized id column, so every row's frame is the same size
+const ID_WIDTH = 7
+
+// the parts of a DataRow message this benchmark consumes, i.e. without depending on its class
+type Row = {
+  fieldCount: number
+  fields: (string | null)[]
+}
+
+type Mode = {
+  name: string
+  /** Returns a number, so that the work of consuming a row cannot be optimized away. */
+  consume: (row: Row) => number
+}
+
+const modes: Mode[] = [
+  // a consumer that never reads a cell, e.g. one decoding lazily on its own
+  { name: 'fieldCount only', consume: (row) => row.fieldCount },
+  // one `fields` read per row, i.e. what pg itself does
+  { name: 'fields array', consume: (row) => row.fields.length },
+  // every cell read, i.e. what building a row object ends up doing
+  { name: 'every cell', consume: sumCellLengths },
+]
+
+main()
+
+/** Times each consumption mode over the same synthesized chunks. */
+async function main(): Promise<void> {
+  const chunks = synthesizeChunks(ROWS, COLUMNS, CHUNK_SIZE)
+  const bytes = chunks.reduce((total, chunk) => total + chunk.byteLength, 0)
+  const mib = (bytes / 1024 / 1024).toFixed(1)
+  console.log(`${ROWS} rows x ${COLUMNS} cells = ${mib} MiB in ${chunks.length} chunks of ${CHUNK_SIZE} bytes`)
+
+  for (const mode of modes) {
+    let checksum = run(chunks, mode) // warm up, uncounted
+    let collections = 0
+    let pause = 0
+    const observer = new PerformanceObserver((list) => {
+      for (const entry of list.getEntries()) {
+        collections++
+        pause += entry.duration
+      }
+    })
+    observer.observe({ entryTypes: ['gc'] })
+    const times: number[] = []
+    for (let i = 0; i < RUNS; i++) {
+      const start = performance.now()
+      checksum += run(chunks, mode)
+      times.push(performance.now() - start)
+      // let the observer's callbacks run, so this run's collections are counted
+      await new Promise((resolve) => setImmediate(resolve))
+    }
+    observer.disconnect()
+    times.sort((a, b) => a - b)
+    const best = times[0].toFixed(1)
+    const median = times[Math.floor(times.length / 2)].toFixed(1)
+    // gc counts are an allocation-pressure proxy: a scavenge means another semi-space was filled
+    const gc = `${(collections / RUNS).toFixed(1)} gc/run, ${(pause / RUNS).toFixed(1)}ms gc pause/run`
+    console.log(`  ${mode.name.padEnd(16)} best ${best}ms  median ${median}ms  ${gc}  (checksum ${checksum})`)
+  }
+}
+
+/** Parses every chunk with a fresh parser, returning the consumer's checksum. */
+function run(chunks: Buffer[], mode: Mode): number {
+  const parser = new Parser()
+  let checksum = 0
+  const consume = (msg: BackendMessage) => {
+    checksum += mode.consume(msg as unknown as Row)
+  }
+  for (const chunk of chunks) {
+    parser.parse(chunk, consume)
+  }
+  return checksum
+}
+
+/** Reads every cell of a row, i.e. the most a consumer can ask of `fields`. */
+function sumCellLengths(row: Row): number {
+  const { fields } = row
+  let total = 0
+  for (let i = 0; i < fields.length; i++) {
+    const cell = fields[i]
+    total += cell === null ? 0 : cell.length
+  }
+  return total
+}
+
+/** Builds `rows` DataRow frames and slices them into standalone reads, i.e. as a socket delivers them. */
+function synthesizeChunks(rows: number, columns: number, chunkSize: number): Buffer[] {
+  const cells = sampleCells(columns)
+  const frameLength = 7 + cells.reduce((total, cell) => total + 4 + (cell === null ? 0 : Buffer.byteLength(cell)), 0)
+  const stream = Buffer.allocUnsafe(rows * frameLength)
+  let offset = 0
+  for (let row = 0; row < rows; row++) {
+    cells[0] = String(row).padStart(ID_WIDTH, '0')
+    offset = writeDataRow(stream, offset, cells)
+  }
+  const chunks: Buffer[] = []
+  for (let start = 0; start < offset; start += chunkSize) {
+    // a copy, so each chunk is its own allocation like the ones net.Socket emits
+    chunks.push(Buffer.from(stream.subarray(start, Math.min(start + chunkSize, offset))))
+  }
+  return chunks
+}
+
+/** A row's worth of representative text values, i.e. ids, timestamps, numbers, prose and NULLs. */
+function sampleCells(columns: number): (string | null)[] {
+  const samples = [
+    '0'.repeat(ID_WIDTH),
+    '4c1a0e2e-6b58-4e2f-9b6f-6a1d0f7f1b23',
+    '2026-07-26 11:22:33.456789',
+    '12345',
+    'true',
+    null,
+    '199.95',
+    'a short description of the row',
+  ]
+  const cells: (string | null)[] = new Array(columns)
+  for (let i = 0; i < columns; i++) {
+    cells[i] = samples[i % samples.length]
+  }
+  return cells
+}
+
+/** Writes one DataRow frame at `offset`, returning the offset just past it. */
+function writeDataRow(into: Buffer, offset: number, cells: (string | null)[]): number {
+  const start = offset
+  into[offset] = 0x44 // 'D'
+  offset += 5 // the int32 length is backfilled once the frame is written
+  offset = into.writeInt16BE(cells.length, offset)
+  for (const cell of cells) {
+    if (cell === null) {
+      offset = into.writeInt32BE(-1, offset)
+    } else {
+      offset = into.writeInt32BE(Buffer.byteLength(cell), offset)
+      offset += into.write(cell, offset, 'utf-8')
+    }
+  }
+  // the frame's length covers everything but the code byte
+  into.writeInt32BE(offset - start - 1, start + 1)
+  return offset
+}

packages/pg-protocol/src/inbound-parser.test.ts

@@ -3,7 +3,7 @@ import BufferList from './testing/buffer-list'
 import { parse } from '.'
 import assert from 'assert'
 import { PassThrough } from 'stream'
-import { BackendMessage } from './messages'
+import { BackendMessage, DataRowMessage, Field } from './messages'
 import { Parser } from './parser'
 
 const authOkBuffer = buffers.authenticationOk()
@@ -305,6 +305,40 @@ describe('PgPacketStream', function () {
         fields: ['test'],
       })
     })
+
+    describe('parsing data row with an empty, a multi-byte and a null field', function () {
+      testForMessage(buffers.dataRow(['', 'héllo → 🐘', null, 'ok']), {
+        name: 'dataRow',
+        fieldCount: 4,
+        fields: ['', 'héllo → 🐘', null, 'ok'],
+      })
+    })
+
+    it('decodes fields lazily, and only once', async function () {
+      const [message] = (await parseBuffers([buffers.dataRow(['a', null, 'b'])])) as DataRowMessage[]
+      // fieldCount is known without decoding any cell
+      assert.strictEqual(message.fieldCount, 3)
+      const fields = message.fields
+      assert.deepEqual(fields, ['a', null, 'b'])
+      // the decoded array is memoized, so repeated reads do not re-decode
+      assert.strictEqual(message.fields, fields)
+    })
+
+    it('exposes the undecoded payload', async function () {
+      const [message] = (await parseBuffers([buffers.dataRow(['🐘'])])) as DataRowMessage[]
+      // offset points at the int16 field count, i.e. past the 1 byte code and int32 length
+      assert.strictEqual(message.offset, 5)
+      assert.strictEqual(message.bytes.readInt16BE(message.offset), 1)
+      const len = message.bytes.readInt32BE(message.offset + 2)
+      assert.strictEqual(len, 4)
+      assert.deepEqual(message.bytes.subarray(message.offset + 6, message.offset + 6 + len), Buffer.from('🐘'))
+    })
+
+    it('accepts pre-decoded fields for backwards compatibility', function () {
+      const message = new DataRowMessage(11, ['a', null])
+      assert.strictEqual(message.fieldCount, 2)
+      assert.deepEqual(message.fields, ['a', null])
+    })
   })
 
   describe('notice message', function () {
@@ -521,13 +555,11 @@ describe('PgPacketStream', function () {
 
     const verifyMessages = function (messages: any[]) {
       assert.strictEqual(messages.length, 2)
-      assert.deepEqual(messages[0], {
-        name: 'dataRow',
-        fieldCount: 1,
-        length: 11,
-        fields: ['!'],
-      })
-      assert.equal(messages[0].fields[0], '!')
+      // dataRow decodes its cells lazily, so assert its shape rather than its own properties
+      assert.strictEqual(messages[0].name, 'dataRow')
+      assert.strictEqual(messages[0].length, 11)
+      assert.strictEqual(messages[0].fieldCount, 1)
+      assert.deepEqual(messages[0].fields, ['!'])
       assert.deepEqual(messages[1], {
         name: 'readyForQuery',
         length: 5,
@@ -567,9 +599,125 @@ describe('PgPacketStream', function () {
     })
   })
 
+  // every message is parsed out of the chunk it arrived in, except the one message per chunk that
+  // straddles a chunk boundary, so exercise every way a stream can be cut up
+  describe('chunked parsing', function () {
+    const stream = Buffer.concat([
+      buffers.rowDescription([
+        {
+          name: 'id',
+          tableID: 1,
+          attributeNumber: 1,
+          dataTypeID: 23,
+          dataTypeSize: 4,
+          typeModifier: -1,
+          formatCode: 0,
+        },
+        {
+          name: 'name',
+          tableID: 1,
+          attributeNumber: 2,
+          dataTypeID: 25,
+          dataTypeSize: -1,
+          typeModifier: -1,
+          formatCode: 0,
+        },
+      ]),
+      buffers.dataRow(['one', null, 'héllo → 🐘']),
+      buffers.dataRow([]),
+      // longer than the smaller chunk sizes below, so that one message spans several chunks
+      buffers.dataRow(['x'.repeat(300)]),
+      buffers.commandComplete('SELECT 3'),
+      buffers.readyForQuery(),
+    ])
+
+    const expected = [
+      'rowDescription(id,name)',
+      'dataRow(one,NULL,héllo → 🐘)',
+      'dataRow()',
+      `dataRow(${'x'.repeat(300)})`,
+      'commandComplete(SELECT 3)',
+      'readyForQuery(5)',
+    ]
+
+    it('parses the whole stream in one chunk', function () {
+      assert.deepEqual(parseChunks([stream]), expected)
+    })
+
+    it('parses the stream split at every offset', function () {
+      for (let split = 1; split < stream.byteLength; split++) {
+        const chunks = [copyOf(stream, 0, split), copyOf(stream, split, stream.byteLength)]
+        assert.deepEqual(parseChunks(chunks), expected, `split at ${split}`)
+      }
+    })
+
+    it('parses a message far larger than the chunks carrying it', function () {
+      // big enough that reassembling it has to grow its buffer several times
+      const cell = 'z'.repeat(3 * 1024 * 1024)
+      const bigStream = Buffer.concat([buffers.dataRow([cell, 'small']), buffers.readyForQuery()])
+      const summaries = parseChunks(intoChunks(bigStream, 64 * 1024))
+      assert.deepEqual(summaries, [`dataRow(${cell},small)`, 'readyForQuery(5)'])
+    })
+
+    it('parses the stream in fixed size chunks', function () {
+      for (const size of [1, 2, 3, 4, 5, 7, 13, 64, 128]) {
+        assert.deepEqual(parseChunks(intoChunks(stream, size)), expected, `${size} byte chunks`)
+      }
+    })
+
+    it('keeps messages valid after later chunks are parsed', function () {
+      // messages are views over the bytes they were parsed from, so hold them all and read them
+      // only once the whole stream has gone through the parser
+      const parser = new Parser()
+      const messages: BackendMessage[] = []
+      for (const chunk of intoChunks(stream, 7)) {
+        parser.parse(chunk, (msg) => messages.push(msg))
+      }
+      assert.deepEqual(messages.map(summarize), expected)
+    })
+  })
+
   it('cleans up the reader after handling a packet', function () {
     const parser = new Parser()
     parser.parse(oneFieldBuf, () => {})
     assert.strictEqual((parser as any).reader.buffer.byteLength, 0)
   })
 })
+
+/** Parses `chunks` with one parser, summarizing each message as it is emitted. */
+const parseChunks = function (chunks: Buffer[]): string[] {
+  const parser = new Parser()
+  const summaries: string[] = []
+  for (const chunk of chunks) {
+    parser.parse(chunk, (msg) => summaries.push(summarize(msg)))
+  }
+  return summaries
+}
+
+/** A message's contents as a string, i.e. `dataRow(one,NULL)`, so a whole stream can be compared at once. */
+const summarize = function (msg: any): string {
+  switch (msg.name) {
+    case 'dataRow':
+      return `dataRow(${msg.fields.map((field: string | null) => (field === null ? 'NULL' : field)).join(',')})`
+    case 'rowDescription':
+      return `rowDescription(${msg.fields.map((field: Field) => field.name).join(',')})`
+    case 'commandComplete':
+      return `commandComplete(${msg.text})`
+    default:
+      return `${msg.name}(${msg.length})`
+  }
+}
+
+/** Cuts `buffer` into standalone `size` byte chunks, i.e. the reads a socket would deliver. */
+const intoChunks = function (buffer: Buffer, size: number): Buffer[] {
+  const chunks: Buffer[] = []
+  for (let offset = 0; offset < buffer.byteLength; offset += size) {
+    chunks.push(copyOf(buffer, offset, Math.min(offset + size, buffer.byteLength)))
+  }
+  return chunks
+}
+
+/** Copies `buffer` between the given offsets, i.e. so each chunk is its own allocation. */
+const copyOf = function (buffer: Buffer, start: number, end: number): Buffer {
+  return Buffer.from(buffer.subarray(start, end))
+}