perf(pg): read a range in one query, and only its bytes

CHANGELOG.md

@@ -10,6 +10,10 @@ this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 
 - `S3ReplicaClient.openLtxFile()` reads a range with a single `GetObject`, without
   the `HeadObject` it sent first, which halves the requests of a paged restore.
+- `PgReplicaClient.openLtxFile()` reads a range in one query instead of two (the
+  file row, then its chunks), and Postgres returns only the bytes of the range
+  instead of every 1 MiB chunk it covers. A read of 128 KiB, the size of a paged
+  restore's read for a page, went from 9.7 ms to 1.4 ms against a local server.
 - The first `sync()` after a restore into the replica's own path (as `ensureExists()` does) uploads
   only the changes made since the restore, instead of a snapshot of the whole database. The restore
   writes the first WAL of the restored database: one commit frame holding page 1 unchanged, under

packages/backend-pg/src/replica-client.ts

@@ -136,7 +136,8 @@ function txid(v: bigint): string {
 
 /** Shape of a row from `ltx_chunks`. */
 interface ChunkRow {
-  seq: number;
+  /** Null when the file has no chunk in the range read. */
+  seq: number | null;
   data: Uint8Array;
 }
 
@@ -346,65 +347,73 @@ export class PgReplicaClient implements ReplicaClient {
     const { level, minTXID, maxTXID, offset, size } = opts;
     const path = this.filePathKey(level, minTXID, maxTXID);
 
-    let fileRow;
-    try {
-      fileRow = await this.sql`
-        SELECT size FROM ${this.sql(this.filesTable())}
-        WHERE prefix = ${this.prefix} AND level = ${level}
-          AND min_txid = ${txid(minTXID)} AND max_txid = ${txid(maxTXID)}
-      `;
-    } catch (err) {
-      throw new LTXError('open', path, level, minTXID, maxTXID, err);
-    }
-    if (fileRow.length === 0) {
-      throw new LTXError('open', path, level, minTXID, maxTXID, new Error('not found'));
-    }
-
-    const fileSize = Number(fileRow[0].size);
-    if (offset >= fileSize) return Readable.from([]);
-    const end = size > 0 ? Math.min(offset + size - 1, fileSize - 1) : fileSize - 1;
-    if (end < offset) return Readable.from([]);
-
-    const firstSeq = Math.floor(offset / this.chunkSize);
-    const lastSeq = Math.floor(end / this.chunkSize);
+    // One query, so one round trip for a range of up to 16 chunks: the file
+    // row, joined with the chunks the range covers, each cut to the bytes of
+    // the range by the server. No row means no file; a row without a chunk, a
+    // range past the end of the file. Chunks are stored uncompressed (LTX
+    // pages already are), so Postgres reads only the part of a chunk asked
+    // for. Offsets are bigints: files can be larger than 2 GiB.
+    const chunkSize = this.chunkSize;
+    const firstSeq = Math.floor(offset / chunkSize);
+    const start = String(offset);
+    const chunkStart = this.sql`c.seq::bigint * ${chunkSize}::bigint`;
+    const from = this.sql`(greatest(${start}::bigint - ${chunkStart}, 0) + 1)::int`;
+    const lastSeq = size > 0 ? Math.floor((offset + size - 1) / chunkSize) : null;
+    const data =
+      size > 0
+        ? this.sql`substring(c.data FROM ${from} FOR least(${String(offset + size)}::bigint - greatest(${start}::bigint, ${chunkStart}), ${chunkSize}::bigint)::int)`
+        : this.sql`substring(c.data FROM ${from})`;
 
     // Stream the chunk rows back lazily via a server-side cursor so a large
     // snapshot is never buffered whole in memory (matching the file/S3
     // backends). Breaking out of the stream early closes the portal and
     // releases the connection.
-    let cursor: AsyncIterable<ChunkRow[]>;
+    let batches: AsyncIterator<ChunkRow[]>;
+    let first: IteratorResult<ChunkRow[]>;
     try {
-      cursor = this.sql`
-        SELECT seq, data FROM ${this.sql(this.chunksTable())}
-        WHERE prefix = ${this.prefix} AND level = ${level}
-          AND min_txid = ${txid(minTXID)} AND max_txid = ${txid(maxTXID)}
-          AND seq BETWEEN ${firstSeq} AND ${lastSeq}
-        ORDER BY seq ASC
+      const cursor = this.sql`
+        SELECT c.seq, ${data} AS data
+        FROM ${this.sql(this.filesTable())} f
+        LEFT JOIN ${this.sql(this.chunksTable())} c
+          ON c.prefix = f.prefix AND c.level = f.level
+          AND c.min_txid = f.min_txid AND c.max_txid = f.max_txid
+          AND c.seq >= ${firstSeq}
+          ${lastSeq === null ? this.sql`` : this.sql`AND c.seq <= ${lastSeq}`}
+        WHERE f.prefix = ${this.prefix} AND f.level = ${level}
+          AND f.min_txid = ${txid(minTXID)} AND f.max_txid = ${txid(maxTXID)}
+        ORDER BY c.seq ASC
       `.cursor(16) as unknown as AsyncIterable<ChunkRow[]>;
+      batches = cursor[Symbol.asyncIterator]();
+      // The first batch tells whether the file exists
+      first = await batches.next();
     } catch (err) {
       throw new LTXError('open', path, level, minTXID, maxTXID, err);
     }
+    if (first.done || first.value.length === 0) {
+      throw new LTXError('open', path, level, minTXID, maxTXID, new Error('not found'));
+    }
+    if (first.value[0].seq === null) return Readable.from([]);
 
-    return Readable.from(this.iterateChunkCursor(cursor, offset, end));
+    return Readable.from(this.iterateChunkCursor(first.value, batches));
   }
 
-  /** Streams the requested byte range from a chunk-row cursor. */
+  /** Streams the chunk parts of the first batch of rows, then of the next ones. */
   private async *iterateChunkCursor(
-    cursor: AsyncIterable<ChunkRow[]>,
-    offset: number,
-    end: number,
+    first: ChunkRow[],
+    batches: AsyncIterator<ChunkRow[]>,
   ): AsyncIterable<Uint8Array> {
-    for await (const rows of cursor) {
-      for (const row of rows) {
-        const data = row.data;
-        const chunkStart = Number(row.seq) * this.chunkSize;
-        const chunkEnd = chunkStart + data.length - 1;
-        if (chunkEnd < offset) continue;
-        if (chunkStart > end) return;
-        const s = Math.max(offset, chunkStart) - chunkStart;
-        const e = Math.min(end, chunkEnd) - chunkStart;
-        yield data.subarray(s, e + 1);
+    try {
+      for (let rows = first; ; ) {
+        for (const row of rows) {
+          if (row.data.length > 0) yield row.data;
+        }
+        const next = await batches.next();
+        if (next.done) return;
+        rows = next.value;
       }
+    } finally {
+      // Closes the portal when the stream is destroyed early
+      await batches.return?.();
     }
   }
 

packages/backend-pg/test/pg.test.ts

@@ -18,7 +18,8 @@ import {
   PgReplicaClient,
   type PgReplicaClientOptions,
 } from '../src/index.ts';
-import { connectForTest } from './helpers/db.ts';
+import postgres from 'postgres';
+import { connectForTest, pgUrl } from './helpers/db.ts';
 import { createWALDB } from './helpers/wal-db.ts';
 
 function concatBytes(chunks: readonly Uint8Array[]): Uint8Array {
@@ -211,6 +212,46 @@ test('openLtxFile reads a full file and byte ranges', async (t) => {
   assert.deepEqual([...tail], [...body.subarray(body.length - 10)]);
 });
 
+test('openLtxFile reads ranges across chunks, in one query each', async (t) => {
+  const queries: string[] = [];
+  const sql = postgres(pgUrl(), { onnotice: () => {}, max: 1, debug: (_connection, query) => queries.push(query) });
+  t.after(() => sql.end({ timeout: 1 }));
+  // 40 chunks of 64 bytes: more than a batch of the cursor (16 rows)
+  const client = new PgReplicaClient({ sql, createTables: true, prefix: makePrefix(), chunkSize: 64 });
+  const { body } = makeLtxBody(40 * 64 - 100);
+  await client.writeLtxFile({ level: 0, minTXID: 1n, maxTXID: 1n }, iter(body));
+
+  const read = async (offset: number, size: number): Promise<number[]> => {
+    queries.length = 0;
+    const bytes = await collect(await client.openLtxFile({ level: 0, minTXID: 1n, maxTXID: 1n, offset, size }));
+    // One query: the next batches of a long range continue it
+    assert.equal(queries.length, 1, `queries for ${offset}+${size}: ${queries.length}`);
+    return [...bytes];
+  };
+  const want = (offset: number, size: number): number[] => [...body.subarray(offset, size > 0 ? offset + size : undefined)];
+
+  for (const [offset, size] of [
+    [0, 0], // the whole file
+    [0, 64], // one chunk
+    [10, 20], // inside a chunk
+    [60, 10], // across two chunks
+    [64, 128], // from a chunk boundary
+    [100, 1000], // across many chunks
+    [130, 0], // to the end, from inside a chunk
+    [body.length - 5, 100], // past the end
+    [body.length, 10], // from the end
+    [body.length + 100, 0], // past the end, to the end
+  ]) {
+    assert.deepEqual(await read(offset, size), want(offset, size), `range ${offset}+${size}`);
+  }
+
+  // A stream left early releases the connection, the only one of the pool
+  const stream = await client.openLtxFile({ level: 0, minTXID: 1n, maxTXID: 1n, offset: 0, size: 0 });
+  await stream[Symbol.asyncIterator]().next();
+  stream.destroy();
+  assert.deepEqual(await read(10, 20), want(10, 20));
+});
+
 test('openLtxFile returns an empty stream when offset >= size', async (t) => {
   const client = makeClient(t, { prefix: makePrefix() });
   const { body } = makeLtxBody(5);