diff --git a/yarn-project/foundation/src/config/env_var.ts b/yarn-project/foundation/src/config/env_var.ts index 849b28b777d6..0209a28b4aaa 100644 --- a/yarn-project/foundation/src/config/env_var.ts +++ b/yarn-project/foundation/src/config/env_var.ts @@ -214,6 +214,7 @@ export type EnvVar = | 'PROVER_TEST_DELAY_TYPE' | 'PROVER_TEST_VERIFICATION_DELAY_MS' | 'PXE_AUTO_SYNC' + | 'PXE_CONCURRENT_CONTRACT_SYNC_ENABLED' | 'PXE_L2_BLOCK_BATCH_SIZE' | 'PXE_PROVER_ENABLED' | 'PXE_SYNC_CHAIN_TIP' diff --git a/yarn-project/pxe/src/config/index.ts b/yarn-project/pxe/src/config/index.ts index 9f81c6c098e7..f5129a76e78d 100644 --- a/yarn-project/pxe/src/config/index.ts +++ b/yarn-project/pxe/src/config/index.ts @@ -36,7 +36,25 @@ export interface BlockSynchronizerConfig { autoSync: boolean; } -export type PXEConfig = KernelProverConfig & DataStoreConfig & ChainConfig & BlockSynchronizerConfig; +/** + * Configuration settings for the contract sync service. + */ +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. + * + * Experimental, off by default. + */ + concurrentContractSyncEnabled: boolean; +} + +export type PXEConfig = KernelProverConfig & + DataStoreConfig & + ChainConfig & + BlockSynchronizerConfig & + ContractSyncConfig; export type CliPXEOptions = { /** Custom Aztec Node URL to connect to */ @@ -74,6 +92,12 @@ export const pxeConfigMappings: ConfigMappingsType = { 'Whether PXE syncs with the node automatically before each operation. Disable to let the caller (e.g. a wallet) drive syncs explicitly via pxe.sync().', ...booleanConfigHelper(true), }, + concurrentContractSyncEnabled: { + env: 'PXE_CONCURRENT_CONTRACT_SYNC_ENABLED', + description: + 'Whether PXE speculatively syncs contracts it predicts will follow the one requested, running them concurrently with it. Repeated flows sync faster, but a wrong prediction spends unnecessary node requests. Experimental, off by default.', + ...booleanConfigHelper(false), + }, }; /** diff --git a/yarn-project/pxe/src/contract/contract_call_graph.test.ts b/yarn-project/pxe/src/contract/contract_call_graph.test.ts new file mode 100644 index 000000000000..c3b41b2062f9 --- /dev/null +++ b/yarn-project/pxe/src/contract/contract_call_graph.test.ts @@ -0,0 +1,212 @@ +import { FunctionSelector } from '@aztec/stdlib/abi'; +import { AztecAddress } from '@aztec/stdlib/aztec-address'; + +import { + ContractCallGraph, + type ContractFunction, + MAX_CONFIDENCE, + PREDICTION_THRESHOLD, +} from './contract_call_graph.js'; + +describe('ContractCallGraph', () => { + let callGraph: ContractCallGraph; + + const accountEntrypoint = fn(1, 1); + const accountClaim = fn(1, 2); + const tokenTransfer = fn(2, 1); + const tokenBalance = fn(2, 2); + const fpcFee = fn(3, 1); + + beforeEach(() => { + callGraph = new ContractCallGraph(true); + }); + + it('returns nothing for a function it has never seen', () => { + expect(calleesOf(accountEntrypoint)).toEqual([]); + }); + + it('does not predict a callee until enough committed jobs observe the call', () => { + runJobs({ + count: PREDICTION_THRESHOLD - 1, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + + expect(calleesOf(accountEntrypoint)).toEqual([]); + }); + + it('predicts a callee once enough committed jobs observe the call', () => { + runJobs({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee])); + }); + + it('predicts only direct callees, not callees of callees', () => { + runJobs({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: accountEntrypoint, callee: fpcFee }, + { caller: fpcFee, callee: tokenTransfer }, + ], + }); + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([fpcFee])); + expect(calleesOf(fpcFee)).toEqual(callKeys([tokenTransfer])); + }); + + it('keys calls per function, so a sibling function of the same contract predicts nothing', () => { + runJobs({ 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', + calls: [{ caller: accountEntrypoint, callee: tokenTransfer }], + }); + runJobs({ 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 }] }); + + 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'); + + 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 }] }); + + // 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); + } + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer])); + }); + + it('keeps predicting a callee at full confidence through every miss it tolerates', () => { + runJobs({ + count: MAX_CONFIDENCE, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + runJobs({ + count: MAX_CONFIDENCE - PREDICTION_THRESHOLD, + calls: [{ caller: accountEntrypoint, callee: tokenTransfer }], + }); + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee])); + }); + + it('caps confidence, so a heavily called callee stops being predicted one miss past that tolerance', () => { + runJobs({ + count: MAX_CONFIDENCE * 2, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + runJobs({ + count: MAX_CONFIDENCE - PREDICTION_THRESHOLD + 1, + calls: [{ caller: accountEntrypoint, callee: tokenTransfer }], + }); + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer])); + }); + + it('drops a callee below the threshold on a miss and predicts it again after one hit', () => { + runJobs({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee])); + + runJob({ jobId: 'miss', calls: [{ caller: accountEntrypoint, callee: tokenTransfer }] }); + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer])); + + runJob({ + jobId: 'refresh', + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + + expect(calleesOf(accountEntrypoint)).toEqual(callKeys([tokenTransfer, fpcFee])); + }); + + it('never records calls when disabled', () => { + callGraph = new ContractCallGraph(false); + runJobs({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: accountEntrypoint, callee: tokenTransfer }, + { caller: accountEntrypoint, callee: fpcFee }, + ], + }); + + expect(calleesOf(accountEntrypoint)).toEqual([]); + }); + + /** Runs `count` whole jobs, each observing the given direct calls. */ + function runJobs({ count, calls }: { count: number; calls: Call[] }) { + for (let i = 0; i < count; i++) { + runJob({ jobId: `job${i}`, calls }); + } + } + + /** Runs a whole job: records each direct call and commits. */ + function runJob({ jobId, calls }: { jobId: string; calls: Call[] }) { + for (const { caller, callee } of calls) { + callGraph.recordCall({ jobId, caller, callee }); + } + callGraph.commitJob(jobId); + } + + /** Returns the direct callees predicted for the given function, as sorted `address:selector` strings. */ + function calleesOf(caller: ContractFunction): string[] { + return callKeys(callGraph.predictDirectCallees(caller)); + } +}); + +/** A direct call observed by a job. */ +type Call = { caller: ContractFunction; callee: ContractFunction }; + +function fn(contractIndex: number, functionIndex: number): ContractFunction { + return { address: makeAddress(contractIndex), selector: new FunctionSelector(0x1000 + functionIndex) }; +} + +function makeAddress(index: number): AztecAddress { + return AztecAddress.fromNumberUnsafe(0x1000 + index); +} + +/** Flattens functions to sorted `address:selector` strings, so sets of predictions can be compared. */ +function callKeys(functions: ContractFunction[]): string[] { + return functions.map(({ address, selector }) => `${address.toString()}:${selector.toString()}`).sort(); +} diff --git a/yarn-project/pxe/src/contract/contract_call_graph.ts b/yarn-project/pxe/src/contract/contract_call_graph.ts new file mode 100644 index 000000000000..b66ee96cd82e --- /dev/null +++ b/yarn-project/pxe/src/contract/contract_call_graph.ts @@ -0,0 +1,113 @@ +import { FunctionSelector } from '@aztec/stdlib/abi'; +import { AztecAddress } from '@aztec/stdlib/aztec-address'; + +/** 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. */ +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 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. + * + * 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>> = new Map(); + + // caller function -> function it calls directly -> confidence score + private readonly callConfidence: Map> = 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 { + // 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 callees = callsInJob.get(toCallKey(caller)); + if (!callees) { + callees = new Set(); + callsInJob.set(toCallKey(caller), callees); + } + callees.add(toCallKey(callee)); + } + + /** Predicts the functions `caller` will call directly. */ + predictDirectCallees(caller: ContractFunction): ContractFunction[] { + const callees = this.callConfidence.get(toCallKey(caller)) ?? new Map(); + return [...callees.entries()] + .filter(([, confidence]) => confidence >= PREDICTION_THRESHOLD) + .map(([callee]) => fromCallKey(callee)); + } + + /** + * 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. + */ + commitJob(jobId: JobId): void { + const callsInJob = this.activeJobs.get(jobId); + this.activeJobs.delete(jobId); + if (!callsInJob) { + return; + } + + for (const [caller, observed] of callsInJob) { + const callees = this.callConfidence.get(caller) ?? new Map(); + for (const [callee, confidence] of callees) { + const delta = observed.has(callee) ? 1 : -1; + const updated = Math.min(confidence + delta, MAX_CONFIDENCE); + if (updated === 0) { + callees.delete(callee); + } else { + callees.set(callee, updated); + } + } + // First-time callees enter at 1, below PREDICTION_THRESHOLD: a call must repeat before it is predicted. + [...observed].filter(callee => !callees.has(callee)).forEach(callee => callees.set(callee, 1)); + this.callConfidence.set(caller, callees); + } + } + + /** Drops a discarded job without learning. */ + discardJob(jobId: JobId): void { + this.activeJobs.delete(jobId); + } +} + +/** A specific function of a contract, as observed in a call. */ +export type ContractFunction = { + /** The address of the contract the function belongs to. */ + address: AztecAddress; + /** The selector of the function. */ + 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}`; + +export function toCallKey({ address, selector }: ContractFunction): CallKey { + return `${address.toString()}:${selector.toString()}`; +} + +function fromCallKey(key: CallKey): ContractFunction { + const [address, selector] = key.split(':'); + return { address: AztecAddress.fromStringUnsafe(address), selector: FunctionSelector.fromString(selector) }; +} diff --git a/yarn-project/pxe/src/contract/contract_sync_service.test.ts b/yarn-project/pxe/src/contract/contract_sync_service.test.ts index f93a21813269..bc56b4e99353 100644 --- a/yarn-project/pxe/src/contract/contract_sync_service.test.ts +++ b/yarn-project/pxe/src/contract/contract_sync_service.test.ts @@ -1,6 +1,8 @@ import { Fr } from '@aztec/foundation/curves/bn254'; import { createLogger } from '@aztec/foundation/log'; +import { promiseWithResolvers } from '@aztec/foundation/promise'; import { executeTimeout } from '@aztec/foundation/timer'; +import { TestContractArtifact } from '@aztec/noir-test-contracts.js/Test'; import { FunctionCall, FunctionSelector, FunctionType } from '@aztec/stdlib/abi'; import { AztecAddress } from '@aztec/stdlib/aztec-address'; import type { AztecNode } from '@aztec/stdlib/interfaces/client'; @@ -11,8 +13,9 @@ import { mock } from 'jest-mock-extended'; import type { ContractStore } from '../storage/contract_store/contract_store.js'; import type { NoteStore } from '../storage/note_store/note_store.js'; +import { type ContractFunction, PREDICTION_THRESHOLD } from './contract_call_graph.js'; import type { ContractClassService } from './contract_class_service.js'; -import { ContractSyncService, MAX_CONCURRENT_SCOPE_SYNCS } from './contract_sync_service.js'; +import { ContractSyncService, MAX_CONCURRENT_SCOPE_SYNCS, SYNC_STATE_SELECTOR } from './contract_sync_service.js'; describe('ContractSyncService', () => { let aztecNode: ReturnType>; @@ -63,78 +66,182 @@ describe('ContractSyncService', () => { contractClassService, noteStore, createLogger('test:contract-sync'), + { concurrentContractSyncEnabled: false }, ); }); describe('ensureContractSynced', () => { it('syncs a contract when not yet cached', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA]); }); it('re-syncs after wipe', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); service.wipe(); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeA]); }); it('skips scope-specific syncs after syncing with all scopes', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // [scopeA, scopeB] syncs each scope individually expectSyncedScopes([scopeA], [scopeB]); // After syncing all scopes, scope-specific calls should be skipped - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeB]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); }); it('only syncs unsynced scopes when requesting multiple', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // scopeA is already cached, so only scopeB is synced expectSyncedScopes([scopeA], [scopeB]); }); it('empty scopes array skips sync entirely', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, []); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [], + triggeredBy: undefined, + }); expectNoSync(); }); it('passes only unsynced scopes to the utility executor', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); }); it('concurrent calls for same contract+scope share one sync promise', async () => { - const p1 = service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - ]); - const p2 = service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - ]); + const p1 = service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + const p2 = service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); await Promise.all([p1, p2]); expectSyncedScopes([scopeA]); }); it('concurrent calls for different scopes trigger separate syncs', async () => { - const p1 = service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - ]); - const p2 = service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeB, - ]); + const p1 = service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + const p2 = service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeB], + triggeredBy: undefined, + }); await Promise.all([p1, p2]); expectSyncedScopes([scopeA], [scopeB]); }); @@ -155,13 +262,29 @@ describe('ContractSyncService', () => { utilityExecutor.mockImplementation(async (call, scopes) => { const nested = nestedByOuter.get(call.to.toString()); if (nested) { - await service.ensureContractSynced(nested, null, utilityExecutor, anchorBlockHeader, jobId, scopes); + await service.ensureContractSynced({ + contract: nested, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes, + triggeredBy: undefined, + }); } }); const syncAll = Promise.all( outerContracts.map(outer => - service.ensureContractSynced(outer, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]), + service.ensureContractSynced({ + contract: outer, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }), ), ); @@ -187,14 +310,15 @@ describe('ContractSyncService', () => { ); }); - const syncAll = service.ensureContractSynced( - contractAddress, - null, + const syncAll = service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, utilityExecutor, anchorBlockHeader, jobId, scopes, - ); + triggeredBy: undefined, + }); // The first wave saturates the limiter; the remaining scopes must queue rather than run. await tick(); @@ -214,11 +338,27 @@ describe('ContractSyncService', () => { it('re-syncs if first sync fails', async () => { utilityExecutor.mockRejectedValueOnce(new Error('sync failed')); await expect( - service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]), + service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }), ).rejects.toThrow('sync failed'); utilityExecutor.mockResolvedValueOnce(undefined); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); // the following checks that we attempted sync twice expectSyncedScopes([scopeA], [scopeA]); }); @@ -226,16 +366,40 @@ describe('ContractSyncService', () => { it('propagates sync errors to caller', async () => { utilityExecutor.mockRejectedValue(new Error('boom')); await expect( - service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]), + service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }), ).rejects.toThrow('boom'); }); }); describe('commit', () => { it('does not clear sync cache', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); await service.commit(jobId); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); // We check that the sync cache was not cleared by checking that the sync was triggered only once. expectSyncedScopes([scopeA]); }); @@ -243,9 +407,25 @@ describe('ContractSyncService', () => { describe('discardStaged', () => { it('clears sync cache', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); await service.discardStaged(jobId); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); // We check that the sync cache was cleared by checking that the sync was triggered twice. expectSyncedScopes([scopeA], [scopeA]); }); @@ -253,10 +433,15 @@ describe('ContractSyncService', () => { describe('multi-scope sync batching', () => { it('batches nullifier sync across all unsynced scopes', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expect(noteStore.getNotes).toHaveBeenCalledTimes(1); expect(noteStore.getNotes).toHaveBeenCalledWith( expect.objectContaining({ contractAddress, scopes: [scopeA, scopeB] }), @@ -265,7 +450,15 @@ describe('ContractSyncService', () => { }); it('only includes unsynced scopes in nullifier sync', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expect(noteStore.getNotes).toHaveBeenCalledTimes(1); expect(noteStore.getNotes).toHaveBeenCalledWith( expect.objectContaining({ contractAddress, scopes: [scopeA] }), @@ -273,10 +466,15 @@ describe('ContractSyncService', () => { ); noteStore.getNotes.mockClear(); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // scopeA is already cached, so nullifier sync only runs for scopeB expect(noteStore.getNotes).toHaveBeenCalledTimes(1); expect(noteStore.getNotes).toHaveBeenCalledWith( @@ -286,17 +484,27 @@ describe('ContractSyncService', () => { }); it('re-runs nullifier sync after scope invalidation', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); noteStore.getNotes.mockClear(); service.invalidateContractForScopes(contractAddress, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // Only scopeA was invalidated, so nullifier sync runs for just scopeA expect(noteStore.getNotes).toHaveBeenCalledTimes(1); expect(noteStore.getNotes).toHaveBeenCalledWith( @@ -310,95 +518,531 @@ describe('ContractSyncService', () => { const contract2 = AztecAddress.fromBigIntUnsafe(300n); it('only invalidates the targeted scope', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); service.invalidateContractForScopes(contractAddress, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // Only scopeA should be re-synced, scopeB is still cached. expectSyncedScopes([scopeA], [scopeB], [scopeA]); }); it('invalidates multiple scopes at once', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); service.invalidateContractForScopes(contractAddress, [scopeA, scopeB]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); // Both scopes should be re-synced. expectSyncedScopes([scopeA], [scopeB], [scopeA], [scopeB]); }); it('invalidating one scope does not affect the other', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); // Syncing scopeA is a no-op because it's already cached. - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); // Invalidate scopeA only. service.invalidateContractForScopes(contractAddress, [scopeA]); // Now syncing scopeA triggers a re-sync. - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB], [scopeA]); // Syncing both scopes only re-syncs scopeA (already re-synced above is cached), scopeB is still cached. - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB], [scopeA]); }); it('empty scopes is a no-op', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); service.invalidateContractForScopes(contractAddress, []); // Both scopes should still be cached since no scopes were invalidated. - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [ - scopeA, - scopeB, - ]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA, scopeB], + triggeredBy: undefined, + }); expectSyncedScopes([scopeA], [scopeB]); }); it('does not affect other contracts', async () => { - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); - await service.ensureContractSynced(contract2, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contract2, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedContracts([contractAddress, [scopeA]], [contract2, [scopeA]]); service.invalidateContractForScopes(contractAddress, [scopeA]); - await service.ensureContractSynced(contractAddress, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); - await service.ensureContractSynced(contract2, null, utilityExecutor, anchorBlockHeader, jobId, [scopeA]); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contract2, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes: [scopeA], + triggeredBy: undefined, + }); expectSyncedContracts([contractAddress, [scopeA]], [contract2, [scopeA]], [contractAddress, [scopeA]]); }); }); + describe('speculative sync', () => { + const otherContract = AztecAddress.fromBigIntUnsafe(101n); + // Calls are recorded and predicted per function, so each contract gets its own function with a distinct selector. + const entryFn: ContractFunction = { address: contractAddress, selector: new FunctionSelector(0xe1) }; + const otherFn: ContractFunction = { address: otherContract, selector: new FunctionSelector(0xe2) }; + + beforeEach(() => { + service = new ContractSyncService( + aztecNode, + contractStore, + contractClassService, + noteStore, + createLogger('test:contract-sync'), + { concurrentContractSyncEnabled: true }, + ); + }); + + it('speculatively syncs the whole predicted call tree', async () => { + const grandChild = AztecAddress.fromBigIntUnsafe(102n); + const grandChildFn: ContractFunction = { address: grandChild, selector: new FunctionSelector(0xe3) }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: entryFn, callee: otherFn }, + { caller: otherFn, callee: grandChildFn }, + ], + }); + + // A new job requests only contractAddress: its callee syncs, and so does its callee's callee. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + // The speculative syncs run in the background; yield so they reach the executor. + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]], [grandChild, [scopeA]]); + }); + + it('speculatively syncs the callees of a function of an already-synced contract', async () => { + const secondFn: ContractFunction = { address: contractAddress, selector: new FunctionSelector(0xe4) }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [{ caller: secondFn, callee: otherFn }], + }); + + // The first function syncs the contract; invoking a second function afterwards hits the sync cache, but its + // own predicted callees must still sync. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: secondFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]]); + }); + + it('speculatively syncs the dependencies of the contract sync itself', async () => { + // `sync_state` is only ever invoked by PXE, so its callees (e.g. most contract syncs query the handshake + // registry) are learned and predicted under the universal sync_state selector, not under a requested function. + const syncStateFn: ContractFunction = { address: contractAddress, selector: SYNC_STATE_SELECTOR }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [{ caller: syncStateFn, callee: otherFn }], + }); + + // A direct read syncs the contract without invoking any function, yet sync_state's learned callee still syncs. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]]); + }); + + it("speculatively syncs a predicted callee's own sync dependencies", async () => { + const thirdContract = AztecAddress.fromBigIntUnsafe(102n); + const thirdFn: ContractFunction = { address: thirdContract, selector: new FunctionSelector(0xe3) }; + const otherSyncState: ContractFunction = { address: otherContract, selector: SYNC_STATE_SELECTOR }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: entryFn, callee: otherFn }, + { caller: otherSyncState, callee: thirdFn }, + ], + }); + + // The entry function predicts otherContract, and starting otherContract's sync predicts its own sync_state's + // callee, so the whole chain syncs from a single request. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]], [thirdContract, [scopeA]]); + }); + + it('stops recursing when the known calls form a cycle', async () => { + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: entryFn, callee: otherFn }, + { caller: otherFn, callee: entryFn }, + ], + }); + + // Each contract syncs exactly once: the job's set of already-speculated functions stops the recursion when the + // predicted graph loops back to a function it already speculated from. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]]); + }); + + it('stops recursing when predicted sync dependencies form a cycle', async () => { + const ownSyncState: ContractFunction = { address: contractAddress, selector: SYNC_STATE_SELECTOR }; + const otherSyncState: ContractFunction = { address: otherContract, selector: SYNC_STATE_SELECTOR }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: ownSyncState, callee: otherFn }, + { caller: otherSyncState, callee: entryFn }, + ], + }); + + // Each contract syncs exactly once: when the chain loops back to an already-syncing contract, the warm cache + // and the job's already-speculated set stop it. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + expectSyncedContracts([contractAddress, [scopeA]], [otherContract, [scopeA]]); + }); + }); + + describe('settle', () => { + beforeEach(() => { + service = new ContractSyncService( + aztecNode, + contractStore, + contractClassService, + noteStore, + createLogger('test:contract-sync'), + { concurrentContractSyncEnabled: true }, + ); + }); + + const otherContract = AztecAddress.fromBigIntUnsafe(101n); + const entryFn: ContractFunction = { address: contractAddress, selector: new FunctionSelector(0xe1) }; + const otherFn: ContractFunction = { address: otherContract, selector: new FunctionSelector(0xe2) }; + + it('waits for a speculative sync still in flight', async () => { + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [{ caller: entryFn, callee: otherFn }], + }); + + // otherContract's sync_state hangs until released, keeping its speculative sync in flight. + const { promise: speculativeSync, resolve: releaseSpeculative } = promiseWithResolvers(); + utilityExecutor.mockImplementation(call => { + if (call.to.equals(otherContract)) { + return speculativeSync; + } + return Promise.resolve(); + }); + + // The job only requests contractAddress, so nothing awaits otherContract's speculative sync. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + + let settled = false; + const settlePromise = service.settle('job-3').then(() => { + settled = true; + }); + await tick(); + expect(settled).toBe(false); + + releaseSpeculative(); + await settlePromise; + }); + + it('waits for speculative syncs fired while settling', async () => { + const lateContract = AztecAddress.fromBigIntUnsafe(102n); + const lateFn: ContractFunction = { address: lateContract, selector: new FunctionSelector(0xe3) }; + const thirdContract = AztecAddress.fromBigIntUnsafe(103n); + const thirdFn: ContractFunction = { address: thirdContract, selector: new FunctionSelector(0xe4) }; + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [ + { caller: entryFn, callee: otherFn }, + { caller: lateFn, callee: thirdFn }, + ], + }); + + // otherContract's speculative sync is held until the job is already settling. When released, its sync_state + // makes a nested call to lateContract, whose predicted callee (thirdContract) fires a fresh speculative sync + // mid-drain, hanging until released. + const { promise: otherGate, resolve: releaseOther } = promiseWithResolvers(); + const { promise: thirdSync, resolve: releaseThird } = promiseWithResolvers(); + utilityExecutor.mockImplementation(async call => { + if (call.to.equals(otherContract)) { + await otherGate; + await service.ensureContractSynced({ + contract: lateContract, + functionToInvokeAfterSync: lateFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: otherFn, + }); + return; + } + return call.to.equals(thirdContract) ? thirdSync : Promise.resolve(); + }); + + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + + let settled = false; + const settlePromise = service.settle('job-3').then(() => { + settled = true; + }); + releaseOther(); + await tick(); + expect(settled).toBe(false); + + releaseThird(); + await settlePromise; + }); + + it('resolves immediately when the job started no syncs', async () => { + await expect(service.settle('unknown-job')).resolves.toBeUndefined(); + }); + + it('rejects when a speculative sync failed, even though no request observed the failure', async () => { + await learnDependencies({ + count: PREDICTION_THRESHOLD, + calls: [{ caller: entryFn, callee: otherFn }], + }); + + utilityExecutor.mockImplementation(call => + call.to.equals(otherContract) ? Promise.reject(new Error('speculative boom')) : Promise.resolve(), + ); + + // The job only requests contractAddress, so the failed speculative sync of otherContract rejects no request. + await service.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: entryFn.selector, + utilityExecutor, + anchorBlockHeader, + jobId: 'job-3', + scopes: [scopeA], + triggeredBy: undefined, + }); + await tick(); + + const settleError = await service.settle('job-3').then( + () => undefined, + (err: AggregateError) => err, + ); + expect(settleError).toBeInstanceOf(AggregateError); + expect(settleError!.message).toContain('Speculative syncs failed'); + expect(settleError!.errors.map((err: Error) => err.message)).toEqual(['speculative boom']); + }); + }); + + /** + * Runs `count` committed jobs, each using the first caller as the entry and observing the given direct calls, then + * wipes the sync cache (as an anchor block change would) so the next job's syncs run for real. + */ + const learnDependencies = async ({ count, calls }: { count: number; calls: Call[] }) => { + const sync = (id: string, { address, selector }: ContractFunction, triggeredBy: ContractFunction | undefined) => + service.ensureContractSynced({ + contract: address, + functionToInvokeAfterSync: selector, + utilityExecutor, + anchorBlockHeader, + jobId: id, + scopes: [scopeA], + triggeredBy, + }); + for (let i = 0; i < count; i++) { + const id = `learn-job-${i}`; + await sync(id, calls[0].caller, undefined); + for (const { caller, callee } of calls) { + await sync(id, callee, caller); + } + await service.commit(id); + } + service.wipe(); + utilityExecutor.mockClear(); + }; + /** Asserts the utility executor was called exactly with the given sequence of scope arrays. */ const expectSyncedScopes = (...expectedScopes: AztecAddress[][]) => { expect(utilityExecutor).toHaveBeenCalledTimes(expectedScopes.length); @@ -423,3 +1067,17 @@ describe('ContractSyncService', () => { /** Yields to the macrotask queue, draining all pending microtasks (semaphore acquires/releases) in between. */ const tick = () => new Promise(resolve => setImmediate(resolve)); }); + +describe('SYNC_STATE_SELECTOR', () => { + // Pins the hardcoded selector to the macro's actual output: if the macro ever changes `sync_state`'s signature, + // this fails instead of predictions silently keying on a stale selector. + it('matches the selector of a compiled artifact', async () => { + const syncState = TestContractArtifact.functions.find(f => f.name === 'sync_state'); + expect(syncState).toBeDefined(); + const expected = await FunctionSelector.fromNameAndParameters(syncState!.name, syncState!.parameters); + expect(SYNC_STATE_SELECTOR).toEqual(expected); + }); +}); + +/** A direct call observed by a job. */ +type Call = { caller: ContractFunction; callee: ContractFunction }; diff --git a/yarn-project/pxe/src/contract/contract_sync_service.ts b/yarn-project/pxe/src/contract/contract_sync_service.ts index 35448e069141..64190afb29bf 100644 --- a/yarn-project/pxe/src/contract/contract_sync_service.ts +++ b/yarn-project/pxe/src/contract/contract_sync_service.ts @@ -2,15 +2,17 @@ import type { Logger } from '@aztec/foundation/log'; import { allToCompletion } from '@aztec/foundation/promise'; import { Semaphore } from '@aztec/foundation/queue'; import { isProtocolContract } from '@aztec/protocol-contracts'; -import type { FunctionCall, FunctionSelector } from '@aztec/stdlib/abi'; +import { type FunctionCall, FunctionSelector } from '@aztec/stdlib/abi'; import type { AztecAddress } from '@aztec/stdlib/aztec-address'; import type { AztecNode } from '@aztec/stdlib/interfaces/client'; import type { BlockHeader } from '@aztec/stdlib/tx'; +import type { ContractSyncConfig } from '../config/index.js'; import type { StagedStore } from '../job_coordinator/job_coordinator.js'; import { NoteService } from '../notes/note_service.js'; import type { ContractStore } from '../storage/contract_store/contract_store.js'; import type { NoteStore } from '../storage/note_store/note_store.js'; +import { type CallKey, ContractCallGraph, type ContractFunction, toCallKey } from './contract_call_graph.js'; import type { ContractClassService } from './contract_class_service.js'; import { syncScope } from './helpers.js'; @@ -20,6 +22,12 @@ import { syncScope } from './helpers.js'; */ export const MAX_CONCURRENT_SCOPE_SYNCS = 5; +/** + * Selector of the macro-generated `sync_state` utility function, which is the same for every contract. + * Pinned against a compiled artifact in tests, so a macro signature change fails there. + */ +export const SYNC_STATE_SELECTOR = FunctionSelector.fromString('0x418ef5da'); + /** * Service for syncing the private state of contracts. It uses a cache to avoid redundant sync operations - the cache * is wiped when the anchor block changes. @@ -30,9 +38,14 @@ export class ContractSyncService implements StagedStore { readonly storeName = 'contract_sync'; // Tracks contracts synced since last wipe. The cache is keyed per individual scope address - // (`contractAddress:scopeAddress`), or `contractAddress:*` for all scopes (all accounts). - // The value is a promise that resolves when the contract is synced. - private syncedContracts: Map> = new Map(); + // (`contractAddress:scopeAddress`). The value is a promise that resolves when the contract is synced. + private readonly syncedContracts: Map> = new Map(); + + // Per-job speculation state, dropped when the job commits or discards. + private readonly speculationByJob: Map = new Map(); + + // Predicts a function's callees from the calls observed in past jobs, driving speculative sync. + private readonly callGraph: ContractCallGraph; constructor( private aztecNode: AztecNode, @@ -40,38 +53,63 @@ export class ContractSyncService implements StagedStore { private contractClassService: ContractClassService, private noteStore: NoteStore, private log: Logger, - ) {} + { concurrentContractSyncEnabled }: ContractSyncConfig, + ) { + this.callGraph = new ContractCallGraph(concurrentContractSyncEnabled); + } /** * Ensures a contract's private state is synchronized. * Uses a cache to avoid redundant sync operations - the cache is wiped when the anchor block changes. - * @param contractAddress - The address of the contract to sync. - * @param functionToInvokeAfterSync - The function selector that will be called after sync (used to validate it's - * not sync_state itself). - * @param utilityExecutor - Executor function for running the sync_state utility function. - * @param scopes - Access scopes to pass through to the utility executor (affects whose account's private state is discovered). */ - async ensureContractSynced( - contractAddress: AztecAddress, - functionToInvokeAfterSync: FunctionSelector | null, - utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise, - anchorBlockHeader: BlockHeader, - jobId: string, - scopes: AztecAddress[], - ): Promise { - this.#startSyncIfNeeded(contractAddress, scopes, anchorBlockHeader, jobId, scope => - syncScope( - contractAddress, - this.contractStore, - this.contractClassService, - anchorBlockHeader, - functionToInvokeAfterSync, - utilityExecutor, - scope, - ), + async ensureContractSynced({ + contract, + functionToInvokeAfterSync, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes, + triggeredBy, + }: ContractSyncRequest): Promise { + // A call is recorded only when both functions are known: the invoked callee and the caller that triggered it. + if (functionToInvokeAfterSync && triggeredBy) { + this.callGraph.recordCall({ + jobId, + caller: triggeredBy, + callee: { address: contract, selector: functionToInvokeAfterSync }, + }); + } + + await this.#startSyncIfNeeded( + contract, + functionToInvokeAfterSync, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes, ); + } - await this.#awaitSync(contractAddress, scopes); + /** + * Waits until every speculative sync the job fired has finished, then rejects if any failed, so the job discards + * instead of committing. This is needed because a sync that fails midway can leave partial staged writes, and a + * speculative failure might not be surfaced by any request. + */ + async settle(jobId: JobId): Promise { + // A speculative sync's execution can fire more speculative syncs mid-drain, so loop until no new promises + // appear, and only escalate once nothing is still writing. + const { syncs } = this.#speculationForJob(jobId); + const failures: unknown[] = []; + while (syncs.length > 0) { + const results = await Promise.allSettled(syncs.splice(0)); + failures.push(...results.filter(result => result.status === 'rejected').map(result => result.reason)); + } + if (failures.length > 0) { + throw new AggregateError( + failures, + 'Speculative syncs failed, so the job must discard its staged writes instead of committing', + ); + } } /** Clears sync cache entries for the given scopes of a contract. */ @@ -88,14 +126,18 @@ export class ContractSyncService implements StagedStore { this.syncedContracts.clear(); } - commit(_jobId: string): Promise { + commit(jobId: JobId): Promise { + this.callGraph.commitJob(jobId); + this.speculationByJob.delete(jobId); return Promise.resolve(); } - discardStaged(_jobId: string): Promise { + discardStaged(jobId: JobId): Promise { // We clear the synced contracts cache here because, when the job is discarded, any associated database writes from // the sync are also undone. this.syncedContracts.clear(); + this.callGraph.discardJob(jobId); + this.speculationByJob.delete(jobId); return Promise.resolve(); } @@ -103,38 +145,113 @@ export class ContractSyncService implements StagedStore { * For each unsynced scope, creates a promise that waits on: * 1. Note nullifier sync (shared, batched across all unsynced scopes). * 2. Per-scope sync (individual, semaphore-bounded). + * When concurrent contract sync is enabled, the predicted direct callees of the invoked function and of the + * contract's `sync_state` start speculatively too, once the contract's own syncs have started (see + * {@link #speculativelySync}). + * @returns A promise that resolves once every requested scope is synced, including syncs already in flight from + * concurrent calls. Speculative syncs are not included: those are only awaited by a later request that needs + * their contract, or by the job's {@link settle}. */ - #startSyncIfNeeded( + async #startSyncIfNeeded( contractAddress: AztecAddress, + functionToInvokeAfterSync: FunctionSelector | null, + utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise, + anchorBlockHeader: BlockHeader, + jobId: JobId, scopes: AztecAddress[], + ): Promise { + const scopesToSync = scopes.filter(scope => !this.syncedContracts.has(toKey(contractAddress, scope))); + if (scopesToSync.length > 0) { + this.log.debug(`Syncing contract ${contractAddress} for ${scopesToSync.length} scope(s)`); + + const syncNullifiersPromise = this.#syncNoteNullifiers(contractAddress, anchorBlockHeader, jobId, scopesToSync); + + // We build a new semaphore for each sync call, so it rate-limits the scopes within that single call. We do + // this so that if these scope syncs trigger nested syncs, the nested ones can execute without causing a deadlock. + const syncSlot = new Semaphore(MAX_CONCURRENT_SCOPE_SYNCS); + + for (const scope of scopesToSync) { + const key = toKey(contractAddress, scope); + const syncScopePromise = runBounded(syncSlot, () => + syncScope( + contractAddress, + this.contractStore, + this.contractClassService, + anchorBlockHeader, + functionToInvokeAfterSync, + utilityExecutor, + scope, + ), + ); + // This cached promise is what every later request for this scope awaits, and both branches write staged data, + // so it must run both to completion even when one fails. + const promise = allToCompletion([syncNullifiersPromise, syncScopePromise]) + .then(() => {}) + .catch(err => { + this.syncedContracts.delete(key); + throw err; + }); + this.syncedContracts.set(key, promise); + } + } + + // `sync_state` itself calls other contracts (e.g. most contract syncs query the handshake registry), so its + // predicted callees start syncing alongside the contract's own syncs. + this.#speculativelySync(contractAddress, SYNC_STATE_SELECTOR, utilityExecutor, anchorBlockHeader, jobId, scopes); + this.#speculativelySync( + contractAddress, + functionToInvokeAfterSync, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes, + ); + + await this.#awaitSync(contractAddress, scopes); + } + + /** + * Starts the syncs of the contracts the given function is predicted to call (see {@link ContractCallGraph} for how + * predictions are learned). Each started sync speculates from its own function in turn, so the whole predicted call + * tree syncs in parallel with the contract instead of one contract at a time as execution reaches it. + * + * A wrong prediction is cheap: the extra node requests are batched into round trips the job already makes, and the + * synced data simply goes unused. + */ + #speculativelySync( + contractAddress: AztecAddress, + functionToInvokeAfterSync: FunctionSelector | null, + utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise, anchorBlockHeader: BlockHeader, - jobId: string, - syncScopeFn: (scope: AztecAddress) => Promise, + jobId: JobId, + scopes: AztecAddress[], ): void { - const scopesToSync = scopes.filter(scope => !this.syncedContracts.has(toKey(contractAddress, scope))); - if (scopesToSync.length === 0) { + // Without a function there is no key to predict from (the request is a direct read). + if (!functionToInvokeAfterSync) { return; } - - this.log.debug(`Syncing contract ${contractAddress} for ${scopesToSync.length} scope(s)`); - - const syncNullifiersPromise = this.#syncNoteNullifiers(contractAddress, anchorBlockHeader, jobId, scopesToSync); - - // We build a new semaphore for each sync call, so it rate-limits the scopes within that single call. We do - // this so that if these scope syncs trigger nested syncs, the nested ones can execute without causing a deadlock. - const syncSlot = new Semaphore(MAX_CONCURRENT_SCOPE_SYNCS); - - for (const scope of scopesToSync) { - const key = toKey(contractAddress, scope); - // This cached promise is what every later request for this scope awaits, and both branches write staged data, - // so it must run both to completion even when one fails. - const promise = allToCompletion([syncNullifiersPromise, runBounded(syncSlot, () => syncScopeFn(scope))]) - .then(() => {}) - .catch(err => { - this.syncedContracts.delete(key); - throw err; - }); - this.syncedContracts.set(key, promise); + const speculation = this.#speculationForJob(jobId); + const caller: ContractFunction = { address: contractAddress, selector: functionToInvokeAfterSync }; + for (const callee of this.callGraph.predictDirectCallees(caller)) { + // The job's set of already-speculated functions stops the recursion when the predicted graph has a cycle. + if (speculation.speculated.has(toCallKey(callee))) { + continue; + } + speculation.speculated.add(toCallKey(callee)); + const syncPromise = this.#startSyncIfNeeded( + callee.address, + callee.selector, + utilityExecutor, + anchorBlockHeader, + jobId, + scopes, + ); + speculation.syncs.push(syncPromise); + // `settle` only escalates these failures at the end of the job: catch here so one does not become an unhandled + // rejection before then, and log it. + syncPromise.catch(err => { + this.log.warn(`Speculative sync of ${callee.address} failed`, { jobId, error: err?.message }); + }); } } @@ -142,7 +259,7 @@ export class ContractSyncService implements StagedStore { async #syncNoteNullifiers( contractAddress: AztecAddress, anchorBlockHeader: BlockHeader, - jobId: string, + jobId: JobId, scopes: AztecAddress[], ): Promise { // Protocol contracts don't have private state to sync @@ -155,6 +272,15 @@ export class ContractSyncService implements StagedStore { await noteService.syncNoteNullifiers(contractAddress, scopes); } + #speculationForJob(jobId: JobId): JobSpeculation { + let speculation = this.speculationByJob.get(jobId); + if (!speculation) { + speculation = { speculated: new Set(), syncs: [] }; + this.speculationByJob.set(jobId, speculation); + } + return speculation; + } + /** Collects all relevant scope promises (including in-flight ones from concurrent calls) and awaits them. */ async #awaitSync(contractAddress: AztecAddress, scopes: AztecAddress[]): Promise { const promises = scopes @@ -164,7 +290,44 @@ export class ContractSyncService implements StagedStore { } } -function toKey(contract: AztecAddress, scope: AztecAddress) { +/** A request to synchronize a contract's private state. */ +type ContractSyncRequest = { + /** The contract to sync. */ + contract: AztecAddress; + /** + * The function that will be invoked after the sync, or null when nothing will be invoked (e.g. reading + * notes/events directly). + */ + functionToInvokeAfterSync: FunctionSelector | null; + /** Executes a utility function call under the given scopes. Syncs run each contract's sync_state through it. */ + utilityExecutor: (call: FunctionCall, scopes: AztecAddress[]) => Promise; + /** The anchor block to sync at. */ + anchorBlockHeader: BlockHeader; + /** The job requesting the sync. */ + jobId: JobId; + /** Access scopes to pass through to the utility executor (affects whose account's private state is discovered). */ + scopes: AztecAddress[]; + /** + * The function whose execution triggered this sync request, or undefined when the request is a job's top-level use + * (an entry call or a direct read) rather than a nested call. + */ + triggeredBy: ContractFunction | undefined; +}; + +type JobId = string; + +/** A job's speculation state. */ +type JobSpeculation = { + /** Functions prediction already ran for, so the recursion stops on cycles in the predicted graph. */ + speculated: Set; + /** Every sync fired by prediction, awaited by {@link settle} before the job commits or discards. */ + syncs: Promise[]; +}; + +/** Key of a contract's sync cache entry for a single scope: `contractAddress:scopeAddress`. */ +type SyncKey = `0x${string}:0x${string}`; + +function toKey(contract: AztecAddress, scope: AztecAddress): SyncKey { return `${contract.toString()}:${scope.toString()}`; } diff --git a/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts b/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts index 46b99e09ea28..9e2980868558 100644 --- a/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts +++ b/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution.test.ts @@ -308,13 +308,13 @@ describe('Private Execution test suite', () => { txResolver.resolveTxs.mockResolvedValue([]); // Configure mock to actually perform sync_state calls (needed for nested call tests) contractSyncService.ensureContractSynced.mockImplementation( - async (contractAddress, functionToInvokeAfterSync, utilityExecutor, _anchorBlockHeader, _jobId, scopes) => { + async ({ contract, functionToInvokeAfterSync, utilityExecutor, anchorBlockHeader, scopes }) => { for (const scope of scopes) { await syncScope( - contractAddress, + contract, contractStore, contractClassService, - _anchorBlockHeader, + anchorBlockHeader, functionToInvokeAfterSync, utilityExecutor, scope, diff --git a/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts b/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts index e12de2eec4ac..0d7fc792851a 100644 --- a/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts +++ b/yarn-project/pxe/src/contract_function_simulator/oracle/private_execution_oracle.ts @@ -662,14 +662,15 @@ export class PrivateExecutionOracle extends UtilityExecutionOracle implements IP isStaticCall = isStaticCall || this.callContext.isStaticCall; - await this.contractSyncService.ensureContractSynced( - targetContractAddress, - functionSelector, - this.utilityExecutor, - this.anchorBlockHeader, - this.jobId, - this.scopes, - ); + await this.contractSyncService.ensureContractSynced({ + contract: targetContractAddress, + functionToInvokeAfterSync: functionSelector, + utilityExecutor: this.utilityExecutor, + anchorBlockHeader: this.anchorBlockHeader, + jobId: this.jobId, + scopes: this.scopes, + triggeredBy: { address: this.callContext.contractAddress, selector: this.callContext.functionSelector }, + }); const targetArtifact = await this.anchoredContractData.getFunctionArtifactWithDebugMetadata( targetContractAddress, diff --git a/yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts b/yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts index bb7677c4446e..244ceb604796 100644 --- a/yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts +++ b/yarn-project/pxe/src/contract_function_simulator/oracle/utility_execution_oracle.ts @@ -1070,14 +1070,15 @@ export class UtilityExecutionOracle implements IMiscOracle, IUtilityExecutionOra } } - await this.contractSyncService.ensureContractSynced( - targetContractAddress, - functionSelector, - this.utilityExecutor, - this.anchorBlockHeader, - this.jobId, - this.scopes, - ); + await this.contractSyncService.ensureContractSynced({ + contract: targetContractAddress, + functionToInvokeAfterSync: functionSelector, + utilityExecutor: this.utilityExecutor, + anchorBlockHeader: this.anchorBlockHeader, + jobId: this.jobId, + scopes: this.scopes, + triggeredBy: { address: this.contractAddress, selector: this.callContext.functionSelector }, + }); } this.logger.debug( diff --git a/yarn-project/pxe/src/debug/pxe_debug_utils.ts b/yarn-project/pxe/src/debug/pxe_debug_utils.ts index 69156f484cf8..472db48ce532 100644 --- a/yarn-project/pxe/src/debug/pxe_debug_utils.ts +++ b/yarn-project/pxe/src/debug/pxe_debug_utils.ts @@ -69,15 +69,16 @@ export class PXEDebugUtils { const contractFunctionSimulator = this.#getSimulatorForTx(); - await this.contractSyncService.ensureContractSynced( - filter.contractAddress, - null, - async (privateSyncCall, execScopes) => + await this.contractSyncService.ensureContractSynced({ + contract: filter.contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor: async (privateSyncCall, execScopes) => await this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId), anchorBlockHeader, jobId, - filter.scopes, - ); + scopes: filter.scopes, + triggeredBy: undefined, + }); return this.noteStore.getNotes(filter, jobId); }); diff --git a/yarn-project/pxe/src/job_coordinator/job_coordinator.test.ts b/yarn-project/pxe/src/job_coordinator/job_coordinator.test.ts index aa52728ee98a..697fa91952c0 100644 --- a/yarn-project/pxe/src/job_coordinator/job_coordinator.test.ts +++ b/yarn-project/pxe/src/job_coordinator/job_coordinator.test.ts @@ -1,3 +1,4 @@ +import { promiseWithResolvers } from '@aztec/foundation/promise'; import type { AztecAsyncKVStore } from '@aztec/kv-store'; import { openTmpStore } from '@aztec/kv-store/lmdb-v2'; @@ -65,6 +66,50 @@ describe('JobCoordinator', () => { expect(commitMock).toHaveBeenCalledWith(jobId); }); + + it('waits for stores to settle before committing any of them', async () => { + const { promise: settling, resolve: finishSettling } = promiseWithResolvers(); + const commitMock = jest.fn<() => Promise>().mockResolvedValue(undefined); + coordinator.registerStore({ + storeName: 'settling_store', + commit: () => Promise.resolve(), + discardStaged: () => Promise.resolve(), + settle: () => settling, + }); + coordinator.registerStore({ + storeName: 'other_store', + commit: commitMock, + discardStaged: () => Promise.resolve(), + }); + + const jobId = coordinator.beginJob(); + const commitPromise = coordinator.commitJob(jobId); + await tick(); + expect(commitMock).not.toHaveBeenCalled(); + + finishSettling(); + await commitPromise; + expect(commitMock).toHaveBeenCalledWith(jobId); + }); + + it('propagates a settle rejection without committing any store', async () => { + const commitMock = jest.fn<() => Promise>().mockResolvedValue(undefined); + coordinator.registerStore({ + storeName: 'failing_store', + commit: () => Promise.resolve(), + discardStaged: () => Promise.resolve(), + settle: () => Promise.reject(new Error('settle failed')), + }); + coordinator.registerStore({ + storeName: 'other_store', + commit: commitMock, + discardStaged: () => Promise.resolve(), + }); + + const jobId = coordinator.beginJob(); + await expect(coordinator.commitJob(jobId)).rejects.toThrow('settle failed'); + expect(commitMock).not.toHaveBeenCalled(); + }); }); describe('abortJob', () => { @@ -93,6 +138,53 @@ describe('JobCoordinator', () => { expect(discardStagedMock).toHaveBeenCalledWith(jobId); }); + + it('waits for stores to settle before discarding any of them', async () => { + const { promise: settling, resolve: finishSettling } = promiseWithResolvers(); + const discardStagedMock = jest.fn<() => Promise>().mockResolvedValue(undefined); + coordinator.registerStore({ + storeName: 'settling_store', + commit: () => Promise.resolve(), + discardStaged: () => Promise.resolve(), + settle: () => settling, + }); + coordinator.registerStore({ + storeName: 'other_store', + commit: () => Promise.resolve(), + discardStaged: discardStagedMock, + }); + + const jobId = coordinator.beginJob(); + const abortPromise = coordinator.abortJob(jobId); + await tick(); + expect(discardStagedMock).not.toHaveBeenCalled(); + + finishSettling(); + await abortPromise; + expect(discardStagedMock).toHaveBeenCalledWith(jobId); + }); + + it('discards all stores even when a settle rejects', async () => { + const failingDiscardMock = jest.fn<() => Promise>().mockResolvedValue(undefined); + const otherDiscardMock = jest.fn<() => Promise>().mockResolvedValue(undefined); + coordinator.registerStore({ + storeName: 'failing_store', + commit: () => Promise.resolve(), + discardStaged: failingDiscardMock, + settle: () => Promise.reject(new Error('settle failed')), + }); + coordinator.registerStore({ + storeName: 'other_store', + commit: () => Promise.resolve(), + discardStaged: otherDiscardMock, + }); + + const jobId = coordinator.beginJob(); + await coordinator.abortJob(jobId); + + expect(failingDiscardMock).toHaveBeenCalledWith(jobId); + expect(otherDiscardMock).toHaveBeenCalledWith(jobId); + }); }); describe('registerStore', () => { @@ -110,4 +202,7 @@ describe('JobCoordinator', () => { expect(() => coordinator.registerStore(mockStore)).toThrow(/already registered/); }); }); + + /** Yields to the macrotask queue, draining all pending microtasks in between. */ + const tick = () => new Promise(resolve => setImmediate(resolve)); }); diff --git a/yarn-project/pxe/src/job_coordinator/job_coordinator.ts b/yarn-project/pxe/src/job_coordinator/job_coordinator.ts index 1c013a76b82f..28616738a52a 100644 --- a/yarn-project/pxe/src/job_coordinator/job_coordinator.ts +++ b/yarn-project/pxe/src/job_coordinator/job_coordinator.ts @@ -1,5 +1,6 @@ import { randomBytes } from '@aztec/foundation/crypto/random'; import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log'; +import { allToCompletion } from '@aztec/foundation/promise'; import type { AztecAsyncKVStore } from '@aztec/kv-store'; /** @@ -24,6 +25,16 @@ export interface StagedStore { * @param jobId - The job identifier */ discardStaged(jobId: string): Promise; + + /** + * A store may have pending work that must finish before the job's staged writes are committed or discarded, yet + * commits run inside a transaction that cannot wait for it. Such stores implement this method: it is called before + * every commit and discard, outside the transaction. If settling fails, the commit is cancelled, but a discard + * proceeds. + * + * @param jobId - The job identifier + */ + settle?(jobId: string): Promise; } /** @@ -108,6 +119,9 @@ export class JobCoordinator { this.log.debug(`Committing job ${jobId}`); + // Settling must stay outside the transaction: it can take arbitrarily long. + await allToCompletion([...this.#stores.values()].map(store => store.settle?.(jobId))); + // Commit all stores atomically in a single transaction. // Each store's commit is a no-op if it has no staged data (but that's up to each store to handle). await this.kvStore.transactionAsync(async () => { @@ -133,6 +147,8 @@ export class JobCoordinator { this.log.debug(`Aborting job ${jobId}`); + await this.#settleStoresLoggingFailures(jobId); + for (const store of this.#stores.values()) { await store.discardStaged(jobId); } @@ -147,4 +163,18 @@ export class JobCoordinator { hasJobInProgress(): boolean { return this.#currentJobId !== undefined; } + + /** + * Settles every store, logging failures instead of propagating them. The abort must run to completion no matter what, + * so a store that fails to settle cannot stop the others from discarding or mask the error that aborted the job. + */ + async #settleStoresLoggingFailures(jobId: string): Promise { + await allToCompletion( + [...this.#stores.values()].map(store => + store.settle?.(jobId).catch(err => { + this.log.warn(`Store ${store.storeName} failed to settle while aborting job ${jobId}`, { jobId, err }); + }), + ), + ); + } } diff --git a/yarn-project/pxe/src/pxe.test.ts b/yarn-project/pxe/src/pxe.test.ts index f9e71359e672..dc85297d34b6 100644 --- a/yarn-project/pxe/src/pxe.test.ts +++ b/yarn-project/pxe/src/pxe.test.ts @@ -75,6 +75,7 @@ describe('PXE', () => { l1ChainId: 31337, rollupVersion: 1, autoSync: true, + concurrentContractSyncEnabled: false, }; // Mock getNodeInfo which is called during PXE creation diff --git a/yarn-project/pxe/src/pxe.ts b/yarn-project/pxe/src/pxe.ts index a26720aeb905..9949ba890751 100644 --- a/yarn-project/pxe/src/pxe.ts +++ b/yarn-project/pxe/src/pxe.ts @@ -327,6 +327,7 @@ export class PXE { contractClassService, noteStore, createLogger('pxe:contract_sync', bindings), + config, ); const txResolver = new TxResolverService(readCachedNode); @@ -549,15 +550,16 @@ export class PXE { const { origin: contractAddress, functionSelector } = txRequest; try { - await this.contractSyncService.ensureContractSynced( - contractAddress, - functionSelector, - (privateSyncCall, execScopes) => + await this.contractSyncService.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: functionSelector, + utilityExecutor: (privateSyncCall, execScopes) => this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId), anchorBlockHeader, jobId, scopes, - ); + triggeredBy: undefined, + }); const result = await contractFunctionSimulator.run(txRequest, { anchorBlockHeader, @@ -1381,15 +1383,16 @@ export class PXE { const contractFunctionSimulator = this.#getSimulatorForTx(); const anchorBlockHeader = await this.anchorBlockStore.getBlockHeader(); - await this.contractSyncService.ensureContractSynced( - call.to, - call.selector, - (privateSyncCall, execScopes) => + await this.contractSyncService.ensureContractSynced({ + contract: call.to, + functionToInvokeAfterSync: call.selector, + utilityExecutor: (privateSyncCall, execScopes) => this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId), anchorBlockHeader, jobId, scopes, - ); + triggeredBy: undefined, + }); const { result: executionResult, offchainEffects } = await this.#executeUtility( contractFunctionSimulator, @@ -1459,15 +1462,16 @@ export class PXE { const contractFunctionSimulator = this.#getSimulatorForTx(); - await this.contractSyncService.ensureContractSynced( - filter.contractAddress, - null, - async (privateSyncCall, execScopes) => + await this.contractSyncService.ensureContractSynced({ + contract: filter.contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor: async (privateSyncCall, execScopes) => await this.#executeUtility(contractFunctionSimulator, privateSyncCall, [], execScopes, jobId), anchorBlockHeader, jobId, - filter.scopes, - ); + scopes: filter.scopes, + triggeredBy: undefined, + }); }); // anchorBlockNumber is set during the job and fixed to whatever it is after a block sync diff --git a/yarn-project/txe/src/oracle/txe_oracle_top_level_context.ts b/yarn-project/txe/src/oracle/txe_oracle_top_level_context.ts index 121bb578fad8..f88c77bc01f3 100644 --- a/yarn-project/txe/src/oracle/txe_oracle_top_level_context.ts +++ b/yarn-project/txe/src/oracle/txe_oracle_top_level_context.ts @@ -209,17 +209,18 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl return; } - const blockHeader = await this.stateMachine.anchorBlockStore.getBlockHeader(); - await this.stateMachine.contractSyncService.ensureContractSynced( - contractAddress, - null, - async (call, execScopes) => { + const anchorBlockHeader = await this.stateMachine.anchorBlockStore.getBlockHeader(); + await this.stateMachine.contractSyncService.ensureContractSynced({ + contract: contractAddress, + functionToInvokeAfterSync: null, + utilityExecutor: async (call, execScopes) => { await this.executeUtilityCall(call, { scopes: execScopes, jobId }); }, - blockHeader, + anchorBlockHeader, jobId, - [scope], - ); + scopes: [scope], + triggeredBy: undefined, + }); } async getPrivateEvents(selector: EventSelector, contractAddress: AztecAddress, scope: AztecAddress) { @@ -451,14 +452,15 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl await this.executeUtilityCall(call, { scopes: execScopes, jobId }); }; - await this.stateMachine.contractSyncService.ensureContractSynced( - targetContractAddress, - functionSelector, + await this.stateMachine.contractSyncService.ensureContractSynced({ + contract: targetContractAddress, + functionToInvokeAfterSync: functionSelector, utilityExecutor, - blockHeader, + anchorBlockHeader: blockHeader, jobId, scopes, - ); + triggeredBy: undefined, + }); const blockNumber = await this.getNextBlockNumber(); @@ -859,16 +861,17 @@ export class TXEOracleTopLevelContext implements IMiscOracle, ITxeExecutionOracl } // Sync notes before executing utility function to discover notes from previous transactions - await this.stateMachine.contractSyncService.ensureContractSynced( - targetContractAddress, - functionSelector, - async (call, execScopes) => { + await this.stateMachine.contractSyncService.ensureContractSynced({ + contract: targetContractAddress, + functionToInvokeAfterSync: functionSelector, + utilityExecutor: async (call, execScopes) => { await this.executeUtilityCall(call, { scopes: execScopes, jobId }); }, - blockHeader, + anchorBlockHeader: blockHeader, jobId, - await this.keyStore.getAccounts(), - ); + scopes: await this.keyStore.getAccounts(), + triggeredBy: undefined, + }); const call = FunctionCall.from({ name: artifact.name, diff --git a/yarn-project/txe/src/state_machine/index.ts b/yarn-project/txe/src/state_machine/index.ts index 4f724f3f8130..ddad4c0375ba 100644 --- a/yarn-project/txe/src/state_machine/index.ts +++ b/yarn-project/txe/src/state_machine/index.ts @@ -82,6 +82,7 @@ export class TXEStateMachine { contractClassService, noteStore, createLogger('txe:contract_sync'), + { concurrentContractSyncEnabled: false }, ); const txResolver = new TxResolverService(node); diff --git a/yarn-project/txe/src/txe_session.ts b/yarn-project/txe/src/txe_session.ts index 7e8fa64c16dd..267b6ab8180f 100644 --- a/yarn-project/txe/src/txe_session.ts +++ b/yarn-project/txe/src/txe_session.ts @@ -331,6 +331,10 @@ export class TXESession implements TXESessionStateHandler { const keyStore = new KeyStore(store); const accountStore = new TXEAccountStore(store); + const archiver = new TXEArchiver(store); + const anchorBlockStore = new AnchorBlockStore(store); + const stateMachine = await TXEStateMachine.create(archiver, anchorBlockStore, contractStore, noteStore); + const jobCoordinator = new JobCoordinator(store); jobCoordinator.registerStores([ capsuleStore, @@ -339,12 +343,9 @@ export class TXESession implements TXESessionStateHandler { recipientTaggingStore, privateEventStore, noteStore, + stateMachine.contractSyncService, ]); - const archiver = new TXEArchiver(store); - const anchorBlockStore = new AnchorBlockStore(store); - const stateMachine = await TXEStateMachine.create(archiver, anchorBlockStore, contractStore, noteStore); - const nextBlockTimestamp = BigInt(Math.floor(new Date().getTime() / 1000)); const version = new Fr(await stateMachine.node.getVersion()); const chainId = new Fr(await stateMachine.node.getChainId());