Improve wait_for_new_batch_params

Author: perekopskiy

core/node/state_keeper/src/io/mempool.rs

@@ -206,6 +206,21 @@ impl StateKeeperIO for MempoolIO {
                 .protocol_version_id_by_timestamp(timestamp)
                 .await
                 .context("Failed loading protocol version")?;
+            let previous_protocol_version = storage
+                .blocks_dal()
+                .pending_protocol_version()
+                .await
+                .context("Failed loading previous protocol version")?;
+            let batch_with_upgrade_tx = if previous_protocol_version != protocol_version {
+                storage
+                    .protocol_versions_dal()
+                    .get_protocol_upgrade_tx(protocol_version)
+                    .await
+                    .context("Failed loading protocol upgrade tx")?
+                    .is_some()
+            } else {
+                false
+            };
             drop(storage);
 
             // We create a new filter each time, since parameters may change and a previously
@@ -217,7 +232,8 @@ impl StateKeeperIO for MempoolIO {
             .await
             .context("failed creating L2 transaction filter")?;
 
-            if !self.mempool.has_next(&self.filter) {
+            // We do not populate mempool with upgrade tx so it should be checked separately.
+            if !batch_with_upgrade_tx && !self.mempool.has_next(&self.filter) {
                 tokio::time::sleep(self.delay_interval).await;
                 continue;
             }

core/node/state_keeper/src/io/tests/mod.rs

@@ -18,6 +18,8 @@ use zksync_types::{
     commitment::{L1BatchCommitmentMode, PubdataParams},
     fee_model::{BatchFeeInput, PubdataIndependentBatchFeeModelInput},
     l2::L2Tx,
+    protocol_upgrade::ProtocolUpgradeTx,
+    protocol_version::ProtocolSemanticVersion,
     AccountTreeId, Address, L1BatchNumber, L2BlockNumber, L2ChainId, ProtocolVersion,
     ProtocolVersionId, StorageKey, TransactionTimeRangeConstraint, H256, U256,
 };
@@ -848,6 +850,52 @@ async fn test_mempool_with_timestamp_assertion() {
     );
 }
 
+#[tokio::test]
+async fn test_batch_params_with_protocol_upgrade_tx() {
+    let connection_pool = ConnectionPool::<Core>::constrained_test_pool(2).await;
+    let tester = Tester::new(L1BatchCommitmentMode::Rollup);
+    // Genesis is needed for proper mempool initialization.
+    tester.genesis(&connection_pool).await;
+
+    let (mut mempool, _) = tester.create_test_mempool_io(connection_pool.clone()).await;
+    let (cursor, _) = mempool.initialize().await.unwrap();
+
+    // Check that new batch params are not returned when there is no tx to process.
+    let new_batch_params = mempool
+        .wait_for_new_batch_params(&cursor, Duration::from_millis(100))
+        .await
+        .unwrap();
+    assert!(new_batch_params.is_none());
+
+    // Insert protocol version with upgrade tx.
+    let protocol_upgrade_tx = ProtocolUpgradeTx {
+        execute: Default::default(),
+        common_data: Default::default(),
+        received_timestamp_ms: 0,
+    };
+    let version = ProtocolVersion {
+        version: ProtocolSemanticVersion {
+            minor: ProtocolVersionId::next(),
+            patch: 0.into(),
+        },
+        tx: Some(protocol_upgrade_tx),
+        ..Default::default()
+    };
+    connection_pool
+        .connection()
+        .await
+        .unwrap()
+        .protocol_versions_dal()
+        .save_protocol_version_with_tx(&version)
+        .await
+        .unwrap();
+    let new_batch_params = mempool
+        .wait_for_new_batch_params(&cursor, Duration::from_millis(100))
+        .await
+        .unwrap();
+    assert!(new_batch_params.is_some());
+}
+
 async fn insert_l2_transaction(storage: &mut Connection<'_, Core>, tx: &L2Tx) {
     storage
         .transactions_dal()