Fix imports, add worker death

Author: nventuro

.gitignore

@@ -31,3 +31,6 @@ docs/docs/protocol-specs/public-vm/gen/
 # for those who use Claude Code
 CLAUDE.md
 .claude
+
+# 0x flamegraphs
+*.0x

cspell.json

@@ -125,6 +125,7 @@
     "Fieldable",
     "Filecoin",
     "filestat",
+    "flamegraphs",
     "Flashbots",
     "flatmap",
     "foundryup",

yarn-project/txe/src/bin/index.ts

@@ -1,6 +1,6 @@
 #!/usr/bin/env -S node --no-warnings
-import { createLogger } from '@aztec/aztec.js';
 import { startHttpRpcServer } from '@aztec/foundation/json-rpc/server';
+import { createLogger } from '@aztec/foundation/log';
 
 import { createTXERpcServer } from '../index.js';
 

yarn-project/txe/src/index.ts

@@ -1,19 +1,19 @@
 import { SchnorrAccountContractArtifact } from '@aztec/accounts/schnorr';
-import {
-  AztecAddress,
-  type ContractInstanceWithAddress,
-  Fr,
-  type NoirCompiledContract,
-  PublicKeys,
-  deriveKeys,
-  getContractInstanceFromInstantiationParams,
-  loadContractArtifact,
-} from '@aztec/aztec.js';
+import { Fr } from '@aztec/foundation/fields';
 import { createSafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 import type { Logger } from '@aztec/foundation/log';
+import { Timer } from '@aztec/foundation/timer';
 import { type ProtocolContract, protocolContractNames } from '@aztec/protocol-contracts';
 import { BundledProtocolContractsProvider } from '@aztec/protocol-contracts/providers/bundle';
-import { computeArtifactHash } from '@aztec/stdlib/contract';
+import { loadContractArtifact } from '@aztec/stdlib/abi';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import {
+  type ContractInstanceWithAddress,
+  computeArtifactHash,
+  getContractInstanceFromInstantiationParams,
+} from '@aztec/stdlib/contract';
+import { PublicKeys, deriveKeys } from '@aztec/stdlib/keys';
+import type { NoirCompiledContract } from '@aztec/stdlib/noir';
 import type { ApiSchemaFor, ZodFor } from '@aztec/stdlib/schemas';
 
 import { createHash } from 'crypto';
@@ -41,6 +41,7 @@ import type { ContractArtifactWithHash } from './util/txe_contract_data_provider
 interface SessionWorker {
   worker: Worker;
   processing: boolean;
+  timeout?: NodeJS.Timeout;
 }
 
 const sessionWorkers = new Map<number, SessionWorker>();
@@ -93,21 +94,29 @@ class TXEDispatcher {
   private async createSessionWorker(): Promise<SessionWorker> {
     // Ensure protocol contracts are loaded
     if (!this.protocolContracts) {
+      this.logger.debug('Loading protocol contracts');
+      const loadTimer = new Timer();
       this.protocolContracts = await Promise.all(
         protocolContractNames.map(name => new BundledProtocolContractsProvider().getProtocolContractArtifact(name)),
       );
+      this.logger.debug(`Protocol Contract loading took ${loadTimer.ms()}ms`);
     }
 
     // Serialize protocol contracts for worker
+    const serTimer = new Timer();
     const serializedProtocolContracts = serializeProtocolContracts(this.protocolContracts);
+    this.logger.debug(`Protocol Contract serialization took ${serTimer.ms()}ms`);
 
+    const createTimer = new Timer();
     const worker = new Worker('./dest/session_worker.js', {
       workerData: { serializedProtocolContracts },
     });
 
     // Wait for the worker to signal it's ready
     await this.waitForWorkerReady(worker);
 
+    this.logger.debug(`Worker creation took ${createTimer.ms()}ms`);
+
     return {
       worker,
       processing: false,
@@ -307,8 +316,23 @@ class TXEDispatcher {
     // Get or create a worker for this session
     const sessionWorker = await this.getOrCreateWorker(sessionId);
 
+    if (sessionWorker.timeout) {
+      clearTimeout(sessionWorker.timeout);
+    }
+
     // Send the request to the worker and wait for the response
-    return await this.sendToWorker(sessionWorker, sessionId, functionName, inputs);
+    const result = await this.sendToWorker(sessionWorker, sessionId, functionName, inputs);
+
+    sessionWorker.timeout = setTimeout(() => {
+      // eslint-disable-next-line @typescript-eslint/no-floating-promises
+      (async () => {
+        this.logger.debug(`Killing session ${sessionId}`);
+        sessionWorkers.delete(sessionId);
+        await sessionWorker.worker.terminate();
+      })();
+    }, 1500);
+
+    return result;
   }
 }
 

yarn-project/txe/src/oracle/txe_oracle.ts

@@ -1,4 +1,3 @@
-import { type AztecNode, Body, L2Block, Note } from '@aztec/aztec.js';
 import {
   CONTRACT_INSTANCE_REGISTRY_CONTRACT_ADDRESS,
   DEFAULT_DA_GAS_LIMIT,
@@ -26,7 +25,7 @@ import {
   PrivateEventDataProvider,
   TaggingDataProvider,
   enrichPublicSimulationError,
-} from '@aztec/pxe/server';
+} from '@aztec/pxe/client/lazy';
 import {
   ExecutionNoteCache,
   HashedValuesCache,
@@ -60,6 +59,7 @@ import {
 import { AuthWitness } from '@aztec/stdlib/auth-witness';
 import { PublicDataWrite } from '@aztec/stdlib/avm';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { Body, L2Block } from '@aztec/stdlib/block';
 import { type ContractInstance, type ContractInstanceWithAddress, computePartialAddress } from '@aztec/stdlib/contract';
 import { Gas, GasFees, GasSettings } from '@aztec/stdlib/gas';
 import {
@@ -71,7 +71,7 @@ import {
   siloNoteHash,
   siloNullifier,
 } from '@aztec/stdlib/hash';
-import type { MerkleTreeReadOperations, MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
+import type { AztecNode, MerkleTreeReadOperations, MerkleTreeWriteOperations } from '@aztec/stdlib/interfaces/server';
 import {
   type KeyValidationRequest,
   PartialPrivateTailPublicInputsForPublic,
@@ -81,7 +81,7 @@ import {
   PublicCallRequest,
 } from '@aztec/stdlib/kernel';
 import { ContractClassLog, IndexedTaggingSecret, PrivateLog, type PublicLog } from '@aztec/stdlib/logs';
-import type { NoteStatus } from '@aztec/stdlib/note';
+import { Note, type NoteStatus } from '@aztec/stdlib/note';
 import { ClientIvcProof } from '@aztec/stdlib/proofs';
 import {
   makeAppendOnlyTreeSnapshot,
@@ -335,7 +335,7 @@ export class TXE extends TXETypedOracle {
     inputs.callContext = new CallContext(this.msgSender, this.contractAddress, this.functionSelector, isStaticCall);
     inputs.startSideEffectCounter = sideEffectsCounter;
 
-    this.logger.info(`Created private context for block ${blockNumber}`);
+    this.logger.verbose(`Created private context for block ${blockNumber}`);
 
     return inputs;
   }
@@ -705,7 +705,7 @@ export class TXE extends TXETypedOracle {
     header.globalVariables.version = new Fr(this.ROLLUP_VERSION);
     header.globalVariables.chainId = new Fr(this.CHAIN_ID);
 
-    this.logger.info(`Created block ${blockNumber} with timestamp ${header.globalVariables.timestamp}`);
+    this.logger.verbose(`Created block ${blockNumber} with timestamp ${header.globalVariables.timestamp}`);
 
     l2Block.header = header;
 

yarn-project/txe/src/oracle/txe_typed_oracle.ts

@@ -1,9 +1,10 @@
-import type { CompleteAddress, ContractArtifact, ContractInstanceWithAddress, TxHash } from '@aztec/aztec.js';
 import type { Fr } from '@aztec/foundation/fields';
 import { TypedOracle } from '@aztec/pxe/simulator';
-import type { FunctionSelector } from '@aztec/stdlib/abi';
+import type { ContractArtifact, FunctionSelector } from '@aztec/stdlib/abi';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { CompleteAddress, ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 import type { PrivateContextInputs } from '@aztec/stdlib/kernel';
+import type { TxHash } from '@aztec/stdlib/tx';
 import type { UInt64 } from '@aztec/stdlib/types';
 
 class OracleMethodNotAvailableError extends Error {

yarn-project/txe/src/session_worker.ts

@@ -1,3 +1,4 @@
+import { Timer } from '@aztec/foundation/timer';
 import type { ProtocolContract } from '@aztec/protocol-contracts';
 
 import { parentPort, workerData } from 'worker_threads';

yarn-project/txe/src/state_machine/index.ts

@@ -2,7 +2,7 @@ import { type AztecNodeConfig, AztecNodeService } from '@aztec/aztec-node';
 import { TestCircuitVerifier } from '@aztec/bb-prover/test';
 import { createLogger } from '@aztec/foundation/log';
 import type { AztecAsyncKVStore } from '@aztec/kv-store';
-import { SyncDataProvider } from '@aztec/pxe/server';
+import { SyncDataProvider } from '@aztec/pxe/client/lazy';
 import type { L2Block } from '@aztec/stdlib/block';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { getPackageVersion } from '@aztec/stdlib/update-checker';

yarn-project/txe/src/txe_service/txe_service.ts

@@ -1,7 +1,8 @@
-import { type ContractInstanceWithAddress, Fr, Point } from '@aztec/aztec.js';
+import { Fr, Point } from '@aztec/foundation/fields';
 import { packAsRetrievedNote } from '@aztec/pxe/simulator';
 import { type ContractArtifact, FunctionSelector, NoteSelector } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { ContractInstanceWithAddress } from '@aztec/stdlib/contract';
 import { MerkleTreeId } from '@aztec/stdlib/trees';
 
 import type { TXETypedOracle } from '../oracle/txe_typed_oracle.js';

yarn-project/txe/src/txe_session.ts

@@ -1,4 +1,3 @@
-import { AztecAddress } from '@aztec/aztec.js';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { KeyStore } from '@aztec/key-store';
 import { openTmpStore } from '@aztec/kv-store/lmdb-v2';
@@ -9,7 +8,8 @@ import {
   NoteDataProvider,
   PrivateEventDataProvider,
   TaggingDataProvider,
-} from '@aztec/pxe/server';
+} from '@aztec/pxe/client/lazy';
+import { AztecAddress } from '@aztec/stdlib/aztec-address';
 
 import { TXE } from './oracle/txe_oracle.js';
 import { TXEStateMachine } from './state_machine/index.js';