@aztec/kv-store 0.66.0 → 0.67.1-devnet

package/src/lmdb/store.ts

@@ -1,33 +1,43 @@
-import { createDebugLogger } from '@aztec/foundation/log';
+import { randomBytes } from '@aztec/foundation/crypto';
+import { createLogger } from '@aztec/foundation/log';
 
-import { mkdirSync } from 'fs';
-import { mkdtemp, rm } from 'fs/promises';
-import { type Database, type Key, type RootDatabase, open } from 'lmdb';
+import { promises as fs, mkdirSync } from 'fs';
+import { type Database, type RootDatabase, open } from 'lmdb';
 import { tmpdir } from 'os';
-import { dirname, join } from 'path';
-
-import { type AztecArray } from '../interfaces/array.js';
-import { type AztecCounter } from '../interfaces/counter.js';
-import { type AztecMap, type AztecMultiMap } from '../interfaces/map.js';
-import { type AztecSet } from '../interfaces/set.js';
-import { type AztecSingleton } from '../interfaces/singleton.js';
-import { type AztecKVStore } from '../interfaces/store.js';
+import { join } from 'path';
+
+import { type AztecArray, type AztecAsyncArray } from '../interfaces/array.js';
+import { type Key } from '../interfaces/common.js';
+import { type AztecAsyncCounter, type AztecCounter } from '../interfaces/counter.js';
+import {
+  type AztecAsyncMap,
+  type AztecAsyncMultiMap,
+  type AztecMap,
+  type AztecMapWithSize,
+  type AztecMultiMap,
+  type AztecMultiMapWithSize,
+} from '../interfaces/map.js';
+import { type AztecAsyncSet, type AztecSet } from '../interfaces/set.js';
+import { type AztecAsyncSingleton, type AztecSingleton } from '../interfaces/singleton.js';
+import { type AztecAsyncKVStore, type AztecKVStore } from '../interfaces/store.js';
 import { LmdbAztecArray } from './array.js';
 import { LmdbAztecCounter } from './counter.js';
-import { LmdbAztecMap } from './map.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 {
+export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
+  syncGetters = true as const;
+
   #rootDb: RootDatabase;
   #data: Database<unknown, Key>;
   #multiMapData: Database<unknown, Key>;
-  #log = createDebugLogger('aztec:kv-store:lmdb');
+  #log = createLogger('kv-store:lmdb');
 
