@aztec/pxe 0.0.1-commit.f2ce05ee → 0.0.1-commit.f5a9928

package/src/node/benchmarked_node.ts

@@ -0,0 +1,140 @@
+import { Timer } from '@aztec/foundation/timer';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+import type { NodeStats } from '@aztec/stdlib/tx';
+
+/*
+ * Proxy generator for an AztecNode that tracks the time taken for each RPC call and the number of round trips (actual
+ * blocking waits for node responses).
+ *
+ * A round trip is counted when we transition from 0 to 1 in-flight calls, and ends when all concurrent calls complete.
+ * This means parallel calls in Promise.all count as a single round trip.
+ *
+ * Note that batching of RPC calls in `safe_json_rpc_client.ts` could affect the round trip counts but in places we
+ * currently use this information we do not even use HTTP as we have direct access to the Aztec Node instance in TS
+ * (i.e. not running against external node) so this is not a problem for now.
+ *
+ * If you want to use this against external node and the info gets skewed by batching you can set the `maxBatchSize`
+ * value in `safe_json_rpc_client.ts` to 1 (the main motivation for batching was to get around parallel http requests
+ * limits in web browsers which is not a problem when debugging in node.js).
+ */
+export interface Recording {
+  /** Closes the recording and returns what it saw. Reads served after this are not recorded. */
+  stop(): NodeStats;
+}
+
+/** An {@link AztecNode} wrapper that can report the reads it serves. */
+export interface BenchmarkedAztecNode extends AztecNode {
+  /**
+   * Opens a recording of the reads this wrapper serves, until {@link Recording.stop}. Recordings are independent, so
+   * several may run at once: one measuring a whole run sees the same reads as one measuring a single operation.
+   *
+   * Nothing is measured while none is open, which is what keeps a long-lived wrapper from accumulating a run's worth
+   * of timings nobody asked for.
+   */
+  startRecording(): Recording;
+}
+
+/** Wraps `node` so that the reads it answers can be recorded. */
+export function withRecording(node: AztecNode): BenchmarkedAztecNode {
+  // The stats of every recording currently open
+  const open = new Set<NodeStats>();
+
+  // Round trip tracking
+  let inFlightCount = 0;
+  let currentRoundTripTimer: Timer | null = null;
+  let currentRoundTripMethods: string[] = [];
+
+  return new Proxy(node, {
+    get(target, prop) {
+      if (prop === 'startRecording') {
+        return (): Recording => {
+          const stats: NodeStats = {
+            perMethod: {},
+            roundTrips: { roundTrips: 0, totalBlockingTime: 0, roundTripDurations: [], roundTripMethods: [] },
+          };
+          open.add(stats);
+          return {
+            stop: () => {
+              open.delete(stats);
+              return stats;
+            },
+          };
+        };
+      }
+
+      const value = Reflect.get(target, prop);
+      if (typeof value !== 'function' || typeof prop !== 'string') {
+        return value;
+      }
+
+      return (...args: unknown[]) => {
+        // With no recording open there is nobody to report to, so the read is left untimed.
+        if (open.size === 0) {
+          return value.apply(target, args);
+        }
+
+        // Start of a new round trip batch?
+        if (inFlightCount === 0) {
+          currentRoundTripTimer = new Timer();
+          currentRoundTripMethods = [];
+        }
+        inFlightCount++;
+        currentRoundTripMethods.push(prop);
+
+        const callTimer = new Timer();
+        const result = value.apply(target, args);
+
+        // Handle completion - called when the call finishes (after Promise resolves)
+        const handleCompletion = () => {
+          const callTime = callTimer.ms();
+          for (const stats of open) {
+            timesFor(stats, prop).push(callTime);
+          }
+
+          inFlightCount--;
+
+          // End of round trip batch - all concurrent calls completed
+          if (inFlightCount === 0 && currentRoundTripTimer) {
+            const roundTripTime = currentRoundTripTimer.ms();
+            for (const { roundTrips } of open) {
+              roundTrips.roundTrips++;
+              roundTrips.totalBlockingTime += roundTripTime;
+              roundTrips.roundTripDurations.push(roundTripTime);
+              roundTrips.roundTripMethods.push(currentRoundTripMethods);
+            }
+            currentRoundTripTimer = null;
+            currentRoundTripMethods = [];
+          }
+        };
+
+        // If the result is a Promise, chain the completion handler
+        if (isThenable(result)) {
+          return result.then(
+            resolved => {
+              handleCompletion();
+              return resolved;
+            },
+            error => {
+              handleCompletion();
+              throw error;
+            },
+          );
+        } else {
+          // Synchronous method - handle completion immediately
+          handleCompletion();
+          return result;
+        }
+      };
+    },
+  }) as BenchmarkedAztecNode;
+}
+
+/** The times `stats` holds for `method`, added on first use. */
+function timesFor(stats: NodeStats, method: string) {
+  return (stats.perMethod[method as keyof AztecNode] ??= { times: [] }).times;
+}
+
+/** Whether completion can be chained onto `value` with `.then`, which a promise from any implementation allows. */
+function isThenable(value: unknown): value is PromiseLike<unknown> {
+  return typeof value === 'object' && value !== null && 'then' in value && typeof value.then === 'function';
+}

