@aztec/world-state 0.69.1 → 0.70.0

package/src/native/native_world_state_instance.ts

@@ -11,7 +11,6 @@ import {
   PUBLIC_DATA_TREE_HEIGHT,
 } from '@aztec/circuits.js';
 import { createLogger } from '@aztec/foundation/log';
-import { SerialQueue } from '@aztec/foundation/queue';
 
 import assert from 'assert';
 import bindings from 'bindings';
@@ -25,8 +24,13 @@ import {
   TypedMessage,
   WorldStateMessageType,
   type WorldStateRequest,
+  type WorldStateRequestCategories,
   type WorldStateResponse,
+  isWithCanonical,
+  isWithForkId,
+  isWithRevision,
 } from './message.js';
+import { WorldStateOpsQueue } from './world_state_ops_queue.js';
 
 // small extension to pack an NodeJS Fr instance to a representation that the C++ code can understand
 // this only works for writes. Unpacking from C++ can't create Fr instances because the data is passed
@@ -49,7 +53,10 @@ const NATIVE_MODULE = bindings(NATIVE_LIBRARY_NAME);
 const MAX_WORLD_STATE_THREADS = +(process.env.HARDWARE_CONCURRENCY || '16');
 
 export interface NativeWorldStateInstance {
-  call<T extends WorldStateMessageType>(messageType: T, body: WorldStateRequest[T]): Promise<WorldStateResponse[T]>;
+  call<T extends WorldStateMessageType>(
+    messageType: T,
+    body: WorldStateRequest[T] & WorldStateRequestCategories,
+  ): Promise<WorldStateResponse[T]>;
 }
 
 /**
@@ -78,8 +85,8 @@ export class NativeWorldState implements NativeWorldStateInstance {
   /** The actual native instance */
   private instance: any;
 
-  /** Calls to the same instance are serialized */
-  private queue = new SerialQueue();
+  // We maintain a map of queue to fork
+  private queues = new Map<number, WorldStateOpsQueue>();
 
   /** Creates a new native WorldState instance */
   constructor(
@@ -109,7 +116,8 @@ export class NativeWorldState implements NativeWorldStateInstance {
       dbMapSizeKb,
       threads,
     );
-    this.queue.start();
+    // Manually create the queue for the canonical fork
+    this.queues.set(0, new WorldStateOpsQueue());
   }
 
   /**
@@ -120,25 +128,65 @@ export class NativeWorldState implements NativeWorldStateInstance {
    * @param errorHandler - A callback called on request error, executed on the job queue
    * @returns The response to the message
    */
