@aztec-labs/kv-store 6.0.0-nightly.20260829

package/src/lmdb-v2/read_transaction.ts

@@ -0,0 +1,156 @@
+import { CURSOR_PAGE_SIZE, Database, type LMDBMessageChannel, LMDBMessageType } from './message.js';
+
+export class ReadTransaction {
+  protected open = true;
+
+  constructor(protected channel: LMDBMessageChannel) {}
+
+  public close(): void {
+    if (!this.open) {
+      return;
+    }
+    this.open = false;
+  }
+
+  protected assertIsOpen() {
+    if (!this.open) {
+      throw new Error('Transaction is closed');
+    }
+  }
+
+  public async get(key: Uint8Array): Promise<Uint8Array | undefined> {
+    this.assertIsOpen();
+    const response = await this.channel.sendMessage(LMDBMessageType.GET, { keys: [key], db: Database.DATA });
+    return response.values[0]?.[0] ?? undefined;
+  }
+
+  public async getIndex(key: Uint8Array): Promise<Uint8Array[]> {
+    this.assertIsOpen();
+    const response = await this.channel.sendMessage(LMDBMessageType.GET, { keys: [key], db: Database.INDEX });
+    return response.values[0] ?? [];
+  }
+
+  public async *iterate(
+    startKey: Uint8Array,
+    endKey?: Uint8Array,
+    reverse = false,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array]> {
+    yield* this.#iterate(Database.DATA, startKey, endKey, reverse, limit, vals => vals[0]);
+  }
+
+  public async *iterateIndex(
+    startKey: Uint8Array,
+    endKey?: Uint8Array,
+    reverse = false,
+    limit?: number,
+  ): AsyncIterable<[Uint8Array, Uint8Array[]]> {
+    yield* this.#iterate(Database.INDEX, startKey, endKey, reverse, limit, vals => vals);
+  }
+
+  public countEntries(startKey: Uint8Array, endKey: Uint8Array, reverse: boolean): Promise<number> {
+    return this.#countEntries(Database.DATA, startKey, endKey, reverse);
+  }
+
+  public countEntriesIndex(startKey: Uint8Array, endKey: Uint8Array, reverse: boolean): Promise<number> {
+    return this.#countEntries(Database.INDEX, startKey, endKey, reverse);
+  }
+
+  async *#iterate<T>(
+    db: string,
+    startKey: Uint8Array,
+    endKey: Uint8Array | undefined,
+    reverse: boolean,
+    limit: number | undefined,
+    map: (val: Uint8Array[]) => T,
+  ): AsyncIterable<[Uint8Array, T]> {
+    this.assertIsOpen();
+
+    let cursor: number | undefined;
+    try {
+      const response = await this.channel.sendMessage(LMDBMessageType.START_CURSOR, {
+        key: startKey,
+        reverse,
+        count: typeof limit === 'number' ? Math.min(limit, CURSOR_PAGE_SIZE) : CURSOR_PAGE_SIZE,
+        onePage: typeof limit === 'number' && limit < CURSOR_PAGE_SIZE,
+        db,
+      });
+
+      cursor = response.cursor ?? undefined;
+      let entries = response.entries;
+      let done = typeof cursor !== 'number';
+      let count = 0;
+      // emit the first page and any subsequent pages in a while loop
+      // NB: end contition is in the middle of the while loop
+      while (entries.length > 0) {
+        for (const [key, values] of entries) {
+          if (typeof limit === 'number' && count >= limit) {
+            done = true;
+            break;
+          }
+
+          if (endKey) {
+            const cmp = Buffer.compare(key, endKey);
+            if ((!reverse && cmp >= 0) || (reverse && cmp <= 0)) {
+              done = true;
+              break;
+            }
+          }
+
+          count++;
+          yield [key, map(values)];
+        }
+
+        // cursor is null if DB returned everything in the first page
+        if (typeof cursor !== 'number' || done) {
+          break;
+        }
+
+        const response = await this.channel.sendMessage(LMDBMessageType.ADVANCE_CURSOR, {
+          cursor,
+          count: CURSOR_PAGE_SIZE,
+        });
+
+        done = response.done;
+        entries = response.entries;
+      }
+    } finally {
+      // we might not have anything to close
+      if (typeof cursor === 'number') {
+        await this.channel.sendMessage(LMDBMessageType.CLOSE_CURSOR, { cursor });
+      }
+    }
+  }
+
+  async #countEntries(db: string, startKey: Uint8Array, endKey: Uint8Array, reverse: boolean): Promise<number> {
+    this.assertIsOpen();
+
+    let cursor: number | undefined;
+    try {
+      const response = await this.channel.sendMessage(LMDBMessageType.START_CURSOR, {
+        key: startKey,
+        reverse,
+        count: 0,
+        onePage: false,
+        db,
+      });
+
+      cursor = response.cursor ?? undefined;
+      if (!cursor) {
+        return 0;
+      }
+
+      const advanceResponse = await this.channel.sendMessage(LMDBMessageType.ADVANCE_CURSOR_COUNT, {
+        cursor,
+        endKey: endKey,
+      });
+
+      return advanceResponse.count;
+    } finally {
+      // we might not have anything to close
+      if (typeof cursor === 'number') {
+        await this.channel.sendMessage(LMDBMessageType.CLOSE_CURSOR, { cursor });
+      }
+    }
+  }
+}

