feat(p2p): attestation pool persistence (#10667)

Author: Maddiaa0

yarn-project/kv-store/src/interfaces/map.ts

@@ -62,6 +62,19 @@ export interface AztecMap<K extends Key, V> extends AztecBaseMap<K, V> {
    * @param range - The range of keys to iterate over
    */
   keys(range?: Range<K>): IterableIterator<K>;
+
+  /**
+   * Clears the map.
+   */
+  clear(): Promise<void>;
+}
+
+export interface AztecMapWithSize<K extends Key, V> extends AztecMap<K, V> {
+  /**
+   * Gets the size of the map.
+   * @returns The size of the map
+   */
+  size(): number;
 }
 
 /**
@@ -82,6 +95,14 @@ export interface AztecMultiMap<K extends Key, V> extends AztecMap<K, V> {
   deleteValue(key: K, val: V): Promise<void>;
 }
 
+export interface AztecMultiMapWithSize<K extends Key, V> extends AztecMultiMap<K, V> {
+  /**
+   * Gets the size of the map.
+   * @returns The size of the map
+   */
+  size(): number;
+}
+
 /**
  * A map backed by a persistent store.
  */

yarn-project/kv-store/src/interfaces/store.ts

@@ -1,7 +1,14 @@
 import { type AztecArray, type AztecAsyncArray } from './array.js';
 import { type Key } from './common.js';
 import { type AztecAsyncCounter, type AztecCounter } from './counter.js';
-import { type AztecAsyncMap, type AztecAsyncMultiMap, type AztecMap, type AztecMultiMap } from './map.js';
+import {
+  type AztecAsyncMap,
+  type AztecAsyncMultiMap,
+  type AztecMap,
+  type AztecMapWithSize,
+  type AztecMultiMap,
+  type AztecMultiMapWithSize,
+} from './map.js';
 import { type AztecAsyncSet, type AztecSet } from './set.js';
 import { type AztecAsyncSingleton, type AztecSingleton } from './singleton.js';
 
@@ -29,6 +36,20 @@ export interface AztecKVStore {
    */
   openMultiMap<K extends Key, V>(name: string): AztecMultiMap<K, V>;
 
+  /**
+   * Creates a new multi-map with size.
+   * @param name - The name of the multi-map
+   * @returns The multi-map
+   */
+  openMultiMapWithSize<K extends Key, V>(name: string): AztecMultiMapWithSize<K, V>;
+
+  /**
+   * Creates a new map with size.
+   * @param name - The name of the map
+   * @returns The map
+   */
+  openMapWithSize<K extends Key, V>(name: string): AztecMapWithSize<K, V>;
+
   /**
    * Creates a new array.
    * @param name - The name of the array

yarn-project/kv-store/src/lmdb/map.test.ts

@@ -1,3 +1,6 @@
+import { expect } from 'chai';
+
+import { type AztecMapWithSize, type AztecMultiMapWithSize } from '../interfaces/map.js';
 import { describeAztecMap } from '../interfaces/map_test_suite.js';
 import { openTmpStore } from './index.js';
 
@@ -6,3 +9,53 @@ describe('LMDBMap', () => {
 
   describeAztecMap('Async AztecMap', () => Promise.resolve(openTmpStore(true)), true);
 });
+
+describe('AztecMultiMapWithSize', () => {
+  let map: AztecMultiMapWithSize<string, string>;
+  let map2: AztecMultiMapWithSize<string, string>;
+
+  beforeEach(() => {
+    const store = openTmpStore(true);
+    map = store.openMultiMapWithSize('test');
+    map2 = store.openMultiMapWithSize('test2');
+  });
+
+  it('should be able to delete values', async () => {
+    await map.set('foo', 'bar');
+    await map.set('foo', 'baz');
+
+    await map2.set('foo', 'bar');
+    await map2.set('foo', 'baz');
+
+    expect(map.size()).to.equal(2);
+    expect(map2.size()).to.equal(2);
+
+    await map.deleteValue('foo', 'bar');
+
+    expect(map.size()).to.equal(1);
+    expect(map.get('foo')).to.equal('baz');
+
+    expect(map2.size()).to.equal(2);
+  });
+});
+
+describe('AztecMapWithSize', () => {
+  let map: AztecMapWithSize<string, string>;
+
+  beforeEach(() => {
+    const store = openTmpStore(true);
+    map = store.openMapWithSize('test');
+  });
+
+  it('should be able to delete values', async () => {
+    await map.set('foo', 'bar');
+    await map.set('fizz', 'buzz');
+
+    expect(map.size()).to.equal(2);
+
+    await map.delete('foo');
+
+    expect(map.size()).to.equal(1);
+    expect(map.get('fizz')).to.equal('buzz');
+  });
+});

yarn-project/kv-store/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,24 +23,24 @@ 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> {
     return this.db.close();
   }
 
   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> {
     return Promise.resolve(this.get(key));
   }
 
   *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,38 +53,38 @@ 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> {
     return Promise.resolve(this.has(key));
   }
 
   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;
+  }
 }

yarn-project/kv-store/src/lmdb/store.ts

@@ -9,13 +9,20 @@ 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';
 
@@ -118,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.