@aztec/kv-store 0.0.0-test.0 → 0.0.1-commit.001888fc

package/src/lmdb/map.ts

@@ -1,7 +1,7 @@
 import type { Database, RangeOptions } from 'lmdb';
 
-import type { Key, Range } from '../interfaces/common.js';
-import type { AztecAsyncMultiMap, AztecMapWithSize, AztecMultiMap } from '../interfaces/map.js';
+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];
@@ -9,7 +9,7 @@ type MapValueSlot<K extends Key | Buffer> = ['map', string, 'slot', K];
 /**
  * A map backed by LMDB.
  */
-export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, AztecAsyncMultiMap<K, V> {
+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;
 
@@ -39,26 +39,6 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     return Promise.resolve(this.get(key));
   }
 
-  *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;
-    }
-  }
-
   has(key: K): boolean {
     return this.db.doesExist(this.slot(key));
   }
@@ -71,6 +51,12 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     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);
@@ -90,10 +76,6 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     await this.db.remove(this.slot(key));
   }
 
-  async deleteValue(key: K, val: V): Promise<void> {
-    await this.db.remove(this.slot(key), [key, val]);
-  }
-
   *entries(range: Range<K> = {}): IterableIterator<[K, V]> {
     const transaction = this.db.useReadTransaction();
 
@@ -106,16 +88,16 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
           ? this.slot(range.end)
           : this.endSentinel
         : range.start
-        ? this.slot(range.start)
-        : this.startSentinel;
+          ? 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;
+          ? this.slot(range.end)
+          : this.endSentinel;
 
       const lmdbRange: RangeOptions = {
         start,
@@ -137,7 +119,7 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     }
   }
 
-  async *entriesAsync(range?: Range<K> | undefined): AsyncIterableIterator<[K, V]> {
+  async *entriesAsync(range?: Range<K>): AsyncIterableIterator<[K, V]> {
     for (const entry of this.entries(range)) {
       yield entry;
     }
@@ -155,6 +137,15 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     }
   }
 
+  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;
@@ -184,65 +175,3 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     }
   }
 }
-
-export class LmdbAztecMapWithSize<K extends Key, V>
-  extends LmdbAztecMap<K, V>
-  implements AztecMapWithSize<K, V>, AztecAsyncMultiMap<K, V>
-{
-  #sizeCache?: number;
-
-  constructor(rootDb: Database, mapName: string) {
-    super(rootDb, mapName);
-  }
-
-  override async set(key: K, val: V): Promise<void> {
-    await this.db.childTransaction(() => {
-      const exists = this.db.doesExist(this.slot(key));
-      this.db.putSync(this.slot(key), [key, val], {
-        appendDup: true,
-      });
-      if (!exists) {
-        this.#sizeCache = undefined; // Invalidate cache
-      }
-    });
-  }
-
-  override async delete(key: K): Promise<void> {
-    await this.db.childTransaction(async () => {
-      const exists = this.db.doesExist(this.slot(key));
-      if (exists) {
-        await this.db.remove(this.slot(key));
-        this.#sizeCache = undefined; // Invalidate cache
-      }
-    });
-  }
-
-  override async deleteValue(key: K, val: V): Promise<void> {
-    await this.db.childTransaction(async () => {
-      const exists = this.db.doesExist(this.slot(key));
-      if (exists) {
-        await this.db.remove(this.slot(key), [key, val]);
-        this.#sizeCache = undefined; // Invalidate cache
-      }
-    });
-  }
-
-  /**
-   * Gets the size of the map by counting entries.
-   * @returns The number of entries in the map
-   */
-  size(): number {
-    if (this.#sizeCache === undefined) {
-      this.#sizeCache = this.db.getCount({
-        start: this.startSentinel,
-        end: this.endSentinel,
-      });
-    }
-    return this.#sizeCache;
-  }
-
-  // Reset cache on clear/drop operations
-  clearCache() {
-    this.#sizeCache = undefined;
-  }
-}

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/store.ts

@@ -1,4 +1,4 @@
-import { randomBytes } from '@aztec/foundation/crypto';
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import { createLogger } from '@aztec/foundation/log';
 
 import { promises as fs, mkdirSync } from 'fs';
@@ -7,22 +7,17 @@ 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 { Key, StoreSize, Value } 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 { 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, LmdbAztecMapWithSize } from './map.js';
+import { LmdbAztecMap } from './map.js';
+import { LmdbAztecMultiMap } from './multi_map.js';
 import { LmdbAztecSet } from './set.js';
 import { LmdbAztecSingleton } from './singleton.js';
 
@@ -37,7 +32,11 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
   #multiMapData: Database<unknown, Key>;
   #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