package/src/lmdb-v2/set.ts

@@ -0,0 +1,33 @@
+import type { Key, Range } from '../interfaces/common.js';
+import type { AztecAsyncSet } from '../interfaces/set.js';
+import { LMDBMap } from './map.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+
+/**
+ * A set backed by LMDB.
+ */
+export class LMDBSet<K extends Key> implements AztecAsyncSet<K> {
+  private map: LMDBMap<K, boolean>;
+
+  constructor(store: AztecLMDBStoreV2, mapName: string) {
+    this.map = new LMDBMap(store, mapName);
+  }
+
+  hasAsync(key: K): Promise<boolean> {
+    return this.map.hasAsync(key);
+  }
+
+  add(key: K): Promise<void> {
+    return this.map.set(key, true);
+  }
+
+  delete(key: K): Promise<void> {
+    return this.map.delete(key);
+  }
+
+  async *entriesAsync(range: Range<K> = {}): AsyncIterableIterator<K> {
+    for await (const key of this.map.keysAsync(range)) {
+      yield key;
+    }
+  }
+}

package/src/lmdb-v2/singleton.ts

@@ -0,0 +1,38 @@
+import { Encoder } from 'msgpackr';
+
+import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
+import { serializeKey } from './utils.js';
+
+export class LMDBSingleValue<T> implements AztecAsyncSingleton<T> {
+  private key: Uint8Array;
+  private encoder = new Encoder();
+  constructor(
+    private store: AztecLMDBStoreV2,
+    name: string,
+  ) {
+    this.key = serializeKey(`singleton:${name}`, 'value');
+  }
+
+  getAsync(): Promise<T | undefined> {
+    return execInReadTx(this.store, async tx => {
+      const val = await tx.get(this.key);
+      return val ? this.encoder.unpack(val) : undefined;
+    });
+  }
+
+  set(val: T): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      await tx.set(this.key, this.encoder.pack(val));
+      return true;
+    });
+  }
+
+  delete(): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      await tx.remove(this.key);
+      return true;
+    });
+  }
+}

package/src/lmdb-v2/store.ts

