Fix for loading bank from snapshot

runtime/src/bank.rs

@@ -6829,6 +6829,9 @@ impl Bank {
             Arc::new(reserved_keys)
         };
 
+        // Update the cost-tracker's block limits.
+        self.update_block_limits();
+
         if new_feature_activations.contains(&feature_set::pico_inflation::id()) {
             *self.inflation.write().unwrap() = Inflation::pico();
             self.fee_rate_governor.burn_percent = 50; // 50% fee burn
@@ -6917,15 +6920,6 @@ impl Bank {
                 );
             }
         }
-
-        if new_feature_activations.contains(&feature_set::raise_block_limits_to_50m::id()) {
-            let (account_cost_limit, block_cost_limit, vote_cost_limit) = simd_0207_block_limits();
-            self.write_cost_tracker().unwrap().set_limits(
-                account_cost_limit,
-                block_cost_limit,
-                vote_cost_limit,
-            );
-        }
     }
 
     fn apply_updated_hashes_per_tick(&mut self, hashes_per_tick: u64) {
@@ -6937,6 +6931,27 @@ impl Bank {
         self.hashes_per_tick = Some(hashes_per_tick);
     }
 
+    fn update_block_limits(&self) {
+        let mut cost_tracker = self.write_cost_tracker().unwrap();
+
+        // For tests and benches, if the limit is set to maximum then skip
+        // setting feature-specific limits.
+        #[cfg(feature = "dev-context-only-utils")]
+        {
+            if cost_tracker.get_block_limit() == u64::MAX {
+                return;
+            }
+        }
+
+        if self
+            .feature_set
+            .is_active(&feature_set::raise_block_limits_to_50m::id())
+        {
+            let (account_cost_limit, block_cost_limit, vote_cost_limit) = simd_0207_block_limits();
+            cost_tracker.set_limits(account_cost_limit, block_cost_limit, vote_cost_limit);
+        }
+    }
+
     fn adjust_sysvar_balance_for_rent(&self, account: &mut AccountSharedData) {
         account.set_lamports(
             self.get_minimum_balance_for_rent_exemption(account.data().len())

runtime/src/bank/tests.rs

@@ -8002,9 +8002,10 @@ fn test_reserved_account_keys() {
 
 #[test]
 fn test_block_limits() {
-    let (bank0, _bank_forks) = create_simple_test_arc_bank(100_000);
+    let (bank0, bank_forks) = create_simple_test_arc_bank(100_000);
     let mut bank = Bank::new_from_parent(bank0, &Pubkey::default(), 1);
     bank.feature_set = Arc::new(FeatureSet::default());
+    let bank = Arc::new(bank);
 
     assert_eq!(
         bank.read_cost_tracker().unwrap().get_block_limit(),
@@ -8017,20 +8018,43 @@ fn test_block_limits() {
         &feature_set::raise_block_limits_to_50m::id(),
         &feature::create_account(&Feature::default(), 42),
     );
-    bank.apply_feature_activations(ApplyFeatureActivationsCaller::NewFromParent, true);
+    // warp to next epoch for feature activation
+    let bank = new_from_parent_next_epoch(bank.clone(), &bank_forks, 1);
     assert_eq!(
         bank.read_cost_tracker().unwrap().get_block_limit(),
         MAX_BLOCK_UNITS_SIMD_0207,
         "after activating the feature, bank should have new limit"
     );
 
     // Make sure the limits propagate to the child-bank.
-    let bank = Bank::new_from_parent(Arc::new(bank), &Pubkey::default(), 2);
+    let bank = new_from_parent_next_epoch(bank.clone(), &bank_forks, 1);
     assert_eq!(
         bank.read_cost_tracker().unwrap().get_block_limit(),
         MAX_BLOCK_UNITS_SIMD_0207,
         "child bank should have new limit"
     );
+
+    // If we manually set limits to some other value,
+    // they will get reset upon new_from_parent for new epoch.
+
+    bank.write_cost_tracker().unwrap().set_limits(1, 2, 0);
+    let bank = new_from_parent_next_epoch(bank.clone(), &bank_forks, 1);
+    assert_eq!(
+        bank.read_cost_tracker().unwrap().get_block_limit(),
+        MAX_BLOCK_UNITS_SIMD_0207,
+        "child bank ignores parent, uses feature-set to derive limits"
+    );
+
+    // Test that for tests/benches if set to u64::MAX, the limits will not reset on the epoch.
+    bank.write_cost_tracker()
+        .unwrap()
+        .set_limits(u64::MAX, u64::MAX, u64::MAX);
+    let bank = new_from_parent_next_epoch(bank.clone(), &bank_forks, 1);
+    assert_eq!(
+        bank.read_cost_tracker().unwrap().get_block_limit(),
+        u64::MAX,
+        "child bank should have the same limit"
+    );
 }
 
 #[test]