From bfca557625d1ee64b3e43d0a89c66e6df85ff610 Mon Sep 17 00:00:00 2001 From: Santiago Palladino Date: Fri, 21 Aug 2026 10:35:37 +0100 Subject: [PATCH] fix(archiver): batch multi-key point reads through a single LMDB round trip The lmdb-v2 wire protocol always supported multi-key GET, but every getAsync sent a single-key message. Add getMany to the read and write transactions (overlaying pending batch state) and getManyAsync to AztecAsyncMap, and use it where the archiver did one point read per tx effect or per log: getBlock, getNoteHashesAndNullifiers, and per-block log reads. --- .../archiver/src/store/block_store.ts | 43 +++++++------ yarn-project/archiver/src/store/log_store.ts | 6 +- .../kv-store/src/deprecated/indexeddb/map.ts | 6 +- yarn-project/kv-store/src/interfaces/map.ts | 15 +++++ .../kv-store/src/interfaces/map_test_suite.ts | 31 +++++++++ yarn-project/kv-store/src/lmdb-v2/map.ts | 12 +++- yarn-project/kv-store/src/lmdb-v2/message.ts | 2 + .../kv-store/src/lmdb-v2/multi_map.ts | 5 ++ .../src/lmdb-v2/read_only_transaction.test.ts | 33 ++++++++++ .../src/lmdb-v2/read_transaction.test.ts | 61 ++++++++++++++++++ .../kv-store/src/lmdb-v2/read_transaction.ts | 38 ++++++++++- .../src/lmdb-v2/write_transaction.test.ts | 64 +++++++++++++++++++ .../kv-store/src/lmdb-v2/write_transaction.ts | 43 +++++++++++-- yarn-project/kv-store/src/lmdb/map.ts | 6 +- yarn-project/kv-store/src/sqlite-opfs/map.ts | 6 +- 15 files changed, 340 insertions(+), 31 deletions(-) diff --git a/yarn-project/archiver/src/store/block_store.ts b/yarn-project/archiver/src/store/block_store.ts index fa8de5e82ea..d2548bcd75a 100644 --- a/yarn-project/archiver/src/store/block_store.ts +++ b/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]; + }); } /** diff --git a/yarn-project/archiver/src/store/log_store.ts b/yarn-project/archiver/src/store/log_store.ts index 92083c183c7..490f84eb072 100644 --- a/yarn-project/archiver/src/store/log_store.ts +++ b/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({ diff --git a/yarn-project/kv-store/src/deprecated/indexeddb/map.ts b/yarn-project/kv-store/src/deprecated/indexeddb/map.ts index 2db0bd27301..ae9c3059ab3 100644 --- a/yarn-project/kv-store/src/deprecated/indexeddb/map.ts +++ b/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 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 { const result = (await this.getAsync(key)) !== undefined; return result; diff --git a/yarn-project/kv-store/src/interfaces/map.ts b/yarn-project/kv-store/src/interfaces/map.ts index 8e0ea473547..945cefb1943 100644 --- a/yarn-project/kv-store/src/interfaces/map.ts +++ b/yarn-project/kv-store/src/interfaces/map.ts @@ -77,6 +77,12 @@ export interface AztecMap extends AztecBaseMap extends AztecBaseMap { /** * Gets the value at the given key. @@ -84,6 +90,15 @@ export interface AztecAsyncMap extends AztecBase */ getAsync(key: K): Promise; + /** + * 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 diff --git a/yarn-project/kv-store/src/interfaces/map_test_suite.ts b/yarn-project/kv-store/src/interfaces/map_test_suite.ts index 94edc2bd4b3..16005dc5ec2 100644 --- a/yarn-project/kv-store/src/interfaces/map_test_suite.ts +++ b/yarn-project/kv-store/src/interfaces/map_test_suite.ts @@ -29,6 +29,10 @@ export function describeAztecMap( : await (sut as AztecAsyncMap).getAsync(key); } + async function getMany(keys: Key[], sut: AztecAsyncMap | AztecMap = map) { + return await (sut as AztecAsyncMap).getManyAsync(keys); + } + async function size(sut: AztecAsyncMap | AztecMap = map) { return isSyncStore(store) && !forceAsync ? (sut as AztecMap).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'); diff --git a/yarn-project/kv-store/src/lmdb-v2/map.ts b/yarn-project/kv-store/src/lmdb-v2/map.ts index 28913ea5540..ab9bc181a77 100644 --- a/yarn-project/kv-store/src/lmdb-v2/map.ts +++ b/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 implements AztecAsyncMap { + 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 { return execInReadTx(this.store, async tx => !!(await tx.get(serializeKey(this.prefix, key)))); } diff --git a/yarn-project/kv-store/src/lmdb-v2/message.ts b/yarn-project/kv-store/src/lmdb-v2/message.ts index 80923bb6a26..8c165936b1d 100644 --- a/yarn-project/kv-store/src/lmdb-v2/message.ts +++ b/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, diff --git a/yarn-project/kv-store/src/lmdb-v2/multi_map.ts b/yarn-project/kv-store/src/lmdb-v2/multi_map.ts index e37a7c36450..bfa6b08c520 100644 --- a/yarn-project/kv-store/src/lmdb-v2/multi_map.ts +++ b/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 implements AztecAsyncM }); } + getManyAsync(keys: K[], _opts?: GetManyOptions): Promise<(V | undefined)[]> { + return Promise.all(keys.map(key => this.getAsync(key))); + } + hasAsync(key: K): Promise { return execInReadTx(this.store, async tx => (await tx.getIndex(serializeKey(this.prefix, key))).length > 0); } diff --git a/yarn-project/kv-store/src/lmdb-v2/read_only_transaction.test.ts b/yarn-project/kv-store/src/lmdb-v2/read_only_transaction.test.ts index 8d70c12f9ae..6d719e87477 100644 --- a/yarn-project/kv-store/src/lmdb-v2/read_only_transaction.test.ts +++ b/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('batched'); + await store.transactionAsync(async () => { + await map.set('a', '1'); + await map.set('b', '2'); + }); + + const opened = promiseWithResolvers(); + const writeCommitted = promiseWithResolvers(); + + 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('iteration'); await store.transactionAsync(async () => { diff --git a/yarn-project/kv-store/src/lmdb-v2/read_transaction.test.ts b/yarn-project/kv-store/src/lmdb-v2/read_transaction.test.ts index 1ee85fd614b..05e65a13b7f 100644 --- a/yarn-project/kv-store/src/lmdb-v2/read_transaction.test.ts +++ b/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(); + + 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 () => { diff --git a/yarn-project/kv-store/src/lmdb-v2/read_transaction.ts b/yarn-project/kv-store/src/lmdb-v2/read_transaction.ts index d9b1c364fb8..ce3ef9d3fac 100644 --- a/yarn-project/kv-store/src/lmdb-v2/read_transaction.ts +++ b/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 { this.assertIsOpen(); const response = await this.channel.sendMessage(LMDBMessageType.GET, { diff --git a/yarn-project/kv-store/src/lmdb-v2/write_transaction.test.ts b/yarn-project/kv-store/src/lmdb-v2/write_transaction.test.ts index db0f45f3d9f..6a2a67cc0c3 100644 --- a/yarn-project/kv-store/src/lmdb-v2/write_transaction.test.ts +++ b/yarn-project/kv-store/src/lmdb-v2/write_transaction.test.ts @@ -100,6 +100,70 @@ describe('WriteTransaction', () => { expect(await tx.get(Buffer.from('foo'))).toEqual(undefined); }); + it('correctly overlays pending data on batched reads', async () => { + channel.sendMessage.mockImplementation((type: LMDBMessageType, body: any) => { + if (type === LMDBMessageType.GET) { + const committed: Record = { committed: Buffer.from('c'), overwritten: Buffer.from('old') }; + return Promise.resolve({ + values: body.keys.map((key: Buffer) => { + const value = committed[key.toString()]; + return value ? [value] : null; + }), + }); + } + return Promise.resolve({ ok: true }); + }); + + await tx.set(Buffer.from('overwritten'), Buffer.from('new')); + await tx.set(Buffer.from('added'), Buffer.from('a')); + await tx.remove(Buffer.from('committed')); + + const keys = ['overwritten', 'added', 'committed', 'missing'].map(k => Buffer.from(k)); + expect(await tx.getMany(keys)).toEqual([Buffer.from('new'), Buffer.from('a'), undefined, undefined]); + + // Only the keys the pending batch says nothing about make it to the store. + expect(channel.sendMessage).toHaveBeenCalledWith(LMDBMessageType.GET, { + keys: [Buffer.from('missing')], + db: Database.DATA, + txId: null, + }); + }); + + it('matches sequential gets when reading many with pending writes', async () => { + channel.sendMessage.mockImplementation((type: LMDBMessageType, body: any) => { + if (type === LMDBMessageType.GET) { + const committed: Record = { aaa: Buffer.from('1'), bbb: Buffer.from('2') }; + return Promise.resolve({ + values: body.keys.map((key: Buffer) => { + const value = committed[key.toString()]; + return value ? [value] : null; + }), + }); + } + return Promise.resolve({ ok: true }); + }); + + await tx.set(Buffer.from('ccc'), Buffer.from('3')); + await tx.remove(Buffer.from('aaa')); + + const keys = ['aaa', 'bbb', 'ccc', 'ddd'].map(k => Buffer.from(k)); + const sequential = []; + for (const key of keys) { + sequential.push(await tx.get(key)); + } + expect(await tx.getMany(keys)).toEqual(sequential); + }); + + it('skips the round trip when reading many with no keys', async () => { + expect(await tx.getMany([])).toEqual([]); + expect(channel.sendMessage).not.toHaveBeenCalledWith(LMDBMessageType.GET, expect.anything()); + }); + + it('refuses batched reads if closed', async () => { + tx.close(); + await expect(tx.getMany([Buffer.from('foo')])).rejects.toThrow('Transaction is closed'); + }); + it('correctly meanages pending index reads', async () => { channel.sendMessage.mockResolvedValue({ values: [[Buffer.from('1')]] }); expect(await tx.getIndex(Buffer.from('foo'))).toEqual([Buffer.from('1')]); diff --git a/yarn-project/kv-store/src/lmdb-v2/write_transaction.ts b/yarn-project/kv-store/src/lmdb-v2/write_transaction.ts index df85f3aab14..e5170019418 100644 --- a/yarn-project/kv-store/src/lmdb-v2/write_transaction.ts +++ b/yarn-project/kv-store/src/lmdb-v2/write_transaction.ts @@ -8,6 +8,7 @@ import { removeFromSortedArray, } from '@aztec/foundation/array'; +import type { GetManyOptions } from '../interfaces/map.js'; import { type Batch, Database, LMDBMessageType } from './message.js'; import { ReadTransaction } from './read_transaction.js'; import { keyCmp, singleKeyCmp } from './utils.js'; @@ -58,16 +59,50 @@ export class WriteTransaction extends ReadTransaction { public override async get(key: Buffer): Promise { this.assertIsOpen(); + const pending = this.#getPending(key); + if (pending) { + return pending.value; + } + + return await super.get(key); + } + + public override async getMany(keys: Uint8Array[], opts?: GetManyOptions): Promise<(Uint8Array | undefined)[]> { + this.assertIsOpen(); + + const results: (Uint8Array | undefined)[] = new Array(keys.length); + const unresolvedKeys: Uint8Array[] = []; + const unresolvedIndices: number[] = []; + + for (let i = 0; i < keys.length; i++) { + const pending = this.#getPending(keys[i]); + if (pending) { + results[i] = pending.value; + } else { + unresolvedKeys.push(keys[i]); + unresolvedIndices.push(i); + } + } + + const committed = await super.getMany(unresolvedKeys, opts); + for (let i = 0; i < unresolvedIndices.length; i++) { + results[unresolvedIndices[i]] = committed[i]; + } + + return results; + } + + /** Resolves a key against the pending data batch, returning undefined when the batch says nothing about it. */ + #getPending(key: Uint8Array): { value: Uint8Array | undefined } | undefined { const addEntry = findInSortedArray(this.dataBatch.addEntries, key, singleKeyCmp); if (addEntry) { - return addEntry[1][0]; + return { value: addEntry[1][0] }; } const removeEntryIdx = findIndexInSortedArray(this.dataBatch.removeEntries, key, singleKeyCmp); if (removeEntryIdx > -1) { - return undefined; + return { value: undefined }; } - - return await super.get(key); + return undefined; } setIndex(key: Buffer, ...values: Buffer[]): Promise { diff --git a/yarn-project/kv-store/src/lmdb/map.ts b/yarn-project/kv-store/src/lmdb/map.ts index 5dc97b5d189..1dbe3b771e4 100644 --- a/yarn-project/kv-store/src/lmdb/map.ts +++ b/yarn-project/kv-store/src/lmdb/map.ts @@ -1,7 +1,7 @@ import type { Database, RangeOptions } from 'lmdb'; import type { Key, Range, Value } from '../interfaces/common.js'; -import type { AztecAsyncMap, AztecMap } from '../interfaces/map.js'; +import type { AztecAsyncMap, AztecMap, GetManyOptions } from '../interfaces/map.js'; /** The slot where a key-value entry would be stored */ type MapValueSlot = ['map', string, 'slot', K]; @@ -39,6 +39,10 @@ export class LmdbAztecMap implements AztecMap { + return Promise.all(keys.map(key => this.getAsync(key))); + } + has(key: K): boolean { return this.db.doesExist(this.slot(key)); } diff --git a/yarn-project/kv-store/src/sqlite-opfs/map.ts b/yarn-project/kv-store/src/sqlite-opfs/map.ts index f350f7dfb28..154256e214c 100644 --- a/yarn-project/kv-store/src/sqlite-opfs/map.ts +++ b/yarn-project/kv-store/src/sqlite-opfs/map.ts @@ -4,7 +4,7 @@ import { Buffer } from 'buffer'; 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 { SqlValue } from './messages.js'; import type { AztecSQLiteOPFSStore } from './store.js'; @@ -31,6 +31,10 @@ export class SQLiteOPFSAztecMap implements Aztec return raw == null ? undefined : this.decodeValue(raw); } + getManyAsync(keys: K[], _opts?: GetManyOptions): Promise<(V | undefined)[]> { + return Promise.all(keys.map(key => this.getAsync(key))); + } + async hasAsync(key: K): Promise { const rows = await this.store.allAsync('SELECT 1 FROM data WHERE slot = ? LIMIT 1', [this.slot(key)]); return rows.length > 0;