@@ -80,27 +79,12 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
     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> {
+  openMap<K extends Key, V extends Value>(name: string): AztecMap<K, V> & AztecAsyncMap<K, V> {
     return new LmdbAztecMap(this.#data, name);
   }
 
@@ -118,37 +102,20 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
    * @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);
+  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 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> {
+  openArray<T extends Value>(name: string): AztecArray<T> & AztecAsyncArray<T> {
     return new LmdbAztecArray(this.#data, name);
   }
 
@@ -180,21 +147,25 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
   }
 
   /**
-   * Clears all entries in the store & sub DBs.
+   * Clears all entries in the store & sub DBs atomically within a single transaction.
    */
   async clear() {
-    await this.#data.clearAsync();
-    await this.#multiMapData.clearAsync();
-    await this.#rootDb.clearAsync();
+    await this.#rootDb.transaction(async () => {
+      await this.#data.clearAsync();
+      await this.#multiMapData.clearAsync();
+      await this.#rootDb.clearAsync();
+    });
   }
 
   /**
-   * Drops the database & sub DBs.
+   * Drops the database & sub DBs atomically within a single transaction.
    */
   async drop() {
-    await this.#data.drop();
-    await this.#multiMapData.drop();
-    await this.#rootDb.drop();
+    await this.#rootDb.transaction(async () => {
+      await this.#data.drop();
+      await this.#multiMapData.drop();
+      await this.#rootDb.drop();
+    });
   }
 
   /**
@@ -224,10 +195,15 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
     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,
     });
@@ -260,4 +236,8 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
     }
     return { actualSize, numItems };
   }
+
+  backupTo(_dstPath: string, _compact?: boolean): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
 }

package/src/lmdb-v2/array.ts

@@ -1,17 +1,22 @@
 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 { AztecLMDBStoreV2, execInReadTx, execInWriteTx } from './store.js';
+import type { AztecLMDBStoreV2 } from './store.js';
+import { execInReadTx, execInWriteTx } from './tx-helpers.js';
 import { deserializeKey, serializeKey } from './utils.js';
 
-export class LMDBArray<T> implements AztecAsyncArray<T> {
+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) {
+  constructor(
+    private store: AztecLMDBStoreV2,
+    name: string,
+  ) {
     this.length = store.openSingleton(name + ':length');
     this.prefix = `array:${name}`;
   }

package/src/lmdb-v2/factory.ts

@@ -1,12 +1,12 @@
 import { EthAddress } from '@aztec/foundation/eth-address';
-import { type Logger, createLogger } from '@aztec/foundation/log';
-import { DatabaseVersionManager } from '@aztec/stdlib/database-version';
+import { type LoggerBindings, createLogger } from '@aztec/foundation/log';
+import { DatabaseVersionManager } from '@aztec/stdlib/database-version/manager';
+import type { DataStoreConfig } from '@aztec/stdlib/kv-store';
 
 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;
@@ -15,8 +15,9 @@ export async function createStore(
   name: string,
   schemaVersion: number,
   config: DataStoreConfig,
-  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
   const { dataDirectory, l1Contracts } = config;
 
   let store: AztecLMDBStoreV2;
@@ -28,16 +29,20 @@ export async function createStore(
     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),
-    );
+    const versionManager = new DatabaseVersionManager({
+      schemaVersion,
+      rollupAddress,
+      dataDirectory: subDir,
+      onOpen: dbDirectory =>
+        AztecLMDBStoreV2.new(dbDirectory, config.dataStoreMapSizeKb, MAX_READERS, () => Promise.resolve(), bindings),
+    });
 
     log.info(
-      `Creating ${name} data store at directory ${subDir} with map size ${config.dataStoreMapSizeKB} KB (LMDB v2)`,
+      `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);
+    store = await openTmpStore(name, true, config.dataStoreMapSizeKb, MAX_READERS, bindings);
   }
 
   return store;
@@ -48,8 +53,9 @@ export async function openTmpStore(
   ephemeral: boolean = true,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2:' + name),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2:' + name, bindings);
   const dataDir = await mkdtemp(join(tmpdir(), name + '-'));
   log.debug(`Created temporary data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
 
@@ -69,30 +75,35 @@ export async function openTmpStore(
 
   // 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);
+  return AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, cleanup, bindings);
 }
 
 export async function openStoreAt(
   dataDir: string,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2'),
+  bindings?: LoggerBindings,
 ): Promise<AztecLMDBStoreV2> {
+  const log = createLogger('kv-store:lmdb-v2', bindings);
   log.debug(`Opening data store at: ${dataDir} with size: ${dbMapSizeKb} KB (LMDB v2)`);
-  return await AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, log);
+  return await AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, bindings);
 }
 
 export async function openVersionedStoreAt(
-  dataDir: string,
+  dataDirectory: string,
   schemaVersion: number,
   rollupAddress: EthAddress,
   dbMapSizeKb = 10 * 1_024 * 1_024, // 10GB
   maxReaders = MAX_READERS,
-  log: Logger = createLogger('kv-store:lmdb-v2'),
+  bindings?: LoggerBindings,
 ): 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();
+  const log = createLogger('kv-store:lmdb-v2', bindings);
+  log.debug(`Opening data store at: ${dataDirectory} with size: ${dbMapSizeKb} KB (LMDB v2)`);
+  const [store] = await new DatabaseVersionManager({
+    schemaVersion,
+    rollupAddress,
+    dataDirectory,
+    onOpen: dataDir => AztecLMDBStoreV2.new(dataDir, dbMapSizeKb, maxReaders, undefined, bindings),
+  }).open();
   return store;
 }