feat: PXE adds note processors for stored accounts (AztecProtocol#3673)

This PR makes the PXE add note processors for each stored key in the
keystore when it is started. This way any previously stored accounts
start syncing with the chain immediately.

Author: alexghr

yarn-project/archiver/src/archiver/archiver.ts

@@ -47,11 +47,6 @@ export class Archiver implements L2BlockSource, L2LogsSource, ContractDataSource
    */
   private runningPromise?: RunningPromise;
 
-  /**
-   * Next L1 block number to fetch `L2BlockProcessed` logs from (i.e. `fromBlock` in eth_getLogs).
-   */
-  private nextL2BlockFromL1Block = 0n;
-
   /**
    * Use this to track logged block in order to avoid repeating the same message.
    */
@@ -220,11 +215,21 @@ export class Archiver implements L2BlockSource, L2LogsSource, ContractDataSource
       this.publicClient,
       this.rollupAddress,
       blockUntilSynced,
-      this.nextL2BlockFromL1Block,
+      lastProcessedL1BlockNumber + 1n,
       currentL1BlockNumber,
       nextExpectedL2BlockNum,
     );
 
+    if (retrievedBlocks.retrievedData.length === 0) {
+      return;
+    } else {
+      this.log(
+        `Retrieved ${retrievedBlocks.retrievedData.length} new L2 blocks between L1 blocks ${
+          lastProcessedL1BlockNumber + 1n
+        } and ${currentL1BlockNumber}.`,
+      );
+    }
+
     // create the block number -> block hash mapping to ensure we retrieve the appropriate events
     const blockHashMapping: { [key: number]: Buffer | undefined } = {};
     retrievedBlocks.retrievedData.forEach((block: L2Block) => {
@@ -234,13 +239,10 @@ export class Archiver implements L2BlockSource, L2LogsSource, ContractDataSource
       this.publicClient,
       this.contractDeploymentEmitterAddress,
       blockUntilSynced,
-      this.nextL2BlockFromL1Block,
+      lastProcessedL1BlockNumber + 1n,
       currentL1BlockNumber,
       blockHashMapping,
     );
-    if (retrievedBlocks.retrievedData.length === 0) {
-      return;
-    }
 
     this.log(`Retrieved ${retrievedBlocks.retrievedData.length} block(s) from chain`);
 
@@ -280,9 +282,6 @@ export class Archiver implements L2BlockSource, L2LogsSource, ContractDataSource
         );
       }),
     );
-
-    // set the L1 block for the next search
-    this.nextL2BlockFromL1Block = retrievedBlocks.nextEthBlockNumber;
   }
 
   /**

yarn-project/pxe/src/config/index.ts

@@ -10,7 +10,7 @@ import { fileURLToPath } from 'url';
 export interface PXEServiceConfig {
   /** The interval to wait between polling for new blocks. */
   l2BlockPollingIntervalMS: number;
-  /** L2 block to start scanning from */
+  /** L2 block to start scanning from for new accounts */
   l2StartingBlock: number;
 
   /** Where to store PXE data. If not set will store in memory */

yarn-project/pxe/src/database/kv_pxe_database.ts

@@ -1,24 +1,26 @@
 import { AztecAddress, BlockHeader, CompleteAddress } from '@aztec/circuits.js';
-import { Fr } from '@aztec/foundation/fields';
+import { Fr, Point } from '@aztec/foundation/fields';
 import { AztecArray, AztecKVStore, AztecMap, AztecMultiMap, AztecSingleton } from '@aztec/kv-store';
 import { ContractDao, MerkleTreeId, NoteFilter, PublicKey } from '@aztec/types';
 
 import { NoteDao } from './note_dao.js';
 import { PxeDatabase } from './pxe_database.js';
 
 /** Serialized structure of a block header */
