@aztec/kv-store 0.0.0-test.0

package/src/lmdb/store.ts

@@ -0,0 +1,263 @@
+import { randomBytes } from '@aztec/foundation/crypto';
+import { createLogger } from '@aztec/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 } from '../interfaces/common.js';
+import type { AztecAsyncCounter, AztecCounter } from '../interfaces/counter.js';
+import type {
+  AztecAsyncMap,
+  AztecAsyncMultiMap,
+  AztecMap,
+  AztecMapWithSize,
+  AztecMultiMap,
+  AztecMultiMapWithSize,
+} from '../interfaces/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, LmdbAztecMapWithSize } from './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);
+  }
+
+  /**
+   * Forks the current DB into a new DB by backing it up to a temporary location and opening a new lmdb db.
+   * @returns A new AztecLmdbStore.
+   */
+  async fork() {
+    const baseDir = this.path;
+    this.#log.debug(`Forking store with basedir ${baseDir}`);
+    const forkPath = await fs.mkdtemp(join(baseDir, 'aztec-store-fork-'));
+    this.#log.verbose(`Forking store to ${forkPath}`);
+    await this.#rootDb.backup(forkPath, false);
+    const forkDb = open(forkPath, { noSync: this.isEphemeral });
+    this.#log.debug(`Forked store at ${forkPath} opened successfully`);
+    return new AztecLmdbStore(forkDb, this.isEphemeral, forkPath);
+  }
+
+  /**
+   * Creates a new AztecMap in the store.
+   * @param name - Name of the map
+   * @returns A new AztecMap
+   */
+  openMap<K extends Key, V>(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>(name: string): AztecMultiMap<K, V> & AztecAsyncMultiMap<K, V> {
+    return new LmdbAztecMap(this.#multiMapData, name);
+  }
+
+  openCounter<K extends Key>(name: string): AztecCounter<K> & AztecAsyncCounter<K> {
+    return new LmdbAztecCounter(this.#data, name);
+  }
+  /**
+   * Creates a new AztecMultiMapWithSize in the store. A multi-map with size stores multiple values for a single key automatically.
+   * @param name - Name of the map
+   * @returns A new AztecMultiMapWithSize
+   */
+  openMultiMapWithSize<K extends Key, V>(name: string): AztecMultiMapWithSize<K, V> {
+    return new LmdbAztecMapWithSize(this.#multiMapData, name);
+  }
+
+  /**
+   * Creates a new AztecMapWithSize in the store.
+   * @param name - Name of the map
+   * @returns A new AztecMapWithSize
+   */
+  openMapWithSize<K extends Key, V>(name: string): AztecMapWithSize<K, V> {
+    return new LmdbAztecMapWithSize(this.#data, name);
+  }
+
+  /**
+   * Creates a new AztecArray in the store.
+   * @param name - Name of the array
+   * @returns A new AztecArray
+   */
+  openArray<T>(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.
+   */
+  async clear() {
+    await this.#data.clearAsync();
+    await this.#multiMapData.clearAsync();
+    await this.#rootDb.clearAsync();
+  }
+
+  /**
+   * Drops the database & sub DBs.
+   */
+  async drop() {
+    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;
+    }
+    const dataResult = this.estimateSubDBSize(this.#data);
+    const multiResult = this.estimateSubDBSize(this.#multiMapData);
+    return Promise.resolve({
+      mappingSize: mapSize,
+      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 };
+  }
+}

package/src/lmdb-v2/array.ts

@@ -0,0 +1,115 @@
+import { Encoder } from 'msgpackr/pack';
+
+import type { AztecAsyncArray } from '../interfaces/array.js';
+import type { AztecAsyncSingleton } from '../interfaces/singleton.js';
+import type { ReadTransaction } from './read_transaction.js';
+import { AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import { deserializeKey, serializeKey } from './utils.js';
+
+export class LMDBArray<T> 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();
+  }
+}

package/src/lmdb-v2/factory.ts

@@ -0,0 +1,98 @@
+import { EthAddress } from '@aztec/foundation/eth-address';
+import { type Logger, createLogger } from '@aztec/foundation/log';
+import { DatabaseVersionManager } from '@aztec/stdlib/database-version';
+
+import { mkdir, mkdtemp, rm } from 'fs/promises';
+import { tmpdir } from 'os';
+import { join } from 'path';
+
+import type { DataStoreConfig } from '../config.js';
+import { AztecLMDBStoreV2 } from './store.js';
+
+const MAX_READERS = 16;
+
+export async function createStore(
+  name: string,
+  schemaVersion: number,
+  config: DataStoreConfig,
+  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+): Promise<AztecLMDBStoreV2> {
+  const { dataDirectory, l1Contracts } = config;
+
+  let store: AztecLMDBStoreV2;
+  if (typeof dataDirectory !== 'undefined') {
+    // Get rollup address from contracts config, or use zero address
+    const subDir = join(dataDirectory, name);
+    await mkdir(subDir, { recursive: true });
+
+    const rollupAddress = l1Contracts ? l1Contracts.rollupAddress : EthAddress.ZERO;
+
+    // Create a version manager
+    const versionManager = new DatabaseVersionManager(schemaVersion, rollupAddress, subDir, dbDirectory =>
+      AztecLMDBStoreV2.new(dbDirectory, config.dataStoreMapSizeKB, MAX_READERS, () => Promise.resolve(), log),
+    );
+
+    log.info(
+      `Creating ${name} data store at directory ${subDir} with map size ${config.dataStoreMapSizeKB} KB (LMDB v2)`,
+    );
+    [store] = await versionManager.open();
+  } else {
+    store = await openTmpStore(name, true, config.dataStoreMapSizeKB, MAX_READERS, log);
+  }
+
+  return store;
+}
+
+export async function openTmpStore(
+  name: string,
+  ephemeral: boolean = true,
+  dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
+  maxReaders = MAX_READERS,
+  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+): Promise<AztecLMDBStoreV2> {
+  const dataDir = await mkdtemp(join(tmpdir(), name + '-'));
+  log.debug(`Created temporary data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
+
+  // pass a cleanup callback because process.on('beforeExit', cleanup) does not work under Jest
+  const cleanup = async () => {
+    if (ephemeral) {
+      try {
+        await rm(dataDir, { recursive: true, force: true, maxRetries: 3 });
+        log.debug(`Deleted temporary data store: ${dataDir}`);
+      } catch (err) {
+        log.warn(`Failed to delete temporary data directory (LMDB v2) ${dataDir}: ${err}`);
+      }
+    } else {
+      log.debug(`Leaving temporary data store: ${dataDir}`);
+    }
+  };
+
+  // For temporary stores, we don't need to worry about versioning
+  // as they are ephemeral and get cleaned up after use
+  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, log);
+}
+
+export async function openStoreAt(
+  dataDir: string,
+  dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
+  maxReaders = MAX_READERS,
+  log: Logger = createLogger('kv-store:lmdb-v2'),
+): Promise<AztecLMDBStoreV2> {
+  log.debug(`Opening data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
+  return await AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, log);
+}
+
+export async function openVersionedStoreAt(
+  dataDir: string,
+  schemaVersion: number,
+  rollupAddress: EthAddress,
+  dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
+  maxReaders = MAX_READERS,
+  log: Logger = createLogger('kv-store:lmdb-v2'),
+): Promise<AztecLMDBStoreV2> {
+  log.debug(`Opening data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
+  const [store] = await new DatabaseVersionManager(schemaVersion, rollupAddress, dataDir, dataDir =>
+    AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, log),
+  ).open();
+  return store;
+}

package/src/lmdb-v2/index.ts

@@ -0,0 +1,2 @@
+export * from './store.js';
+export * from './factory.js';