@@ -0,0 +1,220 @@
+import { type Logger, type LoggerBindings, createLogger } from '@aztec-labs/foundation/log';
+import { Semaphore, SerialQueue } from '@aztec-labs/foundation/queue';
+import { MsgpackChannel, NativeLMDBStore } from '@aztec-labs/native';
+import { AsyncLocalStorage } from 'async_hooks';
+import { mkdir, rm } from 'fs/promises';
+
+import type { AztecAsyncArray } from '../interfaces/array.js';
+import type { Key, StoreSize, Value } from '../interfaces/common.js';
+import type { AztecAsyncCounter } from '../interfaces/counter.js';
+import type { AztecAsyncMap } from '../interfaces/map.js';
+import type { AztecAsyncMultiMap } from '../interfaces/multi_map.js';
+import type { AztecAsyncSet } from '../interfaces/set.js';
+import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
+import type { AztecAsyncKVStore } from '../interfaces/store.js';
+import { LMDBArray } from './array.js';
+import { LMDBMap } from './map.js';
+import {
+  Database,
+  type LMDBMessageChannel,
+  LMDBMessageType,
+  type LMDBRequestBody,
+  type LMDBResponseBody,
+} from './message.js';
+import { LMDBMultiMap } from './multi_map.js';
+import { ReadTransaction } from './read_transaction.js';
+import { LMDBSet } from './set.js';
+import { LMDBSingleValue } from './singleton.js';
+import { WriteTransaction } from './write_transaction.js';
+
+export { execInReadTx, execInWriteTx } from './tx-helpers.js';
+
+export class AztecLMDBStoreV2 implements AztecAsyncKVStore, LMDBMessageChannel {
+  private open = false;
+  private channel: MsgpackChannel<LMDBMessageType, LMDBRequestBody, LMDBResponseBody>;
+  private writerCtx = new AsyncLocalStorage<WriteTransaction>();
+  private writerQueue = new SerialQueue();
+  private availableCursors: Semaphore;
+
+  private constructor(
+    private dataDir: string,
+    mapSize: number,
+    maxReaders: number,
+    private log: Logger,
+    private cleanup?: () => Promise<void>,
+    ephemeral: boolean = false,
+  ) {
+    this.log.info(`Starting data store with maxReaders ${maxReaders}`);
+    this.channel = new MsgpackChannel(new NativeLMDBStore(dataDir, mapSize, maxReaders, ephemeral));
+    // leave one reader to always be available for regular, atomic, reads
+    this.availableCursors = new Semaphore(maxReaders - 1);
+  }
+
+  public get dataDirectory(): string {
+    return this.dataDir;
+  }
+
+  private async start() {
+    this.writerQueue.start();
+
+    await this.channel.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.DATA,
+      uniqueKeys: true,
+    });
+
+    await this.channel.sendMessage(LMDBMessageType.OPEN_DATABASE, {
+      db: Database.INDEX,
+      uniqueKeys: false,
+    });
+
+    this.open = true;
+  }
+
+  public static async new(
+    dataDir: string,
+    dbMapSizeKb: number = 10 * 1024 * 1024,
+    maxReaders: number = 16,
+    cleanup?: () => Promise<void>,
+    bindings?: LoggerBindings,
+    ephemeral: boolean = false,
+  ) {
+    const log = createLogger('kv-store:lmdb-v2', bindings);
+    const db = new AztecLMDBStoreV2(dataDir, dbMapSizeKb, maxReaders, log, cleanup, ephemeral);
+    await db.start();
+    return db;
+  }
+
+  public async backupTo(dstPath: string, compact = true) {
+    await mkdir(dstPath, { recursive: true });
+    await this.channel.sendMessage(LMDBMessageType.COPY_STORE, { dstPath, compact });
+  }
+
+  public getReadTx(): ReadTransaction {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+    return new ReadTransaction(this);
+  }
+
+  public getCurrentWriteTx(): WriteTransaction | undefined {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+    const currentWrite = this.writerCtx.getStore();
+    return currentWrite;
+  }
+
+  openMap<K extends Key, V extends Value>(name: string): AztecAsyncMap<K, V> {
+    return new LMDBMap(this, name);
+  }
+
+  openMultiMap<K extends Key, V extends Value>(name: string): AztecAsyncMultiMap<K, V> {
+    return new LMDBMultiMap(this, name);
+  }
+
+  openSingleton<T extends Value>(name: string): AztecAsyncSingleton<T> {
+    return new LMDBSingleValue(this, name);
+  }
+
+  openArray<T extends Value>(name: string): AztecAsyncArray<T> {
+    return new LMDBArray(this, name);
+  }
+
+  openSet<K extends Key>(name: string): AztecAsyncSet<K> {
+    return new LMDBSet(this, name);
+  }
+
+  openCounter<K extends Key>(_name: string): AztecAsyncCounter<K> {
+    throw new Error('Not implemented');
+  }
+
+  async transactionAsync<T extends Exclude<any, Promise<any>>>(
+    callback: (tx: WriteTransaction) => Promise<T>,
+  ): Promise<T> {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+
+    // transactionAsync might be called recursively
+    // send any writes to the parent tx, but don't close it
+    // if the callback throws then the parent tx will rollback automatically
+    const currentTx = this.getCurrentWriteTx();
+    if (currentTx) {
+      return await callback(currentTx);
+    }
+
+    return this.writerQueue.put(async () => {
+      const tx = new WriteTransaction(this);
+      try {
+        const res = await this.writerCtx.run(tx, callback, tx);
+        await tx.commit();
+        return res;
+      } catch (err) {
+        this.log.error(`Failed to commit transaction`, err);
+        throw err;
+      } finally {
+        tx.close();
+      }
+    });
+  }
+
+  clear(): Promise<void> {
+    return Promise.resolve();
+  }
+
+  async delete(): Promise<void> {
+    await this.close();
+    await rm(this.dataDir, { recursive: true, force: true, maxRetries: 3 });
+    this.log.verbose(`Deleted database files at ${this.dataDir}`);
+    await this.cleanup?.();
+  }
+
+  async close() {
+    if (!this.open) {
+      // already closed
+      return;
+    }
+    this.open = false;
+    await this.writerQueue.cancel();
+    await this.channel.sendMessage(LMDBMessageType.CLOSE, undefined);
+  }
+
+  public async sendMessage<T extends LMDBMessageType>(
+    msgType: T,
+    body: LMDBRequestBody[T],
+  ): Promise<LMDBResponseBody[T]> {
+    if (!this.open) {
+      throw new Error('Store is closed');
+    }
+
+    if (msgType === LMDBMessageType.START_CURSOR) {
+      await this.availableCursors.acquire();
+    }
+
+    let response: LMDBResponseBody[T] | undefined = undefined;
+    try {
+      ({ response } = await this.channel.sendMessage(msgType, body));
+      return response;
+    } finally {
+      if (
+        (msgType === LMDBMessageType.START_CURSOR && response === undefined) ||
+        msgType === LMDBMessageType.CLOSE_CURSOR ||
+        // it's possible for a START_CURSOR command to not return a cursor (e.g. db is empty)
+        (msgType === LMDBMessageType.START_CURSOR &&
+          typeof (response as LMDBResponseBody[LMDBMessageType.START_CURSOR]).cursor !== 'number')
+      ) {
+        this.availableCursors.release();
+      }
+    }
+  }
+
+  public async estimateSize(): Promise<StoreSize> {
+    const resp = await this.sendMessage(LMDBMessageType.STATS, undefined);
+    return {
+      mappingSize: Number(resp.dbMapSizeBytes),
+      physicalFileSize: Number(resp.dbPhysicalFileSizeBytes),
+      actualSize: resp.stats.reduce((s, db) => Number(db.totalUsedSize) + s, 0),
+      numItems: resp.stats.reduce((s, db) => Number(db.numDataItems) + s, 0),
+    };
+  }
+}

