feat: persist pxe state (#3628)

This PR adds a persistent database to PXE. I've created a new package in
the monorepo called `@aztec/kv-store` which exports a set of general
data structures that can be used by components to store state in
consistent manner. The only implementation right now is with LMDB (both
persisted on disk and in-memory/temporary file).

This higher level abstraction allowed me to easily add storage to ~PXE's
note processors, its synchronizer and~ the key store too.

Fix #3364 

Changes to the synchronizer and note processors have been taken out and
will be merged in a separate PR (as part of #3570) because it will
require some changes to the benchmarking code.

Synch changes here #3673

yarn-project/end-to-end/src/e2e_p2p_network.test.ts

@@ -150,7 +150,7 @@ describe('e2e_p2p_network', () => {
     numTxs: number,
   ): Promise<NodeContext> => {
     const rpcConfig = getRpcConfig();
-    const pxeService = await createPXEService(node, rpcConfig, {}, true);
+    const pxeService = await createPXEService(node, rpcConfig, true);
 
     const keyPair = ConstantKeyPair.random(new Grumpkin());
     const completeAddress = CompleteAddress.fromPrivateKeyAndPartialAddress(keyPair.getPrivateKey(), Fr.random());

yarn-project/end-to-end/src/fixtures/utils.ts

@@ -151,7 +151,7 @@ export async function setupPXEService(
   logger: DebugLogger;
 }> {
   const pxeServiceConfig = getPXEServiceConfig();
-  const pxe = await createPXEService(aztecNode, pxeServiceConfig, {}, useLogSuffix);
+  const pxe = await createPXEService(aztecNode, pxeServiceConfig, useLogSuffix);
 
   const wallets = await createAccounts(pxe, numberOfAccounts);
 

yarn-project/foundation/src/serialize/buffer_reader.ts

@@ -31,8 +31,16 @@ export class BufferReader {
    * @param bufferOrReader - A Buffer or BufferReader to initialize the BufferReader.
    * @returns An instance of BufferReader.
    */
-  public static asReader(bufferOrReader: Buffer | BufferReader) {
-    return Buffer.isBuffer(bufferOrReader) ? new BufferReader(bufferOrReader) : bufferOrReader;
+  public static asReader(bufferOrReader: Uint8Array | Buffer | BufferReader): BufferReader {
+    if (bufferOrReader instanceof BufferReader) {
+      return bufferOrReader;
+    }
+
+    const buf = Buffer.isBuffer(bufferOrReader)
+      ? bufferOrReader
+      : Buffer.from(bufferOrReader.buffer, bufferOrReader.byteOffset, bufferOrReader.byteLength);
+
+    return new BufferReader(buf);
   }
 
   /**

yarn-project/key-store/package.json

@@ -32,6 +32,7 @@
   "dependencies": {
     "@aztec/circuits.js": "workspace:^",
     "@aztec/foundation": "workspace:^",
+    "@aztec/kv-store": "workspace:^",
     "@aztec/types": "workspace:^",
     "tslib": "^2.4.0"
   },

yarn-project/key-store/src/test_key_store.ts

@@ -1,5 +1,6 @@
-import { GrumpkinPrivateKey } from '@aztec/circuits.js';
+import { GrumpkinPrivateKey, GrumpkinScalar, Point } from '@aztec/circuits.js';
 import { Grumpkin } from '@aztec/circuits.js/barretenberg';
+import { AztecKVStore, AztecMap } from '@aztec/kv-store';
 import { KeyPair, KeyStore, PublicKey } from '@aztec/types';
 
 import { ConstantKeyPair } from './key_pair.js';
@@ -9,30 +10,27 @@ import { ConstantKeyPair } from './key_pair.js';
  * It should be utilized in testing scenarios where secure key management is not required, and ease-of-use is prioritized.
  */
 export class TestKeyStore implements KeyStore {
-  private accounts: KeyPair[] = [];
-  constructor(private curve: Grumpkin) {}
+  #keys: AztecMap<string, Buffer>;
 
-  public addAccount(privKey: GrumpkinPrivateKey): PublicKey {
-    const keyPair = ConstantKeyPair.fromPrivateKey(this.curve, privKey);
-
-    // check if private key has already been used
-    const account = this.accounts.find(a => a.getPublicKey().equals(keyPair.getPublicKey()));
-    if (account) {
-      return account.getPublicKey();
-    }
+  constructor(private curve: Grumpkin, database: AztecKVStore) {
+    this.#keys = database.createMap('key_store');
+  }
 
-    this.accounts.push(keyPair);
+  public async addAccount(privKey: GrumpkinPrivateKey): Promise<PublicKey> {
+    const keyPair = ConstantKeyPair.fromPrivateKey(this.curve, privKey);
+    await this.#keys.setIfNotExists(keyPair.getPublicKey().toString(), keyPair.getPrivateKey().toBuffer());
     return keyPair.getPublicKey();
   }
 
-  public createAccount(): Promise<PublicKey> {
+  public async createAccount(): Promise<PublicKey> {
     const keyPair = ConstantKeyPair.random(this.curve);
-    this.accounts.push(keyPair);
-    return Promise.resolve(keyPair.getPublicKey());
+    await this.#keys.set(keyPair.getPublicKey().toString(), keyPair.getPrivateKey().toBuffer());
+    return keyPair.getPublicKey();
   }
 
   public getAccounts(): Promise<PublicKey[]> {
-    return Promise.resolve(this.accounts.map(a => a.getPublicKey()));
+    const range = Array.from(this.#keys.keys());
+    return Promise.resolve(range.map(key => Point.fromString(key)));
   }
 
   public getAccountPrivateKey(pubKey: PublicKey): Promise<GrumpkinPrivateKey> {
@@ -48,13 +46,13 @@ export class TestKeyStore implements KeyStore {
    * @param pubKey - The public key of the account to retrieve.
    * @returns The KeyPair object associated with the provided key.
    */
-  private getAccount(pubKey: PublicKey) {
-    const account = this.accounts.find(a => a.getPublicKey().equals(pubKey));
-    if (!account) {
+  private getAccount(pubKey: PublicKey): KeyPair {
+    const privKey = this.#keys.get(pubKey.toString());
+    if (!privKey) {
       throw new Error(
         'Unknown account.\nSee docs for context: https://docs.aztec.network/dev_docs/contracts/common_errors#unknown-contract-error',
       );
     }
-    return account;
+    return ConstantKeyPair.fromPrivateKey(this.curve, GrumpkinScalar.fromBuffer(privKey));
   }
 }

yarn-project/key-store/tsconfig.json

@@ -6,6 +6,9 @@
     "tsBuildInfoFile": ".tsbuildinfo"
   },
   "references": [
+    {
+      "path": "../kv-store"
+    },
     {
       "path": "../circuits.js"
     },

yarn-project/kv-store/.eslintrc.cjs

@@ -0,0 +1 @@
+module.exports = require('@aztec/foundation/eslint');

yarn-project/kv-store/README.md

@@ -0,0 +1,10 @@
+# KV Store
+
+The Aztec KV store is an implementation of a durable key-value database with a pluggable backend. THe only supported backend right now is LMDB by using the [`lmdb-js` package](https://github.com/kriszyp/lmdb-js).
+
+This package exports a number of primitive data structures that can be used to build domain-specific databases in each node component (e.g. a PXE database or an Archiver database). The data structures supported:
+
+- singleton - holds a single value. Great for when a value needs to be stored but it's not a collection (e.g. the latest block header or the length of an array)
+- array - works like a normal in-memory JS array. It can't contain holes and it can be used as a stack (push-pop mechanics).
+- map - a hashmap where keys can be numbers or strings
+- multi-map - just like a map but each key holds multiple values. Can be used for indexing into other data structures

yarn-project/kv-store/package.json

@@ -0,0 +1,50 @@
+{
+  "name": "@aztec/kv-store",
+  "version": "0.1.0",
+  "type": "module",
+  "exports": "./dest/index.js",
+  "scripts": {
+    "build": "yarn clean && tsc -b",
+    "build:dev": "tsc -b --watch",
+    "clean": "rm -rf ./dest .tsbuildinfo",
+    "formatting": "run -T prettier --check ./src && run -T eslint ./src",
+    "formatting:fix": "run -T eslint --fix ./src && run -T prettier -w ./src",
+    "test": "NODE_NO_WARNINGS=1 node --experimental-vm-modules $(yarn bin jest) --passWithNoTests",
+    "start": "DEBUG='aztec:*' && node ./dest/bin/index.js"
+  },
+  "inherits": [
+    "../package.common.json"
+  ],
+  "jest": {
+    "preset": "ts-jest/presets/default-esm",
+    "moduleNameMapper": {
+      "^(\\.{1,2}/.*)\\.[cm]?js$": "$1"
+    },
+    "testRegex": "./src/.*\\.test\\.(js|mjs|ts)$",
+    "rootDir": "./src",
+    "workerThreads": true
+  },
+  "dependencies": {
+    "@aztec/foundation": "workspace:^",
+    "lmdb": "^2.9.1"
+  },
+  "devDependencies": {
+    "@jest/globals": "^29.5.0",
+    "@types/jest": "^29.5.0",
+    "@types/node": "^18.7.23",
+    "jest": "^29.5.0",
+    "jest-mock-extended": "^3.0.3",
+    "ts-jest": "^29.1.0",
+    "ts-node": "^10.9.1",
+    "typescript": "^5.0.4"
+  },
+  "files": [
+    "dest",
+    "src",
+    "!*.test.*"
+  ],
+  "types": "./dest/index.d.ts",
+  "engines": {
+    "node": ">=18"
+  }
+}

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

@@ -0,0 +1,5 @@
+export * from './interfaces/array.js';
+export * from './interfaces/map.js';
+export * from './interfaces/singleton.js';
+export * from './interfaces/store.js';
+export * from './lmdb/store.js';

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

@@ -0,0 +1,54 @@
+/**
+ * An array backed by a persistent store. Can not have any holes in it.
+ */
+export interface AztecArray<T> {
+  /**
+   * The size of the array
+   */
+  length: number;
+
+  /**
+   * Pushes values to the end of the array
+   * @param vals - The values to push to the end of the array
+   * @returns The new length of the array
+   */
+  push(...vals: T[]): Promise<number>;
+
+  /**
+   * Pops a value from the end of the array.
+   * @returns The value that was popped, or undefined if the array was empty
+   */
+  pop(): Promise<T | undefined>;
+
+  /**
+   * Gets the value at the given index. Index can be in the range [-length, length - 1).
+   * If the index is negative, it will be treated as an offset from the end of the array.
+   *
+   * @param index - The index to get the value from
+   * @returns The value at the given index or undefined if the index is out of bounds
+   */
+  at(index: number): T | undefined;
+
+  /**
+   * Updates the value at the given index. Index can be in the range [-length, length - 1).
+   * @param index - The index to set the value at
+   * @param val - The value to set
+   * @returns Whether the value was set
+   */
+  setAt(index: number, val: T): Promise<boolean>;
+
+  /**
+   * Iterates over the array with indexes.
+   */
+  entries(): IterableIterator<[number, T]>;
+
+  /**
+   * Iterates over the array.
+   */
+  values(): IterableIterator<T>;
+
+  /**
+   * Iterates over the array.
+   */
+  [Symbol.iterator](): IterableIterator<T>;
+}

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

@@ -0,0 +1,70 @@
+/**
+ * A map backed by a persistent store.
+ */
+export interface AztecMap<K extends string | number, V> {
+  /**
+   * Gets the value at the given key.
+   * @param key - The key to get the value from
+   */
+  get(key: K): V | undefined;
+
+  /**
+   * Checks if a key exists in the map.
+   * @param key - The key to check
+   * @returns True if the key exists, false otherwise
+   */
+  has(key: K): boolean;
+
+  /**
+   * Sets the value at the given key.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  set(key: K, val: V): Promise<boolean>;
+
+  /**
+   * Sets the value at the given key if it does not already exist.
+   * @param key - The key to set the value at
+   * @param val - The value to set
+   */
+  setIfNotExists(key: K, val: V): Promise<boolean>;
+
+  /**
+   * Deletes the value at the given key.
+   * @param key - The key to delete the value at
+   */
+  delete(key: K): Promise<boolean>;
+
+  /**
+   * Iterates over the map's key-value entries
+   */
+  entries(): IterableIterator<[K, V]>;
+
+  /**
+   * Iterates over the map's values
+   */
+  values(): IterableIterator<V>;
+
+  /**
+   * Iterates over the map's keys
+   */
+  keys(): IterableIterator<K>;
+}
+
+/**
+ * A map backed by a persistent store that can have multiple values for a single key.
+ */
+export interface AztecMultiMap<K extends string | number, V> extends AztecMap<K, V> {
+  /**
+   * Gets all the values at the given key.
+   * @param key - The key to get the values from
+   */
+  getValues(key: K): IterableIterator<V>;
+
+  /**
+   * Deletes a specific value at the given key.
+   * @param key - The key to delete the value at
+   * @param val - The value to delete
+   */
+  deleteValue(key: K, val: V): Promise<void>;
+}

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

@@ -0,0 +1,20 @@
+/**
+ * Represents a singleton value in the database.
+ */
+export interface AztecSingleton<T> {
+  /**
+   * Gets the value.
+   */
+  get(): T | undefined;
+
+  /**
+   * Sets the value.
+   * @param val - The new value
+   */
+  set(val: T): Promise<boolean>;
+
+  /**
+   * Deletes the value.
+   */
+  delete(): Promise<boolean>;
+}

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

@@ -0,0 +1,40 @@
+import { AztecArray } from './array.js';
+import { AztecMap, AztecMultiMap } from './map.js';
+import { AztecSingleton } from './singleton.js';
+
+/** A key-value store */
+export interface AztecKVStore {
+  /**
+   * Creates a new map.
+   * @param name - The name of the map
+   * @returns The map
+   */
+  createMap<K extends string | number, V>(name: string): AztecMap<K, V>;
+
+  /**
+   * Creates a new multi-map.
+   * @param name - The name of the multi-map
+   * @returns The multi-map
+   */
+  createMultiMap<K extends string | number, V>(name: string): AztecMultiMap<K, V>;
+
+  /**
+   * Creates a new array.
+   * @param name - The name of the array
+   * @returns The array
+   */
+  createArray<T>(name: string): AztecArray<T>;
+
+  /**
+   * Creates a new singleton.
+   * @param name - The name of the singleton
+   * @returns The singleton
+   */
+  createSingleton<T>(name: string): AztecSingleton<T>;
+
+  /**
+   * Starts a transaction. All calls to read/write data while in a transaction are queued and executed atomically.
+   * @param callback - The callback to execute in a transaction
+   */
+  transaction<T extends Exclude<any, Promise<any>>>(callback: () => T): Promise<T>;
+}