package/src/node/caching_aztec_node.ts

@@ -0,0 +1,407 @@
+import { Fr } from '@aztec/foundation/curves/bn254';
+import { allToCompletion } from '@aztec/foundation/promise';
+import type { AztecAddress } from '@aztec/stdlib/aztec-address';
+import { BlockHash, type BlockParameter, type DataInBlock } from '@aztec/stdlib/block';
+import type { BlockIncludeOptions } from '@aztec/stdlib/interfaces/client';
+import type { AztecNode } from '@aztec/stdlib/interfaces/server';
+import type {
+  LogResult,
+  LogsQueryBase,
+  PrivateLogsQuery,
+  PublicLogsQuery,
+  SiloedTag,
+  Tag,
+  TagQuery,
+} from '@aztec/stdlib/logs';
+import type { MerkleTreeId } from '@aztec/stdlib/trees';
+
+import { type BenchmarkedAztecNode, withRecording } from './benchmarked_node.js';
+
+/**
+ * An {@link AztecNode} wrapper that serves repeated reads from a cache it owns.
+ */
+export interface CachingAztecNode extends BenchmarkedAztecNode {
+  /** Clears the read cache. The owner calls this periodically to bound memory. Cached entries never go stale. */
+  wipeCache(): void;
+}
+
+/**
+ * Wraps `node`, serving repeated reads from a cache the wrapper owns. Consumers share cached reads by sharing the
+ * wrapper.
+ *
+ * Only hash-pinned reads are cached. A block hash names immutable content, so the same call can never correctly answer
+ * differently, and even an `undefined` witness or leaf index is kept as a non-membership fact. That makes the wrapper
+ * safe for any consumer, whatever its anchor block. Uncached methods, and reads that name a block any other way (see
+ * {@link hashReferenceOf}), pass straight through. A few entries below opt out even when hash-pinned, where the answer
+ * can still change, and the tag queries opt in only once the request also stops short of blocks yet to be produced
+ * (see {@link hasImmutableAnswer}).
+ *
+ * Wiping (see {@link CachingAztecNode.wipeCache}) bounds memory. Correctness does not depend on it.
+ */
+export function withCache(node: AztecNode): CachingAztecNode {
+  const cache = new AztecNodeCache();
+  // The recording wrapper sits below the cache, so a recording sees only the reads the node answered: a read served
+  // from the cache appears nowhere in it. Every read below goes to `source` rather than to `node`, and
+  // `startRecording` passes through this wrapper untouched.
+  const source = withRecording(node);
+
+  /** Runs `read` through `cache` when `block` pins a chain position by hash. Everything else goes to the node. */
+  const readCachedIfBlockIsHashPinned = <T>(
+    block: BlockParameter | undefined,
+    key: (blockHash: string) => string,
+    read: () => Promise<T>,
+    options?: { shouldCache?: (value: T) => boolean },
+  ): Promise<T> => {
+    const blockHash = hashReferenceOf(block);
+    return blockHash === undefined ? read() : cache.fetch(key(blockHash.toString()), read, options);
+  };
+
+  const cachedReads: { [K in keyof AztecNode]?: (...args: Parameters<AztecNode[K]>) => ReturnType<AztecNode[K]> } = {
+    getBlock: (block: BlockParameter, options?: BlockIncludeOptions) => {
+      if (options?.includeL1PublishInfo || options?.includeAttestations) {
+        // These payloads keep changing after a block is fixed (its L1 publication status, its incoming committee
+        // attestations), so a repeat can correctly differ; read straight from the node.
+        return source.getBlock(block, options);
+      }
+      return readCachedIfBlockIsHashPinned(
+        block,
+        blockHash => `block:${blockHash}:${keyPart(options)}`,
+        () => source.getBlock(block, options),
+        {
+          // PXE only asks for pinned blocks it has already observed, so an undefined answer means the node has not
+          // served the block yet, not that it never will.
+          shouldCache: response => response !== undefined,
+        },
+      );
+    },
+
+    getContract: (address: AztecAddress, referenceBlock?: BlockParameter) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `contract:${blockHash}:${address.toString()}`,
+        () => source.getContract(address, referenceBlock),
+      ),
+
+    getBlockHashMembershipWitness: (referenceBlock: BlockParameter, blockHash: BlockHash) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        referenceHash => `block-hash-membership-witness:${referenceHash}:${blockHash.toString()}`,
+        () => source.getBlockHashMembershipWitness(referenceBlock, blockHash),
+      ),
+
+    getPublicDataWitness: (referenceBlock: BlockParameter, leafSlot: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `public-data-witness:${blockHash}:${leafSlot.toString()}`,
+        () => source.getPublicDataWitness(referenceBlock, leafSlot),
+      ),
+
+    getNoteHashMembershipWitness: (referenceBlock: BlockParameter, noteHash: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `note-hash-membership-witness:${blockHash}:${noteHash.toString()}`,
+        () => source.getNoteHashMembershipWitness(referenceBlock, noteHash),
+      ),
+
+    getNullifierMembershipWitness: (referenceBlock: BlockParameter, nullifier: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `nullifier-membership-witness:${blockHash}:${nullifier.toString()}`,
+        () => source.getNullifierMembershipWitness(referenceBlock, nullifier),
+      ),
+
+    getLowNullifierMembershipWitness: (referenceBlock: BlockParameter, nullifier: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `low-nullifier-membership-witness:${blockHash}:${nullifier.toString()}`,
+        () => source.getLowNullifierMembershipWitness(referenceBlock, nullifier),
+      ),
+
+    getL1ToL2MessageMembershipWitness: (referenceBlock: BlockParameter, l1ToL2Message: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `l1-to-l2-message-membership-witness:${blockHash}:${l1ToL2Message.toString()}`,
+        () => source.getL1ToL2MessageMembershipWitness(referenceBlock, l1ToL2Message),
+      ),
+
+    getPublicStorageAt: (referenceBlock: BlockParameter, contractAddress: AztecAddress, storageSlot: Fr) =>
+      readCachedIfBlockIsHashPinned(
+        referenceBlock,
+        blockHash => `public-storage:${blockHash}:${contractAddress.toString()}:${storageSlot.toString()}`,
+        () => source.getPublicStorageAt(referenceBlock, contractAddress, storageSlot),
+      ),
+
+    findLeavesIndexes: (referenceBlock: BlockParameter, treeId: MerkleTreeId, leafValues: Fr[]) => {
+      // Cached per leaf: only leaves without a cached result are fetched, in a single batched node call. The per-leaf
+      // keys don't fit single-key readCachedIfBlockIsHashPinned, so the hash gate is inline.
+      const referenceHash = hashReferenceOf(referenceBlock);
+      if (referenceHash === undefined) {
+        return source.findLeavesIndexes(referenceBlock, treeId, leafValues);
+      }
+      return readBatchedPerKey<DataInBlock<bigint> | undefined>(
+        cache,
+        leafValues.map(leaf => `leaf-index:${referenceHash.toString()}:${treeId}:${leaf.toString()}`),
+        missing =>
+          source.findLeavesIndexes(
+            referenceBlock,
+            treeId,
+            missing.map(i => leafValues[i]),
+          ),
+        { method: 'findLeavesIndexes', requested: 'leaves' },
+      );
+    },
+
+    getPrivateLogsByTags: (query: PrivateLogsQuery) => {
+      if (!hasImmutableAnswer(query)) {
+        return source.getPrivateLogsByTags(query);
+      }
+      const keyPrefix = `private-logs:${privateLogsQueryKey(query)}`;
+      return readBatchedPerKey<LogResult[]>(
+        cache,
+        query.tags.map(tag => `${keyPrefix}:${tagQueryKey(tag)}`),
+        missing => source.getPrivateLogsByTags({ ...query, tags: missing.map(i => query.tags[i]) }),
+        { method: 'getPrivateLogsByTags', requested: 'tags' },
+      );
+    },
+
+    getPublicLogsByTags: (query: PublicLogsQuery) => {
+      if (!hasImmutableAnswer(query)) {
+        return source.getPublicLogsByTags(query);
+      }
+      const keyPrefix = `public-logs:${publicLogsQueryKey(query)}`;
+      return readBatchedPerKey<LogResult[]>(
+        cache,
+        query.tags.map(tag => `${keyPrefix}:${tagQueryKey(tag)}`),
+        missing => source.getPublicLogsByTags({ ...query, tags: missing.map(i => query.tags[i]) }),
+        { method: 'getPublicLogsByTags', requested: 'tags' },
+      );
+    },
+  };
+
+  return new Proxy(source, {
+    get(target, prop) {
+      if (prop === 'wipeCache') {
+        return () => cache.wipe();
+      }
+      if (Object.hasOwn(cachedReads, prop)) {
+        return cachedReads[prop as keyof AztecNode];
+      }
+      const value = Reflect.get(target, prop);
+      return typeof value === 'function' ? value.bind(target) : value;
+    },
+  }) as CachingAztecNode;
+}
+
+/**
+ * Store of node read promises keyed by method and arguments.
+ *
+ * Holds no node reference of its own: it only memoizes promises handed to it by the {@link withCache} wrapper that
+ * owns it. Rejected promises are evicted so callers can retry. A settled value is kept unless the caller's
+ * `shouldCache` turns it away.
+ */
+class AztecNodeCache {
+  private readonly cache = new Map<string, Promise<unknown>>();
+
+  /**
+   * Returns the cached promise for `key`, or runs `read` and caches the in-flight promise.
+   *
+   * `shouldCache` runs once the promise settles and answers whether its value is kept. Every settled value is kept by
+   * default. Callers waiting on a value that is then turned away still share the one in-flight read, so declining to
+   * keep a value costs nothing on concurrent reads.
+   */
+  public fetch<T>(
+    key: string,
+    read: () => Promise<T>,
+    { shouldCache }: { shouldCache?: (value: T) => boolean } = {},
+  ): Promise<T> {
+    return this.#get<T>(key) ?? this.#set(key, read(), shouldCache);
+  }
+
+  /**
+   * The cached promise for each of `keys`, `undefined` where there is none. For callers that batch their own misses
+   * rather than reading one key.
+   */
+  public peekAll<T>(keys: string[]): (Promise<T> | undefined)[] {
+    return keys.map(key => this.#get<T>(key));
+  }
+
+  /**
+   * Clears the store.
+   */
+  public wipe(): void {
+    this.cache.clear();
+  }
+
+  #get<T>(key: string): Promise<T> | undefined {
+    return this.cache.get(key) as Promise<T> | undefined;
+  }
+
+  #set<T>(key: string, promise: Promise<T>, shouldCache?: (value: T) => boolean): Promise<T> {
+    const evict = () => {
+      if (this.cache.get(key) === promise) {
+        this.cache.delete(key);
+      }
+    };
+    // A rejected read is evicted so callers can retry; a `shouldCache` that throws is treated the same way. The
+    // trailing catch keeps a throwing predicate from surfacing as an unhandled rejection on this detached promise.
+    promise
+      .then(value => {
+        if (shouldCache && !shouldCache(value)) {
+          evict();
+        }
+      }, evict)
+      .catch(evict);
+    this.cache.set(key, promise);
+    return promise;
+  }
+}
+
+/**
+ * The block hash `block` pins to (a `BlockHash` or `{ hash }` reference), or `undefined` for every other way of naming
+ * a block.
+ *
+ * A number or a tag names a moving chain position, as does naming no block at all: a reorg can put a different block
+ * at the same number, and tags follow the growing chain. An `{ archive }` root is different, since it commits to every
+ * block hash up to its own block and so pins content as tightly as a hash. It stays uncached because no PXE read uses
+ * one: keeping an `undefined` answer as a non-membership fact is only sound because PXE pins blocks it has already
+ * seen, and that holds for the references PXE actually builds.
+ */
+function hashReferenceOf(block: BlockParameter | undefined): BlockHash | undefined {
+  if (block instanceof BlockHash) {
+    return block;
+  }
+  if (typeof block === 'object' && block !== null && 'hash' in block) {
+    return block.hash;
+  }
+  return undefined;
+}
+
+/**
+ * Answers one result per entry of `keys`, taking whatever the cache already holds and fetching the rest in a single
+ * batched call. Each fetched result is cached under its own key, so a later call that overlaps this one fetches only
+ * the entries it has not seen, and a key repeated within `keys` is asked of the node once.
+ *
+ * `fetchMissing` receives the indexes still missing, in order, and must answer one result per requested index.
+ * `shortResponse` names the method and what it requests, for the error raised when it does not.
+ */
+function readBatchedPerKey<T>(
+  cache: AztecNodeCache,
+  keys: string[],
+  fetchMissing: (missingIndexes: number[]) => Promise<T[]>,
+  shortResponse: { method: string; requested: string },
+): Promise<T[]> {
+  const cached = cache.peekAll<T>(keys);
+
+  // One request per missing key rather than per missing position, so a key repeated in `keys` seeds a single entry
+  // that every position carrying it then reads.
+  const requestIndexByKey = new Map<string, number>();
+  cached.forEach((result, index) => {
+    if (result === undefined && !requestIndexByKey.has(keys[index])) {
+      requestIndexByKey.set(keys[index], index);
+    }
+  });
+
+  const fetchedByKey = new Map<string, Promise<T>>();
+  if (requestIndexByKey.size > 0) {
+    const requestedIndexes = [...requestIndexByKey.values()];
+    const batch = fetchMissing(requestedIndexes).then(fetched => {
+      if (fetched.length !== requestedIndexes.length) {
+        // Each entry is filled from the batch response by position. If the node returned fewer results than we asked
+        // for, the positions past the end would read undefined, which the cache keeps as a genuine answer, so a
+        // truncated response would be cached as fact. Rejecting the whole batch evicts those entries instead, so a
+        // retry re-fetches every one of them.
+        throw new Error(
+          `${shortResponse.method} returned ${fetched.length} results for ${requestedIndexes.length} requested ` +
+            `${shortResponse.requested}`,
+        );
+      }
+      return fetched;
+    });
+    [...requestIndexByKey.keys()].forEach((key, batchIndex) => {
+      fetchedByKey.set(
+        key,
+        cache.fetch(key, () => batch.then(fetched => fetched[batchIndex])),
+      );
+    });
+  }
+
+  // Every position is answered: a cache hit, or the entry its key's request seeded just above.
+  return allToCompletion(keys.map((key, index) => cached[index] ?? fetchedByKey.get(key)!));
+}
+
+/**
+ * Whether a tag query's answer can no longer change, and so can be cached.
+ *
+ * Two things would let it change, and the query has to rule out both. `referenceBlock` names the chain the answer
+ * belongs to, so the call fails once that block is gone rather than answering from a chain that reorged. `toBlock`
+ * stops the query below the tip, so blocks yet to be produced cannot add logs to the answer. Where the query starts
+ * does not matter: an open lower end reaches only blocks that are already behind it, which no longer move. PXE's tag
+ * queries get both from the `getAll*LogsByTags` helpers, which take them from the anchor block the query is pinned to.
+ */
+function hasImmutableAnswer(query: LogsQueryBase): boolean {
+  return query.referenceBlock !== undefined && query.toBlock !== undefined;
+}
+
+/** Cache-key segment for the parts of a private tag query that every tag in it shares. */
+function privateLogsQueryKey(query: PrivateLogsQuery): string {
+  const parts: QueryKeyParts<PrivateLogsQuery> = { ...sharedQueryKeyParts(query) };
+  return Object.values(parts).map(keyPart).join(':');
+}
+
+/** Cache-key segment for the parts of a public tag query that every tag in it shares. */
+function publicLogsQueryKey(query: PublicLogsQuery): string {
+  const parts: QueryKeyParts<PublicLogsQuery> = {
+    ...sharedQueryKeyParts(query),
+    contractAddress: query.contractAddress,
+  };
+  return Object.values(parts).map(keyPart).join(':');
+}
+
+/**
+ * The fields of a query that a key has to cover, which is all of them but `tags`: each tag is keyed on its own by
+ * {@link tagQueryKey}.
+ *
+ * The key builders are total over this type, so a field added to a query fails to compile there rather than silently
+ * being left out of the key and colliding with a query that differs only in it.
+ */
+type QueryKeyParts<T extends LogsQueryBase> = { [K in keyof Required<Omit<T, 'tags'>>]: unknown };
+
+/** The key parts {@link PrivateLogsQuery} and {@link PublicLogsQuery} have in common. */
+function sharedQueryKeyParts(query: LogsQueryBase): QueryKeyParts<LogsQueryBase> {
+  return {
+    referenceBlock: query.referenceBlock,
+    fromBlock: query.fromBlock,
+    toBlock: query.toBlock,
+    txHash: query.txHash,
+    includeEffects: query.includeEffects,
+    limitPerTag: query.limitPerTag,
+  };
+}
+
+/**
+ * Cache-key segment for one tag entry, covering the pagination cursor as well as the tag: each page of a tag's stream
+ * is a read of its own, and a repeated query asks for the same pages in the same order.
+ */
+function tagQueryKey(entry: TagQuery<Tag | SiloedTag>): string {
+  return 'tag' in entry ? `${keyPart(entry.tag)}@${keyPart(entry.afterLog)}` : keyPart(entry);
+}
+
+/**
+ * Renders a call argument as a stable cache-key segment, via its own `toString` when it has one, else as JSON.
+ *
+ * Example: the options object `{ includeTransactions: true }` renders as `'{"includeTransactions":true}'`, and an
+ * absent optional argument as `'undefined'`.
+ */
+function keyPart(value: unknown): string {
+  if (['string', 'number', 'bigint', 'boolean'].includes(typeof value)) {
+    return String(value);
+  }
+  if (value && typeof value === 'object') {
+    const toString = (value as { toString?: () => string }).toString;
+    if (toString && toString !== Object.prototype.toString) {
+      return toString.call(value);
+    }
+    return JSON.stringify(value, (_key, nested) => (typeof nested === 'bigint' ? nested.toString() : nested));
+  }
+  return String(value);
+}