package/src/lmdb-v2/tx-helpers.ts

@@ -0,0 +1,29 @@
+import type { ReadTransaction } from './read_transaction.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import type { WriteTransaction } from './write_transaction.js';
+
+export function execInWriteTx<T>(store: AztecLMDBStoreV2, fn: (tx: WriteTransaction) => Promise<T>): Promise<T> {
+  const currentWrite = store.getCurrentWriteTx();
+  if (currentWrite) {
+    return fn(currentWrite);
+  } else {
+    return store.transactionAsync(fn);
+  }
+}
+
+export async function execInReadTx<T>(
+  store: AztecLMDBStoreV2,
+  fn: (tx: ReadTransaction) => T | Promise<T>,
+): Promise<T> {
+  const currentWrite = store.getCurrentWriteTx();
+  if (currentWrite) {
+    return await fn(currentWrite);
+  } else {
+    const tx = store.getReadTx();
+    try {
+      return await fn(tx);
+    } finally {
+      tx.close();
+    }
+  }
+}

package/src/lmdb-v2/utils.ts

@@ -0,0 +1,32 @@
+import { MAXIMUM_KEY, fromBufferKey, toBufferKey } from 'ordered-binary';
+
+import type { Key } from '../interfaces/common.js';
+
+export function keyCmp(a: [Uint8Array, Uint8Array[] | null], b: [Uint8Array, Uint8Array[] | null]): -1 | 0 | 1 {
+  return Buffer.compare(a[0], b[0]);
+}
+
+export function singleKeyCmp(a: [Uint8Array, Uint8Array[] | null], b: Uint8Array): -1 | 0 | 1 {
+  return Buffer.compare(a[0], b);
+}
+
+export function minKey(prefix: string) {
+  return toBufferKey([prefix]);
+}
+
+export function maxKey(prefix: string) {
+  return toBufferKey([prefix, MAXIMUM_KEY]);
+}
+
+export function serializeKey(prefix: string, key: Key): Buffer {
+  return toBufferKey([prefix, key]);
+}
+
+export function deserializeKey<K extends Key>(prefix: string, key: Uint8Array): K | false {
+  const buf = Buffer.from(key);
+  const parsed = fromBufferKey(buf);
+  if (!Array.isArray(parsed) || parsed[0] !== prefix) {
+    return false;
+  }
+  return parsed[1] as K;
+}