-type SerializedBlockHeader = {
+type SynchronizedBlock = {
   /** The tree roots when the block was created */
   roots: Record<MerkleTreeId, string>;
   /** The hash of the global variables */
   globalVariablesHash: string;
+  /** The block number */
+  blockNumber: number;
 };
 
 /**
  * A PXE database backed by LMDB.
  */
 export class KVPxeDatabase implements PxeDatabase {
-  #blockHeader: AztecSingleton<SerializedBlockHeader>;
+  #synchronizedBlock: AztecSingleton<SynchronizedBlock>;
   #addresses: AztecArray<Buffer>;
   #addressIndex: AztecMap<string, number>;
   #authWitnesses: AztecMap<string, Buffer[]>;
@@ -30,6 +32,7 @@ export class KVPxeDatabase implements PxeDatabase {
   #notesByStorageSlot: AztecMultiMap<string, number>;
   #notesByTxHash: AztecMultiMap<string, number>;
   #notesByOwner: AztecMultiMap<string, number>;
+  #syncedBlockPerPublicKey: AztecMap<string, number>;
   #db: AztecKVStore;
 
   constructor(db: AztecKVStore) {
@@ -40,9 +43,11 @@ export class KVPxeDatabase implements PxeDatabase {
 
     this.#authWitnesses = db.createMap('auth_witnesses');
     this.#capsules = db.createArray('capsules');
-    this.#blockHeader = db.createSingleton('block_header');
     this.#contracts = db.createMap('contracts');
 
+    this.#synchronizedBlock = db.createSingleton('block_header');
+    this.#syncedBlockPerPublicKey = db.createMap('synced_block_per_public_key');
+
     this.#notes = db.createArray('notes');
     this.#nullifiedNotes = db.createMap('nullified_notes');
 
@@ -173,7 +178,7 @@ export class KVPxeDatabase implements PxeDatabase {
   }
 
   getTreeRoots(): Record<MerkleTreeId, Fr> {
-    const roots = this.#blockHeader.get()?.roots;
+    const roots = this.#synchronizedBlock.get()?.roots;
     if (!roots) {
       throw new Error(`Tree roots not set`);
     }
@@ -188,8 +193,9 @@ export class KVPxeDatabase implements PxeDatabase {
     };
   }
 
-  async setBlockHeader(blockHeader: BlockHeader): Promise<void> {
-    await this.#blockHeader.set({
+  async setBlockData(blockNumber: number, blockHeader: BlockHeader): Promise<void> {
+    await this.#synchronizedBlock.set({
+      blockNumber,
       globalVariablesHash: blockHeader.globalVariablesHash.toString(),
       roots: {
         [MerkleTreeId.NOTE_HASH_TREE]: blockHeader.noteHashTreeRoot.toString(),
@@ -202,8 +208,12 @@ export class KVPxeDatabase implements PxeDatabase {
     });
   }
 
+  getBlockNumber(): number | undefined {
+    return this.#synchronizedBlock.get()?.blockNumber;
+  }
+
   getBlockHeader(): BlockHeader {
-    const value = this.#blockHeader.get();
+    const value = this.#synchronizedBlock.get();
     if (!value) {
       throw new Error(`Block header not set`);
     }
@@ -261,6 +271,14 @@ export class KVPxeDatabase implements PxeDatabase {
     return Promise.resolve(Array.from(this.#addresses).map(v => CompleteAddress.fromBuffer(v)));
   }
 
+  getSynchedBlockNumberForPublicKey(publicKey: Point): number | undefined {
+    return this.#syncedBlockPerPublicKey.get(publicKey.toString());
+  }
+
+  setSynchedBlockNumberForPublicKey(publicKey: Point, blockNumber: number): Promise<boolean> {
+    return this.#syncedBlockPerPublicKey.set(publicKey.toString(), blockNumber);
+  }
+
   estimateSize(): number {
     const notesSize = Array.from(this.#getAllNonNullifiedNotes()).reduce((sum, note) => sum + note.getSize(), 0);
     const authWitsSize = Array.from(this.#authWitnesses.values()).reduce(

yarn-project/pxe/src/database/memory_db.ts

@@ -1,6 +1,6 @@
 import { BlockHeader, CompleteAddress, PublicKey } from '@aztec/circuits.js';
 import { AztecAddress } from '@aztec/foundation/aztec-address';
-import { Fr } from '@aztec/foundation/fields';
+import { Fr, Point } from '@aztec/foundation/fields';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { MerkleTreeId, NoteFilter } from '@aztec/types';
 
@@ -18,8 +18,10 @@ export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
   private notesTable: NoteDao[] = [];
   private treeRoots: Record<MerkleTreeId, Fr> | undefined;
   private globalVariablesHash: Fr | undefined;
+  private blockNumber: number | undefined;
   private addresses: CompleteAddress[] = [];
   private authWitnesses: Record<string, Fr[]> = {};
+  private syncedBlockPerPublicKey = new Map<string, number>();
   // A capsule is a "blob" of data that is passed to the contract through an oracle.
   // We are using a stack to keep track of the capsules that are passed to the contract.
   private capsuleStack: Fr[][] = [];
@@ -134,8 +136,9 @@ export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
     );
   }
 
-  public setBlockHeader(blockHeader: BlockHeader): Promise<void> {
+  public setBlockData(blockNumber: number, blockHeader: BlockHeader): Promise<void> {
     this.globalVariablesHash = blockHeader.globalVariablesHash;
+    this.blockNumber = blockNumber;
     this.setTreeRoots({
       [MerkleTreeId.NOTE_HASH_TREE]: blockHeader.noteHashTreeRoot,
       [MerkleTreeId.NULLIFIER_TREE]: blockHeader.nullifierTreeRoot,
@@ -148,6 +151,10 @@ export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
     return Promise.resolve();
   }
 
+  public getBlockNumber(): number | undefined {
+    return this.blockNumber;
+  }
+
   public addCompleteAddress(completeAddress: CompleteAddress): Promise<boolean> {
     const accountIndex = this.addresses.findIndex(r => r.address.equals(completeAddress.address));
     if (accountIndex !== -1) {
@@ -174,6 +181,15 @@ export class MemoryDB extends MemoryContractDatabase implements PxeDatabase {
     return Promise.resolve(this.addresses);
   }
 
+  getSynchedBlockNumberForPublicKey(publicKey: Point): number | undefined {
+    return this.syncedBlockPerPublicKey.get(publicKey.toString());
+  }
+
+  setSynchedBlockNumberForPublicKey(publicKey: Point, blockNumber: number): Promise<boolean> {
+    this.syncedBlockPerPublicKey.set(publicKey.toString(), blockNumber);
+    return Promise.resolve(true);
+  }
+
   public estimateSize() {
     const notesSize = this.notesTable.reduce((sum, note) => sum + note.getSize(), 0);
     const treeRootsSize = this.treeRoots ? Object.entries(this.treeRoots).length * Fr.SIZE_IN_BYTES : 0;

yarn-project/pxe/src/database/pxe_database.ts

@@ -79,13 +79,22 @@ export interface PxeDatabase extends ContractDatabase {
    */
   getTreeRoots(): Record<MerkleTreeId, Fr>;
 
+  /**
+   * Gets the most recently processed block number.
+   * @returns The most recently processed block number or undefined if never synched.
+   */
+  getBlockNumber(): number | undefined;
+
   /**
    * Retrieve the stored Block Header from the database.
    * The function returns a Promise that resolves to the Block Header.
    * This data is required to reproduce block attestations.
    * Throws an error if the block header is not available within the database.
    *
    * note: this data is a combination of the tree roots and the global variables hash.
+   *
+   * @returns The Block Header.
+   * @throws If no block have been processed yet.
    */
   getBlockHeader(): BlockHeader;
 
@@ -94,10 +103,11 @@ export interface PxeDatabase extends ContractDatabase {
    * This function updates the 'global variables hash' and `tree roots` property of the instance
    * Note that this will overwrite any existing hash or roots in the database.
    *
+   * @param blockNumber - The block number of the most recent block
    * @param blockHeader - An object containing the most recent block header.
    * @returns A Promise that resolves when the hash has been successfully updated in the database.
    */
-  setBlockHeader(blockHeader: BlockHeader): Promise<void>;
+  setBlockData(blockNumber: number, blockHeader: BlockHeader): Promise<void>;
 
   /**
    * Adds complete address to the database.
@@ -121,6 +131,19 @@ export interface PxeDatabase extends ContractDatabase {
    */
   getCompleteAddresses(): Promise<CompleteAddress[]>;
 
+  /**
+   * Updates up to which block number we have processed notes for a given public key.
+   * @param publicKey - The public key to set the synched block number for.
+   * @param blockNumber - The block number to set.
+   */
+  setSynchedBlockNumberForPublicKey(publicKey: PublicKey, blockNumber: number): Promise<boolean>;
+
+  /**
+   * Get the synched block number for a given public key.
+   * @param publicKey - The public key to get the synched block number for.
+   */
+  getSynchedBlockNumberForPublicKey(publicKey: PublicKey): number | undefined;
+
   /**
    * Returns the estimated size in bytes of this db.
    * @returns The estimated size in bytes of this db.

yarn-project/pxe/src/database/pxe_database_test_suite.ts

@@ -1,6 +1,6 @@
 import { AztecAddress, BlockHeader, CompleteAddress } from '@aztec/circuits.js';
 import { Fr, Point } from '@aztec/foundation/fields';
-import { MerkleTreeId, NoteFilter, randomTxHash } from '@aztec/types';
+import { INITIAL_L2_BLOCK_NUM, MerkleTreeId, NoteFilter, randomTxHash } from '@aztec/types';
 
 import { NoteDao } from './note_dao.js';
 import { randomNoteDao } from './note_dao.test.js';
@@ -155,13 +155,13 @@ export function describePxeDatabase(getDatabase: () => PxeDatabase) {
         const blockHeader = BlockHeader.random();
         blockHeader.privateKernelVkTreeRoot = Fr.zero();
 
-        await database.setBlockHeader(blockHeader);
+        await database.setBlockData(INITIAL_L2_BLOCK_NUM, blockHeader);
         expect(database.getBlockHeader()).toEqual(blockHeader);
       });
 
       it('retrieves the merkle tree roots from the block', async () => {
         const blockHeader = BlockHeader.random();
-        await database.setBlockHeader(blockHeader);
+        await database.setBlockData(INITIAL_L2_BLOCK_NUM, blockHeader);
         expect(database.getTreeRoots()).toEqual({
           [MerkleTreeId.NOTE_HASH_TREE]: blockHeader.noteHashTreeRoot,
           [MerkleTreeId.NULLIFIER_TREE]: blockHeader.nullifierTreeRoot,

yarn-project/pxe/src/note_processor/note_processor.test.ts

@@ -223,4 +223,26 @@ describe('Note Processor', () => {
     addedNoteDaos.forEach(info => nonceSet.add(info.nonce.value));
     expect(nonceSet.size).toBe(notes.length);
   });
+
+  it('advances the block number', async () => {
+    const { blockContexts, encryptedLogsArr } = mockData([[2]]);
+    await noteProcessor.process(blockContexts, encryptedLogsArr);
+    expect(noteProcessor.status.syncedToBlock).toEqual(blockContexts.at(-1)?.block.number);
+  });
+
+  it('should restore the last block number processed and ignore the starting block', async () => {
+    const { blockContexts, encryptedLogsArr } = mockData([[2]]);
+    await noteProcessor.process(blockContexts, encryptedLogsArr);
+
+    const newNoteProcessor = new NoteProcessor(
+      owner.getPublicKey(),
+      keyStore,
+      database,
+      aztecNode,
+      INITIAL_L2_BLOCK_NUM,
+      simulator,
+    );
+
+    expect(newNoteProcessor.status).toEqual(noteProcessor.status);
+  });
 });