refactor(pxe): extract a Rollbackable interface for reorg rollbacks (backport of aztec-labs-eng/aztec-node#40)

yarn-project/pxe/src/block_synchronizer/block_synchronizer.ts

@@ -10,9 +10,7 @@ import type { BlockSynchronizerConfig } from '../config/index.js';
 import type { ContractSyncService } from '../contract/contract_sync_service.js';
 import type { CachingAztecNode } from '../node/caching_aztec_node.js';
 import type { AnchorBlockStore } from '../storage/anchor_block_store/index.js';
-import type { FactStore } from '../storage/fact_store/fact_store.js';
-import type { NoteStore } from '../storage/note_store/index.js';
-import type { PrivateEventStore } from '../storage/private_event_store/private_event_store.js';
+import type { Rollbackable } from '../storage/rollbackable.js';
 import { blockStreamSourceFromAztecNode } from './block_stream_source.js';
 
 /**
@@ -30,9 +28,7 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
     private readonly node: CachingAztecNode,
     private readonly store: AztecAsyncKVStore,
     private readonly anchorBlockStore: AnchorBlockStore,
-    private readonly noteStore: NoteStore,
-    private readonly privateEventStore: PrivateEventStore,
-    private readonly factStore: FactStore,
+    private readonly rollbackables: Rollbackable[],
     private readonly l2TipsStore: L2TipsKVStore,
     private readonly contractSyncService: ContractSyncService,
     private readonly config: Partial<BlockSynchronizerConfig> = {},
@@ -138,9 +134,9 @@ export class BlockSynchronizer implements L2BlockStreamEventHandler {
 
         // Operations are wrapped in a single transaction to ensure atomicity.
         await this.store.transactionAsync(async () => {
-          await this.noteStore.rollback(event.block.number);
-          await this.privateEventStore.rollback(event.block.number);
-          await this.factStore.rollback(event.block.number);
+          for (const rollbackable of this.rollbackables) {
+            await rollbackable.rollbackToBlock(event.block.number);
+          }
           await this.updateAnchorBlockHeader(newAnchorBlockHeader);
         });
         break;

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

@@ -43,7 +43,7 @@ export interface ContractSyncConfig {
   /**
    * Whether PXE speculatively syncs contracts it predicts will follow the one requested, running them concurrently
    * with it instead of waiting for execution to reach them. When enabled, repeated flows sync faster, but a wrong
-   * prediction spends unnecessary node requests syncing contracts the job never uses.
+   * prediction spends unnecessary node requests syncing contracts the operation never uses.
    *
    * Experimental, off by default.
    */

yarn-project/pxe/src/contract/contract_call_graph.test.ts

@@ -1,6 +1,7 @@
 import { FunctionSelector } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
+import type { ChangeSetId } from '../storage/staged_write_coordinator.js';
 import {
   ContractCallGraph,
   type ContractFunction,
@@ -25,8 +26,8 @@ describe('ContractCallGraph', () => {
     expect(calleesOf(accountEntrypoint)).toEqual([]);
   });
 
-  it('does not predict a callee until enough committed jobs observe the call', () => {
-    runJobs({
+  it('does not predict a callee until enough committed change sets observe the call', () => {
+    runChangeSets({
       count: PREDICTION_THRESHOLD - 1,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
@@ -37,8 +38,8 @@ describe('ContractCallGraph', () => {
     expect(calleesOf(accountEntrypoint)).toEqual([]);
   });
 
-  it('predicts a callee once enough committed jobs observe the call', () => {
-    runJobs({
+  it('predicts a callee once enough committed change sets observe the call', () => {
+    runChangeSets({
       count: PREDICTION_THRESHOLD,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
@@ -50,7 +51,7 @@ describe('ContractCallGraph', () => {
   });
 
   it('predicts only direct callees, not callees of callees', () => {
-    runJobs({
+    runChangeSets({
       count: PREDICTION_THRESHOLD,
       calls: [
         { caller: accountEntrypoint, callee: fpcFee },
@@ -63,57 +64,57 @@ describe('ContractCallGraph', () => {
   });
 
   it('keys calls per function, so a sibling function of the same contract predicts nothing', () => {
-    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+    runChangeSets({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
 
     expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
     expect(calleesOf(accountClaim)).toEqual([]);
   });
 
   it("predicts a function's callees even when its own callers rarely call it", () => {
-    runJob({
-      jobId: 'rare',
+    runChangeSet({
+      changeSetId: 'rare',
       calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
     });
-    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: fpcFee }] });
+    runChangeSets({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: fpcFee }] });
 
     expect(calleesOf(accountEntrypoint)).toEqual([]);
     expect(calleesOf(tokenTransfer)).toEqual(callKeys([fpcFee]));
   });
 
   it('ignores same-contract calls', () => {
-    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: tokenBalance }] });
+    runChangeSets({ count: PREDICTION_THRESHOLD, calls: [{ caller: tokenTransfer, callee: tokenBalance }] });
 
     expect(calleesOf(tokenTransfer)).toEqual([]);
   });
 
-  it('does not learn from discarded jobs', () => {
-    runJobs({ count: PREDICTION_THRESHOLD - 1, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
-    callGraph.recordCall({ jobId: 'discarded', caller: accountEntrypoint, callee: tokenTransfer });
-    callGraph.discardJob('discarded');
+  it('does not learn from discarded change sets', () => {
+    runChangeSets({ count: PREDICTION_THRESHOLD - 1, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+    callGraph.recordCall({ changeSetId: 'discarded', caller: accountEntrypoint, callee: tokenTransfer });
+    callGraph.discard('discarded');
 
     expect(calleesOf(accountEntrypoint)).toEqual([]);
   });
 
-  it('leaves confidence untouched by jobs in which the caller makes no calls', () => {
-    runJobs({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+  it('leaves confidence untouched by change sets in which the caller makes no calls', () => {
+    runChangeSets({ count: PREDICTION_THRESHOLD, calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
 
-    // The account calls no one in these jobs, so the confidence of the callees it did not call is unaffected.
-    for (const jobId of ['read1', 'read2']) {
-      callGraph.commitJob(jobId);
+    // The account calls no one in these change sets, so the confidence of the callees it did not call is unaffected.
+    for (const changeSetId of ['read1', 'read2']) {
+      callGraph.learn(changeSetId);
     }
 
     expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
   });
 
   it('keeps predicting a callee at full confidence through every miss it tolerates', () => {
-    runJobs({
+    runChangeSets({
       count: MAX_CONFIDENCE,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
         { caller: accountEntrypoint, callee: fpcFee },
       ],
     });
-    runJobs({
+    runChangeSets({
       count: MAX_CONFIDENCE - PREDICTION_THRESHOLD,
       calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
     });
@@ -122,14 +123,14 @@ describe('ContractCallGraph', () => {
   });
 
   it('caps confidence, so a heavily called callee stops being predicted one miss past that tolerance', () => {
-    runJobs({
+    runChangeSets({
       count: MAX_CONFIDENCE * 2,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
         { caller: accountEntrypoint, callee: fpcFee },
       ],
     });
-    runJobs({
+    runChangeSets({
       count: MAX_CONFIDENCE - PREDICTION_THRESHOLD + 1,
       calls: [{ caller: accountEntrypoint, callee: tokenTransfer }],
     });
@@ -138,7 +139,7 @@ describe('ContractCallGraph', () => {
   });
 
   it('drops a callee below the threshold on a miss and predicts it again after one hit', () => {
-    runJobs({
+    runChangeSets({
       count: PREDICTION_THRESHOLD,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
@@ -147,11 +148,11 @@ describe('ContractCallGraph', () => {
     });
     expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee]));
 
-    runJob({ jobId: 'miss', calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
+    runChangeSet({ changeSetId: 'miss', calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] });
     expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer]));
 
-    runJob({
-      jobId: 'refresh',
+    runChangeSet({
+      changeSetId: 'refresh',
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
         { caller: accountEntrypoint, callee: fpcFee },
@@ -163,7 +164,7 @@ describe('ContractCallGraph', () => {
 
   it('never records calls when disabled', () => {
     callGraph = new ContractCallGraph(false);
-    runJobs({
+    runChangeSets({
       count: PREDICTION_THRESHOLD,
       calls: [
         { caller: accountEntrypoint, callee: tokenTransfer },
@@ -174,19 +175,19 @@ describe('ContractCallGraph', () => {
     expect(calleesOf(accountEntrypoint)).toEqual([]);
   });
 
-  /** Runs `count` whole jobs, each observing the given direct calls. */
-  function runJobs({ count, calls }: { count: number; calls: Call[] }) {
+  /** Runs `count` whole change sets, each observing the given direct calls. */
+  function runChangeSets({ count, calls }: { count: number; calls: Call[] }) {
     for (let i = 0; i < count; i++) {
-      runJob({ jobId: `job${i}`, calls });
+      runChangeSet({ changeSetId: `change-set-${i}`, calls });
     }
   }
 
-  /** Runs a whole job: records each direct call and commits. */
-  function runJob({ jobId, calls }: { jobId: string; calls: Call[] }) {
+  /** Runs a whole change set: records each direct call and learns from it as committed. */
+  function runChangeSet({ changeSetId, calls }: { changeSetId: ChangeSetId; calls: Call[] }) {
     for (const { caller, callee } of calls) {
-      callGraph.recordCall({ jobId, caller, callee });
+      callGraph.recordCall({ changeSetId, caller, callee });
     }
-    callGraph.commitJob(jobId);
+    callGraph.learn(changeSetId);
   }
 
   /** Returns the direct callees predicted for the given function, as sorted `address:selector` strings. */
@@ -195,7 +196,7 @@ describe('ContractCallGraph', () => {
   }
 });
 
-/** A direct call observed by a job. */
+/** A direct call observed by a change set. */
 type Call = { caller: ContractFunction; callee: ContractFunction };
 
 function fn(contractIndex: number, functionIndex: number): ContractFunction {

yarn-project/pxe/src/contract/contract_call_graph.ts

@@ -1,49 +1,59 @@
 import { FunctionSelector } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
+import type { ChangeSetId } from '../storage/staged_write_coordinator.js';
+
 /** Confidence a call must reach to be predicted. */
 export const PREDICTION_THRESHOLD = 2;
 
-/** Cap on a call's confidence, so a function called by many jobs is still dropped within a few missed ones. */
+/** Cap on a call's confidence, so a function called by many operations is still dropped within a few missed ones. */
 export const MAX_CONFIDENCE = 5;
 
 /**
- * A call graph over contract functions - who calls whom - learned from the direct calls observed in past jobs, so
+ * A call graph over contract functions - who calls whom - learned from the direct calls observed in past operations, so
  * a function's predicted callees can sync their contracts before execution reaches them.
  *
- * A function's direct calls tend to repeat across jobs: constrained delivery calls the handshake registry, a transfer
- * may call an authwit, an AMM calls its tokens. The same function does not always make the same calls, though: they
- * can depend on context or storage state. A call must therefore repeat often enough to earn confidence before it is
- * predicted. Calls are keyed per function, not per contract, since different functions of a contract call different
- * contracts. See {@link commitJob} for how each call's confidence is learned from committed jobs.
+ * A function's direct calls tend to repeat across operations: constrained delivery calls the handshake registry, a
+ * transfer may call an authwit, an AMM calls its tokens. The same function does not always make the same calls, though:
+ * they can depend on context or storage state. A call must therefore repeat often enough to earn confidence before it
+ * is predicted. Calls are keyed per function, not per contract, since different functions of a contract call different
+ * contracts. See {@link learn} for how each call's confidence is learned from committed operations.
  *
  * Purely in-memory bookkeeping: the graph is lost when PXE is rebuilt (e.g. on restart).
  */
 export class ContractCallGraph {
-  // job -> caller function -> functions it called directly
-  private readonly activeJobs: Map<JobId, Map<CallKey, Set<CallKey>>> = new Map();
+  // change set -> caller function -> functions it called directly
+  private readonly activeChangeSets: Map<ChangeSetId, Map<CallKey, Set<CallKey>>> = new Map();
 
   // caller function -> function it calls directly -> confidence score
   private readonly callConfidence: Map<CallKey, Map<CallKey, number>> = new Map();
 
   constructor(private readonly enabled: boolean) {}
 
-  /** Records that `caller` directly called `callee` in the given job. */
-  recordCall({ jobId, caller, callee }: { jobId: JobId; caller: ContractFunction; callee: ContractFunction }): void {
+  /** Records that `caller` directly called `callee` in the given change set. */
+  recordCall({
+    changeSetId,
+    caller,
+    callee,
+  }: {
+    changeSetId: ChangeSetId;
+    caller: ContractFunction;
+    callee: ContractFunction;
+  }): void {
     // Same-contract calls are ignored: our goal is to warm a callee's contract ahead of use, and the target of such
     // a call is already warm.
     if (!this.enabled || caller.address.equals(callee.address)) {
       return;
     }
-    let callsInJob = this.activeJobs.get(jobId);
-    if (!callsInJob) {
-      callsInJob = new Map();
-      this.activeJobs.set(jobId, callsInJob);
+    let callsInChangeSet = this.activeChangeSets.get(changeSetId);
+    if (!callsInChangeSet) {
+      callsInChangeSet = new Map();
+      this.activeChangeSets.set(changeSetId, callsInChangeSet);
     }
-    let callees = callsInJob.get(toCallKey(caller));
+    let callees = callsInChangeSet.get(toCallKey(caller));
     if (!callees) {
       callees = new Set();
-      callsInJob.set(toCallKey(caller), callees);
+      callsInChangeSet.set(toCallKey(caller), callees);
     }
     callees.add(toCallKey(callee));
   }
@@ -57,17 +67,17 @@ export class ContractCallGraph {
   }
 
   /**
-   * Commits the job so the calls it observed are recorded and learned from. A function that called nothing keeps its
-   * callees untouched, so read-only uses (e.g. reading notes or events) erode nothing.
+   * Learns from a committed change set: the calls it observed update each caller's confidence. A function that called
+   * nothing keeps its callees untouched, so read-only uses (e.g. reading notes or events) erode nothing.
    */
-  commitJob(jobId: JobId): void {
-    const callsInJob = this.activeJobs.get(jobId);
-    this.activeJobs.delete(jobId);
-    if (!callsInJob) {
+  learn(changeSetId: ChangeSetId): void {
+    const callsInChangeSet = this.activeChangeSets.get(changeSetId);
+    this.activeChangeSets.delete(changeSetId);
+    if (!callsInChangeSet) {
       return;
     }
 
-    for (const [caller, observed] of callsInJob) {
+    for (const [caller, observed] of callsInChangeSet) {
       const callees = this.callConfidence.get(caller) ?? new Map<CallKey, number>();
       for (const [callee, confidence] of callees) {
         const delta = observed.has(callee) ? 1 : -1;
@@ -84,9 +94,9 @@ export class ContractCallGraph {
     }
   }
 
-  /** Drops a discarded job without learning. */
-  discardJob(jobId: JobId): void {
-    this.activeJobs.delete(jobId);
+  /** Drops a discarded change set without learning. */
+  discard(changeSetId: ChangeSetId): void {
+    this.activeChangeSets.delete(changeSetId);
   }
 }
 
@@ -98,8 +108,6 @@ export type ContractFunction = {
   selector: FunctionSelector;
 };
 
-type JobId = string;
-
 /** A {@link ContractFunction} flattened to a `contractAddress:selector` string, so maps can key on it. */
 export type CallKey = `0x${string}:${string}`;