From a1f4a5b9ddc23065879864938b5fac97b968154b Mon Sep 17 00:00:00 2001
From: Santiago Palladino <santiago@aztecprotocol.com>
Date: Wed, 7 Aug 2024 10:10:43 -0300
Subject: [PATCH] Address comments by @alexghr

---
 .../l2_block_downloader.ts                    | 19 ++++++++++++-------
 .../server_world_state_synchronizer.ts        | 16 +++++++++-------
 2 files changed, 21 insertions(+), 14 deletions(-)

diff --git a/yarn-project/circuit-types/src/l2_block_downloader/l2_block_downloader.ts b/yarn-project/circuit-types/src/l2_block_downloader/l2_block_downloader.ts
index d8994b52ccd..8c43ce23925 100644
--- a/yarn-project/circuit-types/src/l2_block_downloader/l2_block_downloader.ts
+++ b/yarn-project/circuit-types/src/l2_block_downloader/l2_block_downloader.ts
@@ -68,10 +68,11 @@ export class L2BlockDownloader {
   /**
    * Repeatedly queries the block source and adds the received blocks to the block queue.
    * Stops when no further blocks are received.
-   * @param targetBlockNumber - Optional block number to stop at. If given, will download even unproven blocks to get to it.
+   * @param targetBlockNumber - Optional block number to stop at.
+   * @param proven - Optional override of the default "proven" setting.
    * @returns The total number of blocks added to the block queue.
    */
-  private async collectBlocks(targetBlockNumber?: number) {
+  private async collectBlocks(targetBlockNumber?: number, onlyProven?: boolean) {
     let totalBlocks = 0;
     while (true) {
       // If we have a target and have reached it, return
@@ -80,9 +81,9 @@ export class L2BlockDownloader {
         return totalBlocks;
       }
 
-      // If we have a target, then request at most the number of blocks to get to it, bypassing 'proven' restriction
-      const [proven, limit] =
-        targetBlockNumber !== undefined ? [false, Math.min(targetBlockNumber - this.from + 1, 10)] : [this.proven, 10];
+      // If we have a target, then request at most the number of blocks to get to it
+      const limit = targetBlockNumber !== undefined ? Math.min(targetBlockNumber - this.from + 1, 10) : 10;
+      const proven = onlyProven === undefined ? this.proven : onlyProven;
 
       // Hit the archiver for blocks
       const blocks = await this.l2BlockSource.getBlocks(this.from, limit, proven);
@@ -136,9 +137,13 @@ export class L2BlockDownloader {
 
   /**
    * Forces an immediate request for blocks.
+   * Repeatedly queries the block source and adds the received blocks to the block queue.
+   * Stops when no further blocks are received.
+   * @param targetBlockNumber - Optional block number to stop at.
+   * @param proven - Optional override of the default "proven" setting.
    * @returns A promise that fulfills once the poll is complete
    */
-  public pollImmediate(targetBlockNumber?: number): Promise<number> {
-    return this.jobQueue.put(() => this.collectBlocks(targetBlockNumber));
+  public pollImmediate(targetBlockNumber?: number, onlyProven?: boolean): Promise<number> {
+    return this.jobQueue.put(() => this.collectBlocks(targetBlockNumber, onlyProven));
   }
 }
diff --git a/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts b/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts
index d43bdd88634..fd99f6ba9ce 100644
--- a/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts
+++ b/yarn-project/world-state/src/synchronizer/server_world_state_synchronizer.ts
@@ -191,8 +191,8 @@ export class ServerWorldStateSynchronizer implements WorldStateSynchronizer {
       if (targetBlockNumber !== undefined && targetBlockNumber <= this.currentL2BlockNum) {
         return this.currentL2BlockNum;
       }
-      // Poll for more blocks
-      const numBlocks = await this.l2BlockDownloader.pollImmediate(targetBlockNumber);
+      // Poll for more blocks, requesting even unproven blocks.
+      const numBlocks = await this.l2BlockDownloader.pollImmediate(targetBlockNumber, false);
       this.log.debug(`Block download immediate poll yielded ${numBlocks} blocks`);
       if (numBlocks) {
         // More blocks were received, process them and go round again
@@ -213,11 +213,13 @@ export class ServerWorldStateSynchronizer implements WorldStateSynchronizer {
     targetBlockNumber: number,
     forkIncludeUncommitted: boolean,
   ): Promise<MerkleTreeOperationsFacade> {
-    await this.pause();
-    await this.syncImmediate(targetBlockNumber);
-    const fork = await this.getFork(forkIncludeUncommitted);
-    this.resume();
-    return fork;
+    try {
+      await this.pause();
+      await this.syncImmediate(targetBlockNumber);
+      return await this.getFork(forkIncludeUncommitted);
+    } finally {
+      this.resume();
+    }
   }
 
   /**