-  public call<T extends WorldStateMessageType>(
+  public async call<T extends WorldStateMessageType>(
     messageType: T,
-    body: WorldStateRequest[T],
+    body: WorldStateRequest[T] & WorldStateRequestCategories,
     // allows for the pre-processing of responses on the job queue before being passed back
     responseHandler = (response: WorldStateResponse[T]): WorldStateResponse[T] => response,
     errorHandler = (_: string) => {},
   ): Promise<WorldStateResponse[T]> {
-    return this.queue.put(async () => {
-      assert.notEqual(messageType, WorldStateMessageType.CLOSE, 'Use close() to close the native instance');
-      assert.equal(this.open, true, 'Native instance is closed');
-      let response: WorldStateResponse[T];
-      try {
-        response = await this._sendMessage(messageType, body);
-      } catch (error: any) {
-        errorHandler(error.message);
-        throw error;
-      }
-      return responseHandler(response);
-    });
+    // Here we determine which fork the request is being executed against and whether it requires uncommitted data
+    // We use the fork Id to select the appropriate request queue and the uncommitted data flag to pass to the queue
+    let forkId = -1;
+    // We assume it includes uncommitted unless explicitly told otherwise
+    let committedOnly = false;
+
+    // Canonical requests ALWAYS go against the canonical fork
+    // These include things like block syncs/unwinds etc
+    // These requests don't contain a fork ID
+    if (isWithCanonical(body)) {
+      forkId = 0;
+    } else if (isWithForkId(body)) {
+      forkId = body.forkId;
+    } else if (isWithRevision(body)) {
+      forkId = body.revision.forkId;
+      committedOnly = body.revision.includeUncommitted === false;
+    } else {
+      const _: never = body;
+      throw new Error(`Unable to determine forkId for message=${WorldStateMessageType[messageType]}`);
+    }
+
+    // Get the queue or create a new one
+    let requestQueue = this.queues.get(forkId);
+    if (requestQueue === undefined) {
+      requestQueue = new WorldStateOpsQueue();
+      this.queues.set(forkId, requestQueue);
+    }
+
+    // Enqueue the request and wait for the response
+    const response = await requestQueue.execute(
+      async () => {
+        assert.notEqual(messageType, WorldStateMessageType.CLOSE, 'Use close() to close the native instance');
+        assert.equal(this.open, true, 'Native instance is closed');
+        let response: WorldStateResponse[T];
+        try {
+          response = await this._sendMessage(messageType, body);
+        } catch (error: any) {
+          errorHandler(error.message);
+          throw error;
+        }
+        return responseHandler(response);
+      },
+      messageType,
+      committedOnly,
+    );
+
+    // If the request was to delete the fork then we clean it up here
+    if (messageType === WorldStateMessageType.DELETE_FORK) {
+      await requestQueue.stop();
+      this.queues.delete(forkId);
+    }
+    return response;
   }
 
   /**
@@ -149,13 +197,21 @@ export class NativeWorldState implements NativeWorldStateInstance {
       return;
     }
     this.open = false;
-    await this._sendMessage(WorldStateMessageType.CLOSE, undefined);
-    await this.queue.end();
+    const queue = this.queues.get(0)!;
+
+    await queue.execute(
+      async () => {
+        await this._sendMessage(WorldStateMessageType.CLOSE, { canonical: true });
+      },
+      WorldStateMessageType.CLOSE,
+      false,
+    );
+    await queue.stop();
   }
 
   private async _sendMessage<T extends WorldStateMessageType>(
     messageType: T,
-    body: WorldStateRequest[T],
+    body: WorldStateRequest[T] & WorldStateRequestCategories,
   ): Promise<WorldStateResponse[T]> {
     const messageId = this.nextMessageId++;
     if (body) {

package/src/native/world_state_ops_queue.ts

@@ -0,0 +1,187 @@
+import { promiseWithResolvers } from '@aztec/foundation/promise';
+
+import { WorldStateMessageType } from './message.js';
+
+/**
+ * This is the implementation for queueing requests to the world state.
+ * Requests need to be queued for the world state to ensure that writes are correctly ordered
+ * and reads return the correct data.
+ * Due to the nature of the NAPI we can't really do this there.
+ *
+ * The rules for queueing are as follows:
+ *
+ * 1. Reads of committed state never need to be queued. LMDB uses MVCC to ensure readers see a consistent view of the DB.
+ * 2. Reads of uncommitted state can happen concurrently with other reads of uncommitted state on the same fork (or reads of committed state)
+ * 3. All writes require exclusive access to their respective fork
+ *
+ */
+
+type WorldStateOp = {
+  requestId: number;
+  mutating: boolean;
+  request: () => Promise<any>;
+  promise: PromiseWithResolvers<any>;
+};
+
+// These are the set of message types that implement mutating operations
+// Messages of these types require exclusive access to their given forks
+export const MUTATING_MSG_TYPES = new Set([
+  WorldStateMessageType.APPEND_LEAVES,
+  WorldStateMessageType.BATCH_INSERT,
+  WorldStateMessageType.SEQUENTIAL_INSERT,
+  WorldStateMessageType.UPDATE_ARCHIVE,
+  WorldStateMessageType.COMMIT,
+  WorldStateMessageType.ROLLBACK,
+  WorldStateMessageType.SYNC_BLOCK,
+  WorldStateMessageType.CREATE_FORK,
+  WorldStateMessageType.DELETE_FORK,
+  WorldStateMessageType.FINALISE_BLOCKS,
+  WorldStateMessageType.UNWIND_BLOCKS,
+  WorldStateMessageType.REMOVE_HISTORICAL_BLOCKS,
+]);
+
+// This class implements the per-fork operation queue
+export class WorldStateOpsQueue {
+  private requests: WorldStateOp[] = [];
+  private inFlightMutatingCount = 0;
+  private inFlightCount = 0;
+  private stopPromise?: Promise<void>;
+  private stopResolve?: () => void;
+  private requestId = 0;
+  private ops: Map<number, WorldStateOp> = new Map();
+
+  // The primary public api, this is where an operation is queued
+  // We return a promise that will ultimately be resolved/rejected with the response/error generated by the 'request' argument
+  public execute(request: () => Promise<any>, messageType: WorldStateMessageType, committedOnly: boolean) {
+    if (this.stopResolve !== undefined) {
+      throw new Error('Unable to send request to world state, queue already stopped');
+    }
+
+    const op: WorldStateOp = {
+      requestId: this.requestId++,
+      mutating: MUTATING_MSG_TYPES.has(messageType),
+      request,
+      promise: promiseWithResolvers(),
+    };
+    this.ops.set(op.requestId, op);
+
+    // Perform the appropriate action based upon the queueing rules
+    if (op.mutating) {
+      this.executeMutating(op);
+    } else if (committedOnly === false) {
+      this.executeNonMutatingUncommitted(op);
+    } else {
+      this.executeNonMutatingCommitted(op);
+    }
+    return op.promise.promise;
+  }
+
+  // Mutating requests need exclusive access
+  private executeMutating(op: WorldStateOp) {
+    // If nothing is in flight then we send the request immediately
+    // Otherwise add to the queue
+    if (this.inFlightCount === 0) {
+      this.sendEnqueuedRequest(op);
+    } else {
+      this.requests.push(op);
+    }
+  }
+
+  // Non mutating requests including uncommitted state
+  private executeNonMutatingUncommitted(op: WorldStateOp) {
+    // If there are no mutating requests in flight and there is nothing queued
+    // then send the request immediately
+    // If a mutating request is in flight then we must wait
+    // If a mutating request is not in flight but something is queued then it must be a mutating request
+    if (this.inFlightMutatingCount == 0 && this.requests.length == 0) {
+      this.sendEnqueuedRequest(op);
+    } else {
+      this.requests.push(op);
+    }
+  }
+
+  private executeNonMutatingCommitted(op: WorldStateOp) {
+    // This is a non-mutating request for committed data
+    // It can always be sent
+    op.request()
+      .then(op.promise.resolve, op.promise.reject)
+      .finally(() => {
+        this.ops.delete(op.requestId);
+      });
+  }
+
+  private checkAndEnqueue(completedOp: WorldStateOp) {
+    // As request has completed
+    // First we decrements the relevant in flight counters
+    if (completedOp.mutating) {
+      --this.inFlightMutatingCount;
+    }
+    --this.inFlightCount;
+
+    // If there are still requests in flight then do nothing further
+    if (this.inFlightCount != 0) {
+      return;
+    }
+
+    // No requests in flight, send next queued requests
+    // We loop and send:
+    // 1 mutating request if it is next in the queue
+    // As many non-mutating requests as we encounter until
+    // we exhaust the queue or we reach a mutating request
+    while (this.requests.length > 0) {
+      const next = this.requests[0];
+      if (next.mutating) {
+        if (this.inFlightCount == 0) {
+          // send the mutating request
+          this.requests.shift();
+          this.sendEnqueuedRequest(next);
+        }
+        // this request is mutating, we need to stop here
+        break;
+      } else {
+        // not mutating, send and go round again
+        this.requests.shift();
+        this.sendEnqueuedRequest(next);
+      }
+    }
+
+    // If the queue is empty, there is nothing in flight and we have been told to stop, then resolve the stop promise
+    if (this.inFlightCount == 0 && this.stopResolve !== undefined) {
+      this.stopResolve();
+    }
+  }
+
+  private sendEnqueuedRequest(op: WorldStateOp) {
+    // Here we increment the in flight counts before sending
+    ++this.inFlightCount;
+    if (op.mutating) {
+      ++this.inFlightMutatingCount;
+    }
+
+    // Make the request and pass the response/error through to the stored promise
+    op.request()
+      .then(op.promise.resolve, op.promise.reject)
+      .finally(() => {
+        this.checkAndEnqueue(op);
+        this.ops.delete(op.requestId);
+      });
+  }
+
+  public stop() {
+    // If there is already a stop promise then return it
+    if (this.stopPromise) {
+      return this.stopPromise;
+    }
+
+    // Otherwise create a new one and capture the resolve method
+    this.stopPromise = new Promise(resolve => {
+      this.stopResolve = resolve;
+    });
+
+    // If no outstanding requests then immediately resolve the promise
+    if (this.requests.length == 0 && this.inFlightCount == 0 && this.stopResolve !== undefined) {
+      this.stopResolve();
+    }
+    return this.stopPromise;
+  }
+}

package/src/test/utils.ts

@@ -3,7 +3,6 @@ import {
   MerkleTreeId,
   type MerkleTreeReadOperations,
   type MerkleTreeWriteOperations,
-  TxEffect,
 } from '@aztec/circuit-types';
 import {
   AppendOnlyTreeSnapshot,
@@ -21,15 +20,9 @@ export async function mockBlock(blockNum: number, size: number, fork: MerkleTree
   const l2Block = L2Block.random(blockNum, size);
   const l1ToL2Messages = Array(16).fill(0).map(Fr.random);
 
-  const paddedTxEffects = padArrayEnd(
-    l2Block.body.txEffects,
-    TxEffect.empty(),
-    l2Block.body.numberOfTxsIncludingPadded,
-  );
-
   // Sync the append only trees
   {
-    const noteHashesPadded = paddedTxEffects.flatMap(txEffect =>
+    const noteHashesPadded = l2Block.body.txEffects.flatMap(txEffect =>
       padArrayEnd(txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
     );
     await fork.appendLeaves(MerkleTreeId.NOTE_HASH_TREE, noteHashesPadded);
@@ -41,7 +34,7 @@ export async function mockBlock(blockNum: number, size: number, fork: MerkleTree
   // Sync the indexed trees
   {
     // We insert the public data tree leaves with one batch per tx to avoid updating the same key twice
-    for (const txEffect of paddedTxEffects) {
+    for (const txEffect of l2Block.body.txEffects) {
       await fork.batchInsert(
         MerkleTreeId.PUBLIC_DATA_TREE,
         txEffect.publicDataWrites.map(write => write.toBuffer()),

package/src/world-state-db/merkle_trees.ts

@@ -1,4 +1,4 @@
-import { type L2Block, MerkleTreeId, type SiblingPath, TxEffect } from '@aztec/circuit-types';
+import { type L2Block, MerkleTreeId, type SiblingPath } from '@aztec/circuit-types';
 import {
   type BatchInsertionResult,
   type IndexedTreeId,
@@ -647,17 +647,11 @@ export class MerkleTrees implements MerkleTreeAdminDatabase {
       this.log.verbose(`Block ${l2Block.number} is not ours, rolling back world state and committing state from chain`);
       await this.#rollback();
 
-      // We have to pad both the tx effects and the values within tx effects because that's how the trees are built
-      // by circuits.
-      const paddedTxEffects = padArrayEnd(
-        l2Block.body.txEffects,
-        TxEffect.empty(),
-        l2Block.body.numberOfTxsIncludingPadded,
-      );
+      // We have to pad the values within tx effects because that's how the trees are built by circuits.
 
       // Sync the append only trees
       {
-        const noteHashesPadded = paddedTxEffects.flatMap(txEffect =>
+        const noteHashesPadded = l2Block.body.txEffects.flatMap(txEffect =>
           padArrayEnd(txEffect.noteHashes, Fr.ZERO, MAX_NOTE_HASHES_PER_TX),
         );
         await this.#appendLeaves(MerkleTreeId.NOTE_HASH_TREE, noteHashesPadded);
@@ -668,7 +662,7 @@ export class MerkleTrees implements MerkleTreeAdminDatabase {
 
       // Sync the indexed trees
       {
-        const nullifiersPadded = paddedTxEffects.flatMap(txEffect =>
+        const nullifiersPadded = l2Block.body.txEffects.flatMap(txEffect =>
           padArrayEnd(txEffect.nullifiers, Fr.ZERO, MAX_NULLIFIERS_PER_TX),
         );
         await (this.trees[MerkleTreeId.NULLIFIER_TREE] as StandardIndexedTree).batchInsert(
@@ -679,7 +673,7 @@ export class MerkleTrees implements MerkleTreeAdminDatabase {
         const publicDataTree = this.trees[MerkleTreeId.PUBLIC_DATA_TREE] as StandardIndexedTree;
 
         // We insert the public data tree leaves with one batch per tx to avoid updating the same key twice
-        for (const txEffect of paddedTxEffects) {
+        for (const txEffect of l2Block.body.txEffects) {
           const publicDataWrites = padArrayEnd(
             txEffect.publicDataWrites,
             PublicDataWrite.empty(),