@aztec/kv-store 0.67.1 → 0.68.0

package/src/lmdb/map.ts

@@ -1,7 +1,7 @@
 import { type Database, type RangeOptions } from 'lmdb';
 
 import { type Key, type Range } from '../interfaces/common.js';
-import { type AztecAsyncMultiMap, type AztecMultiMap } from '../interfaces/map.js';
+import { type AztecAsyncMultiMap, type AztecMapWithSize, type AztecMultiMap } from '../interfaces/map.js';
 
 /** The slot where a key-value entry would be stored */
 type MapValueSlot<K extends Key | Buffer> = ['map', string, 'slot', K];
@@ -13,8 +13,8 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
   protected db: Database<[K, V], MapValueSlot<K>>;
   protected name: string;
 
-  #startSentinel: MapValueSlot<Buffer>;
-  #endSentinel: MapValueSlot<Buffer>;
+  protected startSentinel: MapValueSlot<Buffer>;
+  protected endSentinel: MapValueSlot<Buffer>;
 
   constructor(rootDb: Database, mapName: string) {
     this.name = mapName;
@@ -23,8 +23,8 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     // 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])];
+    this.startSentinel = ['map', this.name, 'slot', Buffer.from([])];
+    this.endSentinel = ['map', this.name, 'slot', Buffer.from([255])];
   }
 
   close(): Promise<void> {
@@ -32,7 +32,7 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
   }
 
   get(key: K): V | undefined {
-    return this.db.get(this.#slot(key))?.[1];
+    return this.db.get(this.slot(key))?.[1];
   }
 
   getAsync(key: K): Promise<V | undefined> {
@@ -40,7 +40,7 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
   }
 
   *getValues(key: K): IterableIterator<V> {
-    const values = this.db.getValues(this.#slot(key));
+    const values = this.db.getValues(this.slot(key));
     for (const value of values) {
       yield value?.[1];
     }
@@ -53,7 +53,7 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
   }
 
   has(key: K): boolean {
-    return this.db.doesExist(this.#slot(key));
+    return this.db.doesExist(this.slot(key));
   }
 
   hasAsync(key: K): Promise<boolean> {
@@ -61,30 +61,30 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
   }
 
   async set(key: K, val: V): Promise<void> {
-    await this.db.put(this.#slot(key), [key, val]);
+    await this.db.put(this.slot(key), [key, val]);
   }
 
   swap(key: K, fn: (val: V | undefined) => V): Promise<void> {
     return this.db.childTransaction(() => {
-      const slot = this.#slot(key);
+      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);
+    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));
+    await this.db.remove(this.slot(key));
   }
 
   async deleteValue(key: K, val: V): Promise<void> {
-    await this.db.remove(this.#slot(key), [key, val]);
+    await this.db.remove(this.slot(key), [key, val]);
   }
 
   *entries(range: Range<K> = {}): IterableIterator<[K, V]> {
@@ -93,19 +93,19 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     // in that case, we need to swap the start and end sentinels
     const start = reverse
       ? range.end
-        ? this.#slot(range.end)
-        : this.#endSentinel
+        ? 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
+        ? this.slot(range.start)
+        : this.startSentinel
       : range.end
-      ? this.#slot(range.end)
-      : this.#endSentinel;
+      ? this.slot(range.end)
+      : this.endSentinel;
 
     const lmdbRange: RangeOptions = {
       start,
@@ -153,7 +153,82 @@ export class LmdbAztecMap<K extends Key, V> implements AztecMultiMap<K, V>, Azte
     }
   }
 
-  #slot(key: K): MapValueSlot<K> {
+  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);
+    }
+  }
+}
+
+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/store.ts

@@ -1,20 +1,28 @@
+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 { dirname, join } from 'path';
+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 AztecMultiMap } from '../interfaces/map.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';
 
@@ -29,7 +37,7 @@ 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
@@ -64,13 +72,12 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
     ephemeral: boolean = false,
     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);
   }
 
   /**
@@ -78,10 +85,9 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
    * @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 fs.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 });
@@ -119,6 +125,23 @@ export class AztecLmdbStore implements AztecKVStore, AztecAsyncKVStore {
   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.

package/src/stores/l2_tips_store.ts

@@ -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);