Fix loading expected priority operation in state keeper

core/bin/zksync_core/src/state_keeper/init_params.rs

@@ -52,8 +52,7 @@ impl ZkSyncStateInitParams {
     ) -> Self {
         let (last_block_number, tree, acc_id_by_addr) = Self::load_account_tree(storage).await;
 
-        let unprocessed_priority_op =
-            Self::unprocessed_priority_op_id(storage, last_block_number).await;
+        let unprocessed_priority_op = Self::unprocessed_priority_op_id(storage).await;
         let nfts = Self::load_nft_tokens(storage, last_block_number).await;
 
         let root_hash_jobs = Self::load_root_hash_jobs(storage).await;
@@ -193,17 +192,12 @@ impl ZkSyncStateInitParams {
             .collect()
     }
 
-    async fn unprocessed_priority_op_id(
-        storage: &mut zksync_storage::StorageProcessor<'_>,
-        block_number: BlockNumber,
-    ) -> u64 {
+    async fn unprocessed_priority_op_id(storage: &mut zksync_storage::StorageProcessor<'_>) -> u64 {
         storage
             .chain()
             .block_schema()
-            .get_block(block_number)
+            .next_expected_serial_id()
             .await
             .expect("Unable to load the last block to get unprocessed priority operation")
-            .map(|block| block.processed_priority_ops.1)
-            .unwrap_or(0)
     }
 }

core/bin/zksync_core/src/state_keeper/mod.rs

@@ -105,9 +105,14 @@ impl ZkSyncStateKeeper {
         for job in &initial_state.root_hash_jobs {
             // Ensure that all jobs are sequential and there are no gaps.
             assert_eq!(
-                job.block, last_block + 1,
-                "Unexpected incomplete block number. Started from block {}, got unexpected block {} instead of expected {}, root hash jobs queue: {:?}",
-                initial_state.last_block_number, job.block, last_block + 1, &initial_state.root_hash_jobs
+                job.block,
+                last_block + 1,
+                "Unexpected incomplete block number. Started from block {}, \
+                got unexpected block {} instead of expected {}, root hash jobs queue: {:?}",
+                initial_state.last_block_number,
+                job.block,
+                last_block + 1,
+                &initial_state.root_hash_jobs
             );
             last_block = job.block;
 

core/lib/storage/src/chain/block/mod.rs

@@ -1033,6 +1033,27 @@ impl<'a, 'c> BlockSchema<'a, 'c> {
         Ok(())
     }
 
+    /// Returns the ID of the next expected priority operation.
+    /// Performs a lookup in both incomplete and complete block tables.
+    pub async fn next_expected_serial_id(&mut self) -> QueryResult<u64> {
+        let start = Instant::now();
+
+        let next_expected_serial_id = sqlx::query!(
+            "SELECT GREATEST(
+                (SELECT MAX(unprocessed_prior_op_after) FROM incomplete_blocks),
+                (SELECT MAX(unprocessed_prior_op_after) FROM blocks)
+            )",
+        )
+        .fetch_one(self.0.conn())
+        .await?
+        .greatest
+        .map(|val| val as u64)
+        .unwrap_or_default();
+
+        metrics::histogram!("sql.chain.block.next_expected_serial_id", start.elapsed());
+        Ok(next_expected_serial_id)
+    }
+
     /// Returns the range of existing incomplete blocks.
     ///
     /// Returned range is *inclusive*, meaning that both returned blocks (if they were returned)

core/lib/storage/src/test_data.rs

@@ -197,7 +197,7 @@ pub fn gen_sample_incomplete_block(
         block_number,
         fee_account: AccountId(0),
         block_transactions: txs,
-        processed_priority_ops: (0, 0),
+        processed_priority_ops: (0, 1),
         block_chunks_size,
         commit_gas_limit: 1_000_000.into(),
         verify_gas_limit: 1_500_000.into(),

core/lib/storage/src/tests/chain/block.rs

@@ -1417,6 +1417,13 @@ async fn test_incomplete_block_logic(mut storage: StorageProcessor<'_>) -> Query
         "Pending block should be saved"
     );
 
+    // Pending block doesn't change next expected serial ID.
+    assert_eq!(
+        schema.next_expected_serial_id().await?,
+        0,
+        "There should be no serial ID"
+    );
+
     schema.save_incomplete_block(&incomplete_block).await?;
 
     // Pending block should be removed now.
@@ -1431,6 +1438,13 @@ async fn test_incomplete_block_logic(mut storage: StorageProcessor<'_>) -> Query
         "Block should not exist"
     );
 
+    // Next expected serial ID should be loaded from incomplete block.
+    assert_eq!(
+        schema.next_expected_serial_id().await?,
+        1,
+        "Serial ID should be loaded from the incomplete blocks table"
+    );
+
     // Finish the block and ensure it's created correctly.
     schema
         .finish_incomplete_block(complete_block.clone())
@@ -1446,5 +1460,12 @@ async fn test_incomplete_block_logic(mut storage: StorageProcessor<'_>) -> Query
         complete_block.new_root_hash
     );
 
+    // Next expected serial ID should be loaded from complete block.
+    assert_eq!(
+        schema.next_expected_serial_id().await?,
+        1,
+        "Serial ID should be loaded from the complete blocks table"
+    );
+
     Ok(())
 }