From b60724174dc2edb08e44e55f3574ba3c533e4bbe Mon Sep 17 00:00:00 2001
From: Eric Semeniuc <3838856+esemeniuc@users.noreply.github.com>
Date: Thu, 30 Nov 2023 18:56:20 +0100
Subject: [PATCH] [JIT-1713] Fix bundle's blockspace preallocation (#489)

---
 .../bundle_reserved_space_manager.rs          | 56 ++++++++++++++++++-
 1 file changed, 53 insertions(+), 3 deletions(-)

diff --git a/core/src/bundle_stage/bundle_reserved_space_manager.rs b/core/src/bundle_stage/bundle_reserved_space_manager.rs
index 42cb7adeb6..a6b858d3b2 100644
--- a/core/src/bundle_stage/bundle_reserved_space_manager.rs
+++ b/core/src/bundle_stage/bundle_reserved_space_manager.rs
@@ -56,7 +56,7 @@ impl BundleReservedSpaceManager {
 
     /// return true if the bank is still in the period where block_cost_limits is reduced
     pub fn is_in_reserved_tick_period(&self, bank: &Bank) -> bool {
-        bank.tick_height() < self.reserved_ticks
+        bank.tick_height() % bank.ticks_per_slot() < self.reserved_ticks
     }
 
     /// return the block_cost_limits as determined by the tick height of the bank
@@ -82,8 +82,10 @@ impl BundleReservedSpaceManager {
 mod tests {
     use {
         crate::bundle_stage::bundle_reserved_space_manager::BundleReservedSpaceManager,
-        solana_ledger::genesis_utils::create_genesis_config, solana_runtime::bank::Bank,
-        solana_sdk::hash::Hash, std::sync::Arc,
+        solana_ledger::genesis_utils::create_genesis_config,
+        solana_runtime::bank::Bank,
+        solana_sdk::{hash::Hash, pubkey::Pubkey},
+        std::sync::Arc,
     };
 
     #[test]
@@ -186,4 +188,52 @@ mod tests {
             block_cost_limits
         );
     }
+
+    #[test]
+    fn test_block_limits_after_first_slot() {
+        const BUNDLE_BLOCK_COST_LIMITS_RESERVATION: u64 = 100;
+        const RESERVED_TICKS: u64 = 5;
+        let genesis_config_info = create_genesis_config(100);
+        let bank = Arc::new(Bank::new_for_tests(&genesis_config_info.genesis_config));
+
+        for _ in 0..genesis_config_info.genesis_config.ticks_per_slot {
+            bank.register_tick(&Hash::default());
+        }
+        assert!(bank.is_complete());
+        bank.freeze();
+        assert_eq!(
+            bank.read_cost_tracker().unwrap().block_cost_limit(),
+            solana_cost_model::block_cost_limits::MAX_BLOCK_UNITS,
+        );
+
+        let bank1 = Arc::new(Bank::new_from_parent(bank.clone(), &Pubkey::default(), 1));
+        assert_eq!(bank1.slot(), 1);
+        assert_eq!(bank1.tick_height(), 64);
+        assert_eq!(bank1.max_tick_height(), 128);
+
+        // reserve space
+        let block_cost_limits = bank1.read_cost_tracker().unwrap().block_cost_limit();
+        let mut reserved_space = BundleReservedSpaceManager::new(
+            block_cost_limits,
+            BUNDLE_BLOCK_COST_LIMITS_RESERVATION,
+            RESERVED_TICKS,
+        );
+        reserved_space.tick(&bank1);
+
+        // wait for reservation to be over
+        (0..RESERVED_TICKS).for_each(|_| {
+            bank1.register_tick(&Hash::default());
+            assert_eq!(
+                bank1.read_cost_tracker().unwrap().block_cost_limit(),
+                block_cost_limits - BUNDLE_BLOCK_COST_LIMITS_RESERVATION
+            );
+        });
+        reserved_space.tick(&bank1);
+
+        // after reservation, revert back to normal limit
+        assert_eq!(
+            bank1.read_cost_tracker().unwrap().block_cost_limit(),
+            solana_cost_model::block_cost_limits::MAX_BLOCK_UNITS,
+        );
+    }
 }