fix(archiver): batch multi-key point reads into one LMDB round trip

yarn-project/archiver/src/store/block_store.ts

@@ -1136,13 +1136,18 @@ export class BlockStore {
       return undefined;
     }
 
-    const txEffects: TxEffect[] = [];
+    const txHashes: TxHash[] = [];
     const reader = BufferReader.asReader(blockTxsBuffer);
     while (!reader.isEmpty()) {
-      const txHash = reader.readObject(TxHash);
-      const txEffect = await this.#txEffects.getAsync(txHash.toString());
+      txHashes.push(reader.readObject(TxHash));
+    }
+
+    const storedTxEffects = await this.#txEffects.getManyAsync(txHashes.map(txHash => txHash.toString()));
+    const txEffects: TxEffect[] = [];
+    for (let i = 0; i < storedTxEffects.length; i++) {
+      const txEffect = storedTxEffects[i];
       if (txEffect === undefined) {
-        this.#log.warn(`Could not find tx effect for tx ${txHash} in block ${blockNumber}`);
+        this.#log.warn(`Could not find tx effect for tx ${txHashes[i]} in block ${blockNumber}`);
         return undefined;
       }
       txEffects.push(deserializeIndexedTxEffect(txEffect).data);
@@ -1204,22 +1209,20 @@ export class BlockStore {
    * skip the header, then read the two vectors, and stop — the large tail (`l2ToL1Msgs`,
    * `publicDataWrites`, `privateLogs`, `publicLogs`, `contractClassLogs`) is never touched.
    */
-  public getNoteHashesAndNullifiers(txHashes: TxHash[]): Promise<[Fr[], Fr[]][]> {
-    return Promise.all(
-      txHashes.map(async (txHash): Promise<[Fr[], Fr[]]> => {
-        const buffer = await this.#txEffects.getAsync(txHash.toString());
-        if (!buffer) {
-          return [[], []];
-        }
-        const reader = BufferReader.asReader(buffer);
-        // Skip the fixed-length header: blockHash + l2BlockNumber + txIndexInBlock + slotNumber + revertCode +
-        // txHash + transactionFee.
-        reader.readBytes(32 + 4 + 4 + 4 + 1 + 32 + 32);
-        const noteHashes = reader.readVectorUint8Prefix(Fr);
-        const nullifiers = reader.readVectorUint8Prefix(Fr);
-        return [noteHashes, nullifiers];
-      }),
-    );
+  public async getNoteHashesAndNullifiers(txHashes: TxHash[]): Promise<[Fr[], Fr[]][]> {
+    const buffers = await this.#txEffects.getManyAsync(txHashes.map(txHash => txHash.toString()));
+    return buffers.map((buffer): [Fr[], Fr[]] => {
+      if (!buffer) {
+        return [[], []];
+      }
+      const reader = BufferReader.asReader(buffer);
+      // Skip the fixed-length header: blockHash + l2BlockNumber + txIndexInBlock + slotNumber + revertCode +
+      // txHash + transactionFee.
+      reader.readBytes(32 + 4 + 4 + 4 + 1 + 32 + 32);
+      const noteHashes = reader.readVectorUint8Prefix(Fr);
+      const nullifiers = reader.readVectorUint8Prefix(Fr);
+      return [noteHashes, nullifiers];
+    });
   }
 
   /**

yarn-project/archiver/src/store/log_store.ts

@@ -344,12 +344,14 @@ export class LogStore {
     if (!keys || keys.length === 0) {
       return [];
     }
+    const raws = await primaryMap.getManyAsync(keys);
     const results: LogResult[] = [];
-    for (const key of keys) {
-      const raw = await primaryMap.getAsync(key);
+    for (let i = 0; i < keys.length; i++) {
+      const raw = raws[i];
       if (!raw) {
         continue;
       }
+      const key = keys[i];
       const tail = decodeKeyTail(key);
       const value = decodeValue(raw);
       results.push({

yarn-project/kv-store/src/deprecated/indexeddb/map.ts

@@ -2,7 +2,7 @@ import type { IDBPDatabase, IDBPObjectStore } from 'idb';
 import { hash } from 'ohash';
 
 import type { Key, Range, Value } from '../../interfaces/common.js';
-import type { AztecAsyncMap } from '../../interfaces/map.js';
+import type { AztecAsyncMap, GetManyOptions } from '../../interfaces/map.js';
 import type { AztecIDBSchema } from './store.js';
 
 /**
@@ -34,6 +34,10 @@ export class IndexedDBAztecMap<K extends Key, V extends Value> implements AztecA
     return data ? this.restoreBuffers(data.value as V) : undefined;
   }
 
+  getManyAsync(keys: K[], _opts?: GetManyOptions): Promise<(V | undefined)[]> {
+    return Promise.all(keys.map(key => this.getAsync(key)));
+  }
+
   async hasAsync(key: K): Promise<boolean> {
     const result = (await this.getAsync(key)) !== undefined;
     return result;

yarn-project/kv-store/src/interfaces/map.ts

@@ -77,13 +77,28 @@ export interface AztecMap<K extends Key, V extends Value> extends AztecBaseMap<K
 /**
  * A map backed by a persistent store.
  */
+/** Options for {@link AztecAsyncMap.getManyAsync}. */
+export type GetManyOptions = {
+  /** Maximum keys per backend request. */
+  chunkSize?: number;
+};
+
 export interface AztecAsyncMap<K extends Key, V extends Value> extends AztecBaseMap<K, V> {
   /**
    * Gets the value at the given key.
    * @param key - The key to get the value from
    */
   getAsync(key: K): Promise<V | undefined>;
 
+  /**
+   * Gets the values at the given keys, batching them into as few round trips as the backend supports.
+   * @param keys - The keys to get the values from
+   * @param opts - `chunkSize` caps the keys sent per backend request (default 1024); ignored by backends that read
+   * keys individually
+   * @returns One entry per input key, in input order, `undefined` where the key is absent
+   */
+  getManyAsync(keys: K[], opts?: GetManyOptions): Promise<(V | undefined)[]>;
+
   /**
    * Checks if a key exists in the map.
    * @param key - The key to check

yarn-project/kv-store/src/interfaces/map_test_suite.ts

@@ -29,6 +29,10 @@ export function describeAztecMap(
         : await (sut as AztecAsyncMap<any, any>).getAsync(key);
     }
 
+    async function getMany(keys: Key[], sut: AztecAsyncMap<any, any> | AztecMap<any, any> = map) {
+      return await (sut as AztecAsyncMap<any, any>).getManyAsync(keys);
+    }
+
     async function size(sut: AztecAsyncMap<any, any> | AztecMap<any, any> = map) {
       return isSyncStore(store) && !forceAsync
         ? (sut as AztecMap<any, any>).size()
@@ -71,6 +75,33 @@ export function describeAztecMap(
       }
     });
 
+    it('should be able to get many values at once', async () => {
+      await map.set('foo', 'bar');
+      await map.set('baz', 'qux');
+
+      expect(await getMany(['foo', 'baz'])).toEqual(['bar', 'qux']);
+    });
+
+    it('returns undefined for missing keys when getting many', async () => {
+      await map.set('foo', 'bar');
+
+      expect(await getMany(['missing', 'foo', 'alsoMissing'])).toEqual([undefined, 'bar', undefined]);
+    });
+
+    it('preserves input order when getting many', async () => {
+      const pairs = Array.from({ length: 20 }, (_, i) => ({ key: `key${i}`, value: `value${i}` }));
+      await map.setMany(pairs);
+
+      const requested = [...pairs].reverse().map(({ key }) => key);
+      expect(await getMany(requested)).toEqual([...pairs].reverse().map(({ value }) => value));
+    });
+
+    it('returns an empty array when getting many with no keys', async () => {
+      await map.set('foo', 'bar');
+
+      expect(await getMany([])).toEqual([]);
+    });
+
     it('should be able to overwrite values', async () => {
       await map.set('foo', 'bar');
       await map.set('foo', 'baz');

yarn-project/kv-store/src/lmdb-v2/map.ts

@@ -1,7 +1,7 @@
 import { Encoder } from 'msgpackr';
 
 import type { Key, Range, Value } from '../interfaces/common.js';
-import type { AztecAsyncMap } from '../interfaces/map.js';
+import type { AztecAsyncMap, GetManyOptions } from '../interfaces/map.js';
 import type { AztecLMDBStoreV2 } from './store.js';
 import { acquireReadTx, execInReadTx, execInWriteTx } from './tx-helpers.js';
 import { deserializeKey, maxKey, minKey, serializeKey } from './utils.js';
@@ -69,6 +69,16 @@ export class LMDBMap<K extends Key, V extends Value> implements AztecAsyncMap<K,
     });
   }
 
+  getManyAsync(keys: K[], opts?: GetManyOptions): Promise<(V | undefined)[]> {
+    return execInReadTx(this.store, async tx => {
+      const vals = await tx.getMany(
+        keys.map(key => serializeKey(this.prefix, key)),
+        opts,
+      );
+      return vals.map(val => (val ? this.encoder.unpack(val) : undefined));
+    });
+  }
+
   hasAsync(key: K): Promise<boolean> {
     return execInReadTx(this.store, async tx => !!(await tx.get(serializeKey(this.prefix, key))));
   }

yarn-project/kv-store/src/lmdb-v2/message.ts

@@ -4,6 +4,8 @@ export enum Database {
 }
 
 export const CURSOR_PAGE_SIZE = 10;
+/** Keys per GET message when a batched read is split into chunks; bounds message size and how long one read holds its tx. */
+export const DEFAULT_GET_CHUNK_SIZE = 1024;
 
 export enum LMDBMessageType {
   OPEN_DATABASE = 100,

yarn-project/kv-store/src/lmdb-v2/multi_map.ts

@@ -2,6 +2,7 @@ import { Encoder } from 'msgpackr/pack';
 import { MAXIMUM_KEY, toBufferKey } from 'ordered-binary';
 
 import type { Key, Range, Value } from '../interfaces/common.js';
+import type { GetManyOptions } from '../interfaces/map.js';
 import type { AztecAsyncMultiMap } from '../interfaces/multi_map.js';
 import type { AztecLMDBStoreV2 } from './store.js';
 import { acquireReadTx, execInReadTx, execInWriteTx } from './tx-helpers.js';
@@ -65,6 +66,10 @@ export class LMDBMultiMap<K extends Key, V extends Value> implements AztecAsyncM
     });
   }
 
+  getManyAsync(keys: K[], _opts?: GetManyOptions): Promise<(V | undefined)[]> {
+    return Promise.all(keys.map(key => this.getAsync(key)));
+  }
+
   hasAsync(key: K): Promise<boolean> {
     return execInReadTx(this.store, async tx => (await tx.getIndex(serializeKey(this.prefix, key))).length > 0);
   }

yarn-project/kv-store/src/lmdb-v2/read_only_transaction.test.ts

@@ -69,6 +69,39 @@ describe('AztecLMDBStoreV2 readOnlyTransaction', () => {
     await expect(map.getAsync('k')).resolves.toBe('v2');
   });
 
+  it('serves batched container reads from the snapshot', async () => {
+    const map = store.openMap<string, string>('batched');
+    await store.transactionAsync(async () => {
+      await map.set('a', '1');
+      await map.set('b', '2');
+    });
+
+    const opened = promiseWithResolvers<void>();
+    const writeCommitted = promiseWithResolvers<void>();
+
+    const snapshotReads = store.readOnlyTransaction(async () => {
+      const first = await map.getManyAsync(['a', 'b', 'c']);
+      opened.resolve();
+      await writeCommitted.promise;
+      return [first, await map.getManyAsync(['a', 'b', 'c'])];
+    });
+
+    await opened.promise;
+    await store.transactionAsync(async () => {
+      await map.set('b', '22');
+      await map.set('c', '3');
+    });
+    writeCommitted.resolve();
+
+    await expect(snapshotReads).resolves.toEqual([
+      ['1', '2', undefined],
+      ['1', '2', undefined],
+    ]);
+
+    // outside the snapshot the batched read sees the committed values
+    await expect(map.getManyAsync(['a', 'b', 'c'])).resolves.toEqual(['1', '22', '3']);
+  });
+
   it('does not observe rows committed after the snapshot was taken while iterating', async () => {
     const map = store.openMap<string, string>('iteration');
     await store.transactionAsync(async () => {

yarn-project/kv-store/src/lmdb-v2/read_transaction.test.ts

@@ -51,6 +51,67 @@ describe('ReadTransaction', () => {
       keys: [Buffer.from('test_key1')],
       txId: 7,
     });
+
+    await expect(boundTx.getMany([Buffer.from('test_key1')])).resolves.toEqual([Buffer.from('foo')]);
+    expect(channel.sendMessage).toHaveBeenLastCalledWith(LMDBMessageType.GET, {
+      db: Database.DATA,
+      keys: [Buffer.from('test_key1')],
+      txId: 7,
+    });
+  });
+
+  it('splits many keys into chunked GET requests and preserves order', async () => {
+    const keys = Array.from({ length: 5 }, (_, i) => Buffer.from(`key${i}`));
+    channel.sendMessage.mockImplementation((_type, body: any) =>
+      Promise.resolve({ values: body.keys.map((k: Buffer) => (k.equals(keys[3]) ? null : [Buffer.from(`v-${k}`)])) }),
+    );
+
+    const result = await tx.getMany(keys, { chunkSize: 2 });
+
+    expect(channel.sendMessage).toHaveBeenCalledTimes(3);
+    expect(channel.sendMessage).toHaveBeenNthCalledWith(1, LMDBMessageType.GET, {
+      db: Database.DATA,
+      keys: keys.slice(0, 2),
+      txId: null,
+    });
+    expect(channel.sendMessage).toHaveBeenNthCalledWith(3, LMDBMessageType.GET, {
+      db: Database.DATA,
+      keys: keys.slice(4),
+      txId: null,
+    });
+    expect(result).toEqual(keys.map((k, i) => (i === 3 ? undefined : Buffer.from(`v-${k}`))));
+  });
+
+  it('rejects an invalid chunk size', async () => {
+    await expect(tx.getMany([Buffer.from('a')], { chunkSize: 0 })).rejects.toThrow('Invalid getMany chunk size');
+  });
+
+  it('sends a single GET request for many keys', async () => {
+    const getDeferred = promiseWithResolvers<LMDBResponseBody[LMDBMessageType.GET]>();
+
+    channel.sendMessage.mockReturnValue(getDeferred.promise);
+
+    const keys = [Buffer.from('key1'), Buffer.from('key2'), Buffer.from('key3')];
+    const resp = tx.getMany(keys);
+
+    expect(channel.sendMessage).toHaveBeenCalledTimes(1);
+    expect(channel.sendMessage).toHaveBeenCalledWith(LMDBMessageType.GET, { db: Database.DATA, keys, txId: null });
+
+    getDeferred.resolve({
+      values: [[Buffer.from('foo')], null, [Buffer.from('bar')]],
+    });
+
+    expect(await resp).toEqual([Buffer.from('foo'), undefined, Buffer.from('bar')]);
+  });
+
+  it('skips the GET request when asked for no keys', async () => {
+    expect(await tx.getMany([])).toEqual([]);
+    expect(channel.sendMessage).not.toHaveBeenCalled();
+  });
+
+  it('refuses batched reads once closed', async () => {
+    tx.close();
+    await expect(tx.getMany([Buffer.from('foo')])).rejects.toThrow('Transaction is closed');
   });
 
   it('iterates the database', async () => {

yarn-project/kv-store/src/lmdb-v2/read_transaction.ts

@@ -1,4 +1,13 @@
-import { CURSOR_PAGE_SIZE, Database, type LMDBMessageChannel, LMDBMessageType } from './message.js';
+import { chunk } from '@aztec/foundation/collection';
+
+import type { GetManyOptions } from '../interfaces/map.js';
+import {
+  CURSOR_PAGE_SIZE,
+  DEFAULT_GET_CHUNK_SIZE,
+  Database,
+  type LMDBMessageChannel,
+  LMDBMessageType,
+} from './message.js';
 
 /**
  * Reads against the store. When constructed with the id of a native read transaction, every read and iteration is
@@ -36,6 +45,33 @@ export class ReadTransaction {
     return response.values[0]?.[0] ?? undefined;
   }
 
+  /**
+   * Reads many keys, sending at most `opts.chunkSize` keys per GET message (default {@link DEFAULT_GET_CHUNK_SIZE}).
+   * Returns one entry per input key, in input order, `undefined` where the key is absent.
+   */
+  public async getMany(keys: Uint8Array[], opts: GetManyOptions = {}): Promise<(Uint8Array | undefined)[]> {
+    this.assertIsOpen();
+    if (keys.length === 0) {
+      return [];
+    }
+    const chunkSize = opts.chunkSize ?? DEFAULT_GET_CHUNK_SIZE;
+    if (!Number.isInteger(chunkSize) || chunkSize < 1) {
+      throw new Error(`Invalid getMany chunk size: ${chunkSize}`);
+    }
+    const results: (Uint8Array | undefined)[] = [];
+    for (const keysChunk of chunk(keys, chunkSize)) {
+      const response = await this.channel.sendMessage(LMDBMessageType.GET, {
+        keys: keysChunk,
+        db: Database.DATA,
+        txId: this.txId ?? null,
+      });
+      for (let i = 0; i < keysChunk.length; i++) {
+        results.push(response.values[i]?.[0] ?? undefined);
+      }
+    }
+    return results;
+  }
+
   public async getIndex(key: Uint8Array): Promise<Uint8Array[]> {
     this.assertIsOpen();
     const response = await this.channel.sendMessage(LMDBMessageType.GET, {