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

package/src/lmdb/map.ts

@@ -0,0 +1,177 @@
+import type { Database, RangeOptions } from 'lmdb';
+
+import type { Key, Range, Value } from '../interfaces/common.js';
+import type { AztecAsyncMap, AztecMap } from '../interfaces/map.js';
+
+/** The slot where a key-value entry would be stored */
+type MapValueSlot<K extends Key | Buffer> = ['map', string, 'slot', K];
+
+/**
+ * A map backed by LMDB.
+ */
+export class LmdbAztecMap<K extends Key, V extends Value> implements AztecMap<K, V>, AztecAsyncMap<K, V> {
+  protected db: Database<[K, V], MapValueSlot<K>>;
+  protected name: string;
+
+  protected startSentinel: MapValueSlot<Buffer>;
+  protected endSentinel: MapValueSlot<Buffer>;
+
+  constructor(rootDb: Database, mapName: string) {
+    this.name = mapName;
+    this.db = rootDb as Database<[K, V], MapValueSlot<K>>;
+
+    // sentinels are used to define the start and end of the map
+    // with LMDB's key encoding, no _primitive value_ can be "less than" an empty buffer or greater than Byte 255
+    // these will be used later to answer range queries
+    this.startSentinel = ['map', this.name, 'slot', Buffer.from([])];
+    this.endSentinel = ['map', this.name, 'slot', Buffer.from([255])];
+  }
+
+  close(): Promise<void> {
+    return this.db.close();
+  }
+
+  get(key: K): V | undefined {
+    return this.db.get(this.slot(key))?.[1];
+  }
+
+  getAsync(key: K): Promise<V | undefined> {
+    return Promise.resolve(this.get(key));
+  }
+
+  has(key: K): boolean {
+    return this.db.doesExist(this.slot(key));
+  }
+
+  hasAsync(key: K): Promise<boolean> {
+    return Promise.resolve(this.has(key));
+  }
+
+  async set(key: K, val: V): Promise<void> {
+    await this.db.put(this.slot(key), [key, val]);
+  }
+
+  async setMany(entries: { key: K; value: V }[]): Promise<void> {
+    for (const { key, value } of entries) {
+      await this.set(key, value);
+    }
+  }
+
+  swap(key: K, fn: (val: V | undefined) => V): Promise<void> {
+    return this.db.childTransaction(() => {
+      const slot = this.slot(key);
+      const entry = this.db.get(slot);
+      void this.db.put(slot, [key, fn(entry?.[1])]);
+    });
+  }
+
+  setIfNotExists(key: K, val: V): Promise<boolean> {
+    const slot = this.slot(key);
+    return this.db.ifNoExists(slot, () => {
+      void this.db.put(slot, [key, val]);
+    });
+  }
+
+  async delete(key: K): Promise<void> {
+    await this.db.remove(this.slot(key));
+  }
+
+  *entries(range: Range<K> = {}): IterableIterator<[K, V]> {
+    const transaction = this.db.useReadTransaction();
+
+    try {
+      const { reverse = false, limit } = range;
+      // LMDB has a quirk where it expects start > end when reverse=true
+      // in that case, we need to swap the start and end sentinels
+      const start = reverse
+        ? range.end
+          ? this.slot(range.end)
+          : this.endSentinel
+        : range.start
+          ? this.slot(range.start)
+          : this.startSentinel;
+
+      const end = reverse
+        ? range.start
+          ? this.slot(range.start)
+          : this.startSentinel
+        : range.end
+          ? this.slot(range.end)
+          : this.endSentinel;
+
+      const lmdbRange: RangeOptions = {
+        start,
+        end,
+        reverse,
+        limit,
+        transaction,
+      };
+
+      const iterator = this.db.getRange(lmdbRange);
+
+      for (const {
+        value: [key, value],
+      } of iterator) {
+        yield [key, value];
+      }
+    } finally {
+      transaction.done();
+    }
+  }
+
+  async *entriesAsync(range?: Range<K>): AsyncIterableIterator<[K, V]> {
+    for (const entry of this.entries(range)) {
+      yield entry;
+    }
+  }
+
+  *values(range: Range<K> = {}): IterableIterator<V> {
+    for (const [_, value] of this.entries(range)) {
+      yield value;
+    }
+  }
+
+  async *valuesAsync(range: Range<K> = {}): AsyncIterableIterator<V> {
+    for await (const [_, value] of this.entriesAsync(range)) {
+      yield value;
+    }
+  }
+
+  size(): number {
+    const iterator = this.db.getRange({ start: this.startSentinel, end: this.endSentinel });
+    return iterator.asArray.length;
+  }
+
+  sizeAsync(): Promise<number> {
+    return Promise.resolve(this.size());
+  }
+
+  *keys(range: Range<K> = {}): IterableIterator<K> {
+    for (const [key, _] of this.entries(range)) {
+      yield key;
+    }
+  }
+
+  async *keysAsync(range: Range<K> = {}): AsyncIterableIterator<K> {
+    for await (const [key, _] of this.entriesAsync(range)) {
+      yield key;
+    }
+  }
+
+  protected slot(key: K): MapValueSlot<K> {
+    return ['map', this.name, 'slot', key];
+  }
+
+  async clear(): Promise<void> {
+    const lmdbRange: RangeOptions = {
+      start: this.startSentinel,
+      end: this.endSentinel,
+    };
+
+    const iterator = this.db.getRange(lmdbRange);
+
+    for (const { key } of iterator) {
+      await this.db.remove(key);
+    }
+  }
+}