-  constructor(rootDb: RootDatabase, public readonly isEphemeral: boolean, private path?: string) {
+  constructor(rootDb: RootDatabase, public readonly isEphemeral: boolean, private path: string) {
     this.#rootDb = rootDb;
 
     // big bucket to store all the data
@@ -60,15 +70,14 @@ export class AztecLmdbStore implements AztecKVStore {
     path?: string,
     mapSizeKb = 1 * 1024 * 1024, // defaults to 1 GB map size
     ephemeral: boolean = false,
-    log = createDebugLogger('aztec:kv-store:lmdb'),
+    log = createLogger('kv-store:lmdb'),
   ): AztecLmdbStore {
-    if (path) {
-      mkdirSync(path, { recursive: true });
-    }
+    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, noSync: ephemeral, mapSize });
-    return new AztecLmdbStore(rootDb, ephemeral, path);
+    const rootDb = open({ path: dbPath, noSync: ephemeral, mapSize });
+    return new AztecLmdbStore(rootDb, ephemeral, dbPath);
   }
 
   /**
@@ -76,10 +85,9 @@ export class AztecLmdbStore implements AztecKVStore {
    * @returns A new AztecLmdbStore.
    */
   async fork() {
-    const baseDir = this.path ? dirname(this.path) : tmpdir();
+    const baseDir = this.path;
     this.#log.debug(`Forking store with basedir ${baseDir}`);
-    const forkPath =
-      (await mkdtemp(join(baseDir, 'aztec-store-fork-'))) + (this.isEphemeral || !this.path ? '/data.mdb' : '');
+    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 });
@@ -92,7 +100,7 @@ export class AztecLmdbStore implements AztecKVStore {
    * @param name - Name of the map
    * @returns A new AztecMap
    */
-  openMap<K extends string | number, V>(name: string): AztecMap<K, V> {
+  openMap<K extends Key, V>(name: string): AztecMap<K, V> & AztecAsyncMap<K, V> {
     return new LmdbAztecMap(this.#data, name);
   }
 
@@ -101,7 +109,7 @@ export class AztecLmdbStore implements AztecKVStore {
    * @param name - Name of the set
    * @returns A new AztecSet
    */
-  openSet<K extends string | number>(name: string): AztecSet<K> {
+  openSet<K extends Key>(name: string): AztecSet<K> & AztecAsyncSet<K> {
     return new LmdbAztecSet(this.#data, name);
   }
 
@@ -110,20 +118,37 @@ export class AztecLmdbStore implements AztecKVStore {
    * @param name - Name of the map
    * @returns A new AztecMultiMap
    */
-  openMultiMap<K extends string | number, V>(name: string): AztecMultiMap<K, V> {
+  openMultiMap<K extends Key, V>(name: string): AztecMultiMap<K, V> & AztecAsyncMultiMap<K, V> {
     return new LmdbAztecMap(this.#multiMapData, name);
   }
 
-  openCounter<K extends string | number | Array<string | number>>(name: string): AztecCounter<K> {
+  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> {
+  openArray<T>(name: string): AztecArray<T> & AztecAsyncArray<T> {
     return new LmdbAztecArray(this.#data, name);
   }
 
@@ -132,7 +157,7 @@ export class AztecLmdbStore implements AztecKVStore {
    * @param name - Name of the singleton
    * @returns A new AztecSingleton
    */
-  openSingleton<T>(name: string): AztecSingleton<T> {
+  openSingleton<T>(name: string): AztecSingleton<T> & AztecAsyncSingleton<T> {
     return new LmdbAztecSingleton(this.#data, name);
   }
 
@@ -145,6 +170,15 @@ export class AztecLmdbStore implements AztecKVStore {
     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.
    */
@@ -177,7 +211,7 @@ export class AztecLmdbStore implements AztecKVStore {
     await this.drop();
     await this.close();
     if (this.path) {
-      await rm(this.path, { recursive: true, force: true });
+      await fs.rm(this.path, { recursive: true, force: true });
       this.#log.verbose(`Deleted database files at ${this.path}`);
     }
   }
@@ -201,11 +235,7 @@ export class AztecLmdbStore implements AztecKVStore {
 
   private estimateSubDBSize(db: Database<unknown, Key>): { actualSize: number; numItems: number } {
     const stats = db.getStats();
-    let branchPages = 0;
-    let leafPages = 0;
-    let overflowPages = 0;
-    let pageSize = 0;
-    let totalSize = 0;
+    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') {
@@ -222,12 +252,12 @@ export class AztecLmdbStore implements AztecKVStore {
       'pageSize' in stats &&
       typeof stats.pageSize === 'number'
     ) {
-      branchPages = stats.treeBranchPageCount;
-      leafPages = stats.treeLeafPageCount;
-      overflowPages = stats.overflowPages;
-      pageSize = stats.pageSize;
-      totalSize = (branchPages + leafPages + overflowPages) * pageSize;
+      const branchPages = stats.treeBranchPageCount;
+      const leafPages = stats.treeLeafPageCount;
+      const overflowPages = stats.overflowPages;
+      const pageSize = stats.pageSize;
+      actualSize = (branchPages + leafPages + overflowPages) * pageSize;
     }
-    return { actualSize: totalSize, numItems };
+    return { actualSize, numItems };
   }
 }

package/src/stores/l2_tips_store.ts

@@ -7,37 +7,37 @@ import {
   type L2Tips,
 } from '@aztec/circuit-types';
 
-import { type AztecMap } from '../interfaces/map.js';
-import { type AztecKVStore } from '../interfaces/store.js';
+import { type AztecAsyncMap } from '../interfaces/map.js';
+import { type AztecAsyncKVStore } from '../interfaces/store.js';
 
 /** Stores currently synced L2 tips and unfinalized block hashes. */
 export class L2TipsStore implements L2BlockStreamEventHandler, L2BlockStreamLocalDataProvider {
-  private readonly l2TipsStore: AztecMap<L2BlockTag, number>;
-  private readonly l2BlockHashesStore: AztecMap<number, string>;
+  private readonly l2TipsStore: AztecAsyncMap<L2BlockTag, number>;
+  private readonly l2BlockHashesStore: AztecAsyncMap<number, string>;
 
-  constructor(store: AztecKVStore, namespace: string) {
+  constructor(store: AztecAsyncKVStore, namespace: string) {
     this.l2TipsStore = store.openMap([namespace, 'l2_tips'].join('_'));
     this.l2BlockHashesStore = store.openMap([namespace, 'l2_block_hashes'].join('_'));
   }
 
   public getL2BlockHash(number: number): Promise<string | undefined> {
-    return Promise.resolve(this.l2BlockHashesStore.get(number));
+    return this.l2BlockHashesStore.getAsync(number);
   }
 
-  public getL2Tips(): Promise<L2Tips> {
-    return Promise.resolve({
-      latest: this.getL2Tip('latest'),
-      finalized: this.getL2Tip('finalized'),
-      proven: this.getL2Tip('proven'),
-    });
+  public async getL2Tips(): Promise<L2Tips> {
+    return {
+      latest: await this.getL2Tip('latest'),
+      finalized: await this.getL2Tip('finalized'),
+      proven: await this.getL2Tip('proven'),
+    };
   }
 
-  private getL2Tip(tag: L2BlockTag): L2BlockId {
-    const blockNumber = this.l2TipsStore.get(tag);
+  private async getL2Tip(tag: L2BlockTag): Promise<L2BlockId> {
+    const blockNumber = await this.l2TipsStore.getAsync(tag);
     if (blockNumber === undefined || blockNumber === 0) {
       return { number: 0, hash: undefined };
     }
-    const blockHash = this.l2BlockHashesStore.get(blockNumber);
+    const blockHash = await this.l2BlockHashesStore.getAsync(blockNumber);
     if (!blockHash) {
       throw new Error(`Block hash not found for block number ${blockNumber}`);
     }
@@ -47,10 +47,10 @@ export class L2TipsStore implements L2BlockStreamEventHandler, L2BlockStreamLoca
   public async handleBlockStreamEvent(event: L2BlockStreamEvent): Promise<void> {
     switch (event.type) {
       case 'blocks-added':
-        await this.l2TipsStore.set('latest', event.blocks.at(-1)!.number);
         for (const block of event.blocks) {
           await this.l2BlockHashesStore.set(block.number, block.header.hash().toString());
         }
+        await this.l2TipsStore.set('latest', event.blocks.at(-1)!.number);
         break;
       case 'chain-pruned':
         await this.l2TipsStore.set('latest', event.blockNumber);
@@ -60,7 +60,7 @@ export class L2TipsStore implements L2BlockStreamEventHandler, L2BlockStreamLoca
         break;
       case 'chain-finalized':
         await this.l2TipsStore.set('finalized', event.blockNumber);
-        for (const key of this.l2BlockHashesStore.keys({ end: event.blockNumber })) {
+        for await (const key of this.l2BlockHashesStore.keysAsync({ end: event.blockNumber })) {
           await this.l2BlockHashesStore.delete(key);
         }
         break;

package/src/utils.ts

@@ -1,30 +1,9 @@
 import { type EthAddress } from '@aztec/foundation/eth-address';
-import { type Logger, createDebugLogger } from '@aztec/foundation/log';
+import { type Logger } from '@aztec/foundation/log';
 
-import { join } from 'path';
-
-import { type DataStoreConfig } from './config.js';
-import { type AztecKVStore } from './interfaces/store.js';
-import { AztecLmdbStore } from './lmdb/store.js';
-
-export function createStore(name: string, config: DataStoreConfig, log: Logger = createDebugLogger('aztec:kv-store')) {
-  let { dataDirectory } = config;
-  if (typeof dataDirectory !== 'undefined') {
-    dataDirectory = join(dataDirectory, name);
-  }
-
-  log.info(
-    dataDirectory
-      ? `Creating ${name} data store at directory ${dataDirectory} with map size ${config.dataStoreMapSizeKB} KB`
-      : `Creating ${name} ephemeral data store with map size ${config.dataStoreMapSizeKB} KB`,
-  );
-
-  const store = AztecLmdbStore.open(dataDirectory, config.dataStoreMapSizeKB, false);
-  if (config.l1Contracts?.rollupAddress) {
-    return initStoreForRollup(store, config.l1Contracts.rollupAddress, log);
-  }
-  return store;
-}
+import { type AztecAsyncSingleton, type AztecSingleton } from './interfaces/singleton.js';
+import { type AztecAsyncKVStore, type AztecKVStore } from './interfaces/store.js';
+import { isSyncStore } from './interfaces/utils.js';
 
 /**
  * Clears the store if the rollup address does not match the one stored in the database.
@@ -33,7 +12,7 @@ export function createStore(name: string, config: DataStoreConfig, log: Logger =
  * @param rollupAddress - The ETH address of the rollup contract
  * @returns A promise that resolves when the store is cleared, or rejects if the rollup address does not match
  */
-async function initStoreForRollup<T extends AztecKVStore>(
+export async function initStoreForRollup<T extends AztecKVStore | AztecAsyncKVStore>(
   store: T,
   rollupAddress: EthAddress,
   log?: Logger,
@@ -43,7 +22,9 @@ async function initStoreForRollup<T extends AztecKVStore>(
   }
   const rollupAddressValue = store.openSingleton<ReturnType<EthAddress['toString']>>('rollupAddress');
   const rollupAddressString = rollupAddress.toString();
-  const storedRollupAddressString = rollupAddressValue.get();
+  const storedRollupAddressString = isSyncStore(store)
+    ? (rollupAddressValue as AztecSingleton<ReturnType<EthAddress['toString']>>).get()
+    : await (rollupAddressValue as AztecAsyncSingleton<ReturnType<EthAddress['toString']>>).getAsync();
 
   if (typeof storedRollupAddressString !== 'undefined' && storedRollupAddressString !== rollupAddressString) {
     log?.warn(`Rollup address mismatch. Clearing entire database...`, {
@@ -57,13 +38,3 @@ async function initStoreForRollup<T extends AztecKVStore>(
   await rollupAddressValue.set(rollupAddressString);
   return store;
 }
-
-/**
- * Opens a temporary store for testing purposes.
- * @param ephemeral - true if the store should only exist in memory and not automatically be flushed to disk. Optional
- * @returns A new store
- */
-export function openTmpStore(ephemeral: boolean = false): AztecLmdbStore {
-  const mapSize = 1024 * 1024 * 10; // 10 GB map size
-  return AztecLmdbStore.open(undefined, mapSize, ephemeral);
-}