package/src/lmdb/multi_map.ts

@@ -0,0 +1,51 @@
+import type { Key, Value } from '../interfaces/common.js';
+import type { AztecAsyncMultiMap, AztecMultiMap } from '../interfaces/multi_map.js';
+import { LmdbAztecMap } from './map.js';
+
+/**
+ * A map backed by LMDB.
+ */
+export class LmdbAztecMultiMap<K extends Key, V extends Value>
+  extends LmdbAztecMap<K, V>
+  implements AztecMultiMap<K, V>, AztecAsyncMultiMap<K, V>
+{
+  *getValues(key: K): IterableIterator<V> {
+    const transaction = this.db.useReadTransaction();
+    try {
+      const values = this.db.getValues(this.slot(key), {
+        transaction,
+      });
+      for (const value of values) {
+        yield value?.[1];
+      }
+    } finally {
+      transaction.done();
+    }
+  }
+
+  async *getValuesAsync(key: K): AsyncIterableIterator<V> {
+    for (const value of this.getValues(key)) {
+      yield value;
+    }
+  }
+
+  getValueCountAsync(key: K): Promise<number> {
+    const transaction = this.db.useReadTransaction();
+    try {
+      const values = this.db.getValues(this.slot(key), {
+        transaction,
+      });
+      let count = 0;
+      for (const _ of values) {
+        count++;
+      }
+      return Promise.resolve(count);
+    } finally {
+      transaction.done();
+    }
+  }
+
+  async deleteValue(key: K, val: V): Promise<void> {
+    await this.db.remove(this.slot(key), [key, val]);
+  }
+}

package/src/lmdb/set.ts

@@ -0,0 +1,45 @@
+import type { Database } from 'lmdb';
+
+import type { Key, Range } from '../interfaces/common.js';
+import type { AztecAsyncSet, AztecSet } from '../interfaces/set.js';
+import { LmdbAztecMap } from './map.js';
+
+/**
+ * A set backed by LMDB.
+ */
+export class LmdbAztecSet<K extends Key> implements AztecSet<K>, AztecAsyncSet<K> {
+  private map: LmdbAztecMap<K, boolean>;
+  constructor(rootDb: Database, mapName: string) {
+    this.map = new LmdbAztecMap(rootDb, mapName);
+  }
+
+  close(): Promise<void> {
+    return this.map.close();
+  }
+
+  has(key: K): boolean {
+    return this.map.has(key);
+  }
+
+  hasAsync(key: K): Promise<boolean> {
+    return Promise.resolve(this.has(key));
+  }
+
+  add(key: K): Promise<void> {
+    return this.map.set(key, true);
+  }
+
+  delete(key: K): Promise<void> {
+    return this.map.delete(key);
+  }
+
+  entries(range: Range<K> = {}): IterableIterator<K> {
+    return this.map.keys(range);
+  }
+
+  async *entriesAsync(range: Range<K> = {}): AsyncIterableIterator<K> {
+    for await (const key of this.map.keysAsync(range)) {
+      yield key;
+    }
+  }
+}

package/src/lmdb/singleton.ts

@@ -0,0 +1,35 @@
+import type { Database, Key } from 'lmdb';
+
+import type { AztecAsyncSingleton, AztecSingleton } from '../interfaces/singleton.js';
+
+/** The slot where this singleton will store its value */
+type ValueSlot = ['singleton', string, 'value'];
+
+/**
+ * Stores a single value in LMDB.
+ */
+export class LmdbAztecSingleton<T> implements AztecSingleton<T>, AztecAsyncSingleton<T> {
+  #db: Database<T, ValueSlot>;
+  #slot: ValueSlot;
+
+  constructor(db: Database<unknown, Key>, name: string) {
+    this.#db = db as Database<T, ValueSlot>;
+    this.#slot = ['singleton', name, 'value'];
+  }
+
+  get(): T | undefined {
+    return this.#db.get(this.#slot);
+  }
+
+  getAsync(): Promise<T | undefined> {
+    return Promise.resolve(this.get());
+  }
+
+  set(val: T): Promise<boolean> {
+    return this.#db.put(this.#slot, val);
+  }
+
+  delete(): Promise<boolean> {
+    return this.#db.remove(this.#slot);
+  }
+}

package/src/lmdb/store.ts

@@ -0,0 +1,242 @@
+import { randomBytes } from '@aztec-labs/foundation/crypto/random';
+import { createLogger } from '@aztec-labs/foundation/log';
+import { promises as fs, mkdirSync } from 'fs';
+import { type Database, type RootDatabase, open } from 'lmdb';
+import { tmpdir } from 'os';
+import { join } from 'path';
+
+import type { AztecArray, AztecAsyncArray } from '../interfaces/array.js';
+import type { Key, StoreSize, Value } from '../interfaces/common.js';
+import type { AztecAsyncCounter, AztecCounter } from '../interfaces/counter.js';
+import type { AztecAsyncMap, AztecMap } from '../interfaces/map.js';
+import type { AztecAsyncMultiMap, AztecMultiMap } from '../interfaces/multi_map.js';
+import type { AztecAsyncSet, AztecSet } from '../interfaces/set.js';
+import type { AztecAsyncSingleton, AztecSingleton } from '../interfaces/singleton.js';
+import type { AztecAsyncKVStore, AztecKVStore } from '../interfaces/store.js';
+import { LmdbAztecArray } from './array.js';
+import { LmdbAztecCounter } from './counter.js';
+import { LmdbAztecMap } from './map.js';
+import { LmdbAztecMultiMap } from './multi_map.js';
+import { LmdbAztecSet } from './set.js';
+import { LmdbAztecSingleton } from './singleton.js';
+
+/**
+ * A key-value store backed by LMDB.
+ */
+export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
+  syncGetters = true as const;
+
+  #rootDb: RootDatabase;
+  #data: Database<unknown, Key>;
+  #multiMapData: Database<unknown, Key>;
+  #log = createLogger('kv-store:lmdb');
+
+  constructor(
+    rootDb: RootDatabase,
+    public readonly isEphemeral: boolean,
+    private path: string,
+  ) {
+    this.#rootDb = rootDb;
+
+    // big bucket to store all the data
+    this.#data = rootDb.openDB('data', {
+      encoding: 'msgpack',
+      keyEncoding: 'ordered-binary',
+    });
+
+    this.#multiMapData = rootDb.openDB('data_dup_sort', {
+      encoding: 'ordered-binary',
+      keyEncoding: 'ordered-binary',
+      dupSort: true,
+    });
+  }
+
+  /**
+   * Creates a new AztecKVStore backed by LMDB. The path to the database is optional. If not provided,
+   * the database will be stored in a temporary location and be deleted when the process exists.
+   *
+   * The `rollupAddress` passed is checked against what is stored in the database. If they do not match,
+   * the database is cleared before returning the store. This way data is not accidentally shared between
+   * different rollup instances.
+   *
+   * @param path - A path on the disk to store the database. Optional
+   * @param ephemeral - true if the store should only exist in memory and not automatically be flushed to disk. Optional
+   * @param log - A logger to use. Optional
+   * @returns The store
+   */
+  static open(
+    path?: string,
+    mapSizeKb = 1 * 1024 * 1024, // defaults to 1 GB map size
+    ephemeral: boolean = false,
+    log = createLogger('kv-store:lmdb'),
+  ): AztecLmdbStore {
+    const dbPath = path ?? join(tmpdir(), randomBytes(8).toString('hex'));
+    mkdirSync(dbPath, { recursive: true });
+    const mapSize = 1024 * mapSizeKb;
+    log.debug(`Opening LMDB database at ${path || 'temporary location'} with map size ${mapSize}`);
+    const rootDb = open({ path: dbPath, noSync: ephemeral, mapSize });
+    return new AztecLmdbStore(rootDb, ephemeral, dbPath);
+  }
+
+  /**
+   * Creates a new AztecMap in the store.
+   * @param name - Name of the map
+   * @returns A new AztecMap
+   */
+  openMap<K extends Key, V extends Value>(name: string): AztecMap<K, V> & AztecAsyncMap<K, V> {
+    return new LmdbAztecMap(this.#data, name);
+  }
+
+  /**
+   * Creates a new AztecSet in the store.
+   * @param name - Name of the set
+   * @returns A new AztecSet
+   */
+  openSet<K extends Key>(name: string): AztecSet<K> & AztecAsyncSet<K> {
+    return new LmdbAztecSet(this.#data, name);
+  }
+
+  /**
+   * Creates a new AztecMultiMap in the store. A multi-map stores multiple values for a single key automatically.
+   * @param name - Name of the map
+   * @returns A new AztecMultiMap
+   */
+  openMultiMap<K extends Key, V extends Value>(name: string): AztecMultiMap<K, V> & AztecAsyncMultiMap<K, V> {
+    return new LmdbAztecMultiMap(this.#multiMapData, name);
+  }
+
+  openCounter<K extends Key>(name: string): AztecCounter<K> & AztecAsyncCounter<K> {
+    return new LmdbAztecCounter(this.#data, name);
+  }
+
+  /**
+   * Creates a new AztecArray in the store.
+   * @param name - Name of the array
+   * @returns A new AztecArray
+   */
+  openArray<T extends Value>(name: string): AztecArray<T> & AztecAsyncArray<T> {
+    return new LmdbAztecArray(this.#data, name);
+  }
+
+  /**
+   * Creates a new AztecSingleton in the store.
+   * @param name - Name of the singleton
+   * @returns A new AztecSingleton
+   */
+  openSingleton<T>(name: string): AztecSingleton<T> & AztecAsyncSingleton<T> {
+    return new LmdbAztecSingleton(this.#data, name);
+  }
+
+  /**
+   * Runs a callback in a transaction.
+   * @param callback - Function to execute in a transaction
+   * @returns A promise that resolves to the return value of the callback
+   */
+  transaction<T>(callback: () => T): Promise<T> {
+    return this.#rootDb.transaction(callback);
+  }
+
+  /**
+   * Runs a callback in a transaction.
+   * @param callback - Function to execute in a transaction
+   * @returns A promise that resolves to the return value of the callback
+   */
+  async transactionAsync<T>(callback: () => Promise<T>): Promise<T> {
+    return await this.#rootDb.transaction(callback);
+  }
+
+  /**
+   * Clears all entries in the store & sub DBs atomically within a single transaction.
+   */
+  async clear() {
+    await this.#rootDb.transaction(async () => {
+      await this.#data.clearAsync();
+      await this.#multiMapData.clearAsync();
+      await this.#rootDb.clearAsync();
+    });
+  }
+
+  /**
+   * Drops the database & sub DBs atomically within a single transaction.
+   */
+  async drop() {
+    await this.#rootDb.transaction(async () => {
+      await this.#data.drop();
+      await this.#multiMapData.drop();
+      await this.#rootDb.drop();
+    });
+  }
+
+  /**
+   * Close the database. Note, once this is closed we can no longer interact with the DB.
+   */
+  async close() {
+    await this.#data.close();
+    await this.#multiMapData.close();
+    await this.#rootDb.close();
+  }
+
+  /** Deletes this store and removes the database files from disk */
+  async delete() {
+    await this.drop();
+    await this.close();
+    if (this.path) {
+      await fs.rm(this.path, { recursive: true, force: true, maxRetries: 3 });
+      this.#log.verbose(`Deleted database files at ${this.path}`);
+    }
+  }
+
+  estimateSize(): Promise<StoreSize> {
+    const stats = this.#rootDb.getStats();
+    // The 'mapSize' is the total amount of virtual address space allocated to the DB (effectively the maximum possible size)
+    // http://www.lmdb.tech/doc/group__mdb.html#a4bde3c8b676457342cba2fe27aed5fbd
+    let mapSize = 0;
+    if ('mapSize' in stats && typeof stats.mapSize === 'number') {
+      mapSize = stats.mapSize;
+    }
+    let physicalFileSize = 0;
+    if ('physicalFileSize' in stats && typeof stats.physicalFileSize === 'number') {
+      physicalFileSize = stats.physicalFileSize;
+    }
+    const dataResult = this.estimateSubDBSize(this.#data);
+    const multiResult = this.estimateSubDBSize(this.#multiMapData);
+    return Promise.resolve({
+      mappingSize: mapSize,
+      physicalFileSize: physicalFileSize,
+      actualSize: dataResult.actualSize + multiResult.actualSize,
+      numItems: dataResult.numItems + multiResult.numItems,
+    });
+  }
+
+  private estimateSubDBSize(db: Database<unknown, Key>): { actualSize: number; numItems: number } {
+    const stats = db.getStats();
+    let actualSize = 0;
+    let numItems = 0;
+    // This is the total number of key/value pairs present in the DB
+    if ('entryCount' in stats && typeof stats.entryCount === 'number') {
+      numItems = stats.entryCount;
+    }
+    // The closest value we can get to the actual size of the database is the number of consumed pages * the page size
+    if (
+      'treeBranchPageCount' in stats &&
+      typeof stats.treeBranchPageCount === 'number' &&
+      'treeLeafPageCount' in stats &&
+      typeof stats.treeLeafPageCount === 'number' &&
+      'overflowPages' in stats &&
+      typeof stats.overflowPages === 'number' &&
+      'pageSize' in stats &&
+      typeof stats.pageSize === 'number'
+    ) {
+      const branchPages = stats.treeBranchPageCount;
+      const leafPages = stats.treeLeafPageCount;
+      const overflowPages = stats.overflowPages;
+      const pageSize = stats.pageSize;
+      actualSize = (branchPages + leafPages + overflowPages) * pageSize;
+    }
+    return { actualSize, numItems };
+  }
+
+  backupTo(_dstPath: string, _compact?: boolean): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
+}

package/src/lmdb-v2/array.ts

@@ -0,0 +1,120 @@
+import { Encoder } from 'msgpackr/pack';
+
+import type { AztecAsyncArray } from '../interfaces/array.js';
+import type { Value } from '../interfaces/common.js';
+import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
+import type { ReadTransaction } from './read_transaction.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
+import { deserializeKey, serializeKey } from './utils.js';
+
+export class LMDBArray<T extends Value> implements AztecAsyncArray<T> {
+  private length: AztecAsyncSingleton<number>;
+  private encoder = new Encoder();
+  private prefix: string;
+
+  constructor(
+    private store: AztecLMDBStoreV2,
+    name: string,
+  ) {
+    this.length = store.openSingleton(name + ':length');
+    this.prefix = `array:${name}`;
+  }
+
+  pop(): Promise<T | undefined> {
+    return execInWriteTx(this.store, async tx => {
+      const length = await this.lengthAsync();
+      if (length === 0) {
+        return;
+      }
+
+      const val = await tx.get(serializeKey(this.prefix, length - 1));
+      await tx.remove(serializeKey(this.prefix, length - 1));
+
+      await this.length.set(length - 1);
+
+      return val ? this.encoder.unpack(val) : undefined;
+    });
+  }
+
+  push(...vals: T[]): Promise<number> {
+    return execInWriteTx(this.store, async tx => {
+      let length = await this.lengthAsync();
+      for (const val of vals) {
+        await tx.set(serializeKey(this.prefix, length++), this.encoder.pack(val));
+      }
+      await this.length.set(length);
+      return length;
+    });
+  }
+
+  setAt(index: number, val: T): Promise<boolean> {
+    return execInWriteTx(this.store, async tx => {
+      const length = await this.lengthAsync();
+      if (index < 0) {
+        index += length;
+      }
+      if (index < 0 || index >= length) {
+        return false;
+      }
+      await tx.set(serializeKey(this.prefix, index), this.encoder.pack(val));
+      return true;
+    });
+  }
+
+  atAsync(index: number): Promise<T | undefined> {
+    return execInReadTx(this.store, async tx => {
+      const length = await this.lengthAsync();
+      if (index < 0) {
+        index += length;
+      }
+      if (index < 0 || index >= length) {
+        return undefined;
+      }
+
+      const val = await tx.get(serializeKey(this.prefix, index));
+      return val ? this.encoder.unpack(val) : undefined;
+    });
+  }
+
+  async lengthAsync(): Promise<number> {
+    return (await this.length.getAsync()) ?? 0;
+  }
+
+  async *entriesAsync(): AsyncIterableIterator<[number, T]> {
+    // pin array length so that pushes don't affect iteration
+    const length = await this.lengthAsync();
+    if (length === 0) {
+      return;
+    }
+
+    let tx: ReadTransaction | undefined = this.store.getCurrentWriteTx();
+    const shouldClose = !tx;
+    tx ??= this.store.getReadTx();
+
+    try {
+      for await (const [key, val] of tx.iterate(serializeKey(this.prefix, 0), undefined, false, length)) {
+        const deserializedKey = deserializeKey<number>(this.prefix, key);
+        // if pops happened while iterating we may have read too much. Terminate early
+        if (deserializedKey === false) {
+          break;
+        }
+        yield [deserializedKey, this.encoder.unpack(val)];
+      }
+    } finally {
+      if (shouldClose) {
+        tx.close();
+      }
+    }
+  }
+
+  async *valuesAsync(): AsyncIterableIterator<T> {
+    for await (const [_, value] of this.entriesAsync()) {
+      yield value;
+    }
+  }
+
+  [Symbol.asyncIterator](): AsyncIterableIterator<T> {
+    return this.valuesAsync();
+  }
+}