Fix compute budget overflow on SlotHistory reading

programs/validator-history/src/instructions/copy_cluster_info.rs

@@ -6,7 +6,11 @@ use {
     },
     anchor_lang::{
         prelude::*,
-        solana_program::{clock::Clock, slot_history::Check},
+        solana_program::{
+            clock::Clock,
+            log::sol_log_compute_units,
+            slot_history::{Check, MAX_ENTRIES},
+        },
     },
 };
 
@@ -30,32 +34,72 @@ pub fn handle_copy_cluster_info(ctx: Context<CopyClusterInfo>) -> Result<()> {
     let slot_history: Box<SlotHistory> =
         Box::new(bincode::deserialize(&ctx.accounts.slot_history.try_borrow_data()?).unwrap());
 
+    sol_log_compute_units();
     let clock = Clock::get()?;
 
     let epoch = cast_epoch(clock.epoch)?;
 
     let epoch_start_timestamp = cast_epoch_start_timestamp(clock.epoch_start_timestamp);
 
-    // Sets the slot history for the previous epoch, since the current epoch is not yet complete.
-    if epoch > 0 {
-        cluster_history_account
-            .set_blocks(epoch - 1, blocks_in_epoch(epoch - 1, &slot_history)?)?;
-    }
-    cluster_history_account.set_blocks(epoch, blocks_in_epoch(epoch, &slot_history)?)?;
+    let epoch_schedule = EpochSchedule::get()?;
+
+    let slot_history = if epoch > 0 {
+        let slot_history_next_slot = slot_history.next_slot;
+        let start_slot = epoch_schedule.get_first_slot_in_epoch((epoch - 1).into());
+        let end_slot = epoch_schedule.get_last_slot_in_epoch((epoch - 1).into());
+        let (num_blocks, bitvec_inner) =
+            confirmed_blocks_in_epoch(start_slot, end_slot, slot_history)?;
+        // Sets the slot history for the previous epoch, since the total number of blocks in the epoch is now final
+        cluster_history_account.set_blocks(epoch - 1, num_blocks)?;
+        // recreates SlotHistory with same heap memory chunk
+        Box::new(SlotHistory {
+            bits: bitvec_inner.into(),
+            next_slot: slot_history_next_slot,
+        })
+    } else {
+        slot_history
+    };
+
+    let start_slot = epoch_schedule.get_first_slot_in_epoch(epoch.into());
+    let end_slot = epoch_schedule.get_last_slot_in_epoch(epoch.into());
+    let (num_blocks, _) = confirmed_blocks_in_epoch(start_slot, end_slot, slot_history)?;
+    cluster_history_account.set_blocks(epoch, num_blocks)?;
     cluster_history_account.set_epoch_start_timestamp(epoch, epoch_start_timestamp)?;
 
     cluster_history_account.cluster_history_last_update_slot = clock.slot;
 
     Ok(())
 }
 
-fn blocks_in_epoch(epoch: u16, slot_history: &SlotHistory) -> Result<u32> {
-    let epoch_schedule = EpochSchedule::get()?;
-    let start_slot = epoch_schedule.get_first_slot_in_epoch(epoch as u64);
-    let end_slot = epoch_schedule.get_last_slot_in_epoch(epoch as u64);
+const BITVEC_BLOCK_SIZE: u64 = 64;
+
+pub fn confirmed_blocks_in_epoch(
+    start_slot: u64,
+    end_slot: u64,
+    slot_history: Box<SlotHistory>,
+) -> Result<(u32, Box<[u64]>)> {
+    // The SlotHistory BitVec wraps a slice of "Blocks", usizes representing 64 slots each (different than solana blocks).
+    // Iterating through each slot uses too much compute, but we can count the bits of each u64 altogether efficiently
+    // with `.count_ones()`.
+    // The epoch is not guaranteed to align perfectly with Blocks so we need to count the first and last partial Blocks separately.
+    // The bitvec inner data is taken ownership of, then returned to be reused.
 
     let mut blocks_in_epoch = 0;
-    for i in start_slot..=end_slot {
+
+    let first_full_block_slot = if (start_slot % MAX_ENTRIES) % BITVEC_BLOCK_SIZE == 0 {
+        start_slot
+    } else {
+        start_slot + (64 - (start_slot % MAX_ENTRIES) % BITVEC_BLOCK_SIZE)
+    };
+
+    let last_full_block_slot = if (end_slot % MAX_ENTRIES) % BITVEC_BLOCK_SIZE == 0 {
+        end_slot
+    } else {
+        end_slot - (end_slot % MAX_ENTRIES) % BITVEC_BLOCK_SIZE
+    };
+
+    // First and last slots, in partial blocks
+    for i in (start_slot..first_full_block_slot).chain(last_full_block_slot..=end_slot) {
         match slot_history.check(i) {
             Check::Found => {
                 blocks_in_epoch += 1;
@@ -72,5 +116,12 @@ fn blocks_in_epoch(epoch: u16, slot_history: &SlotHistory) -> Result<u32> {
         };
     }
 
-    Ok(blocks_in_epoch)
+    let inner_bitvec = slot_history.bits.into_boxed_slice();
+
+    for i in (first_full_block_slot..last_full_block_slot).step_by(BITVEC_BLOCK_SIZE as usize) {
+        let block_index = (i % MAX_ENTRIES) / BITVEC_BLOCK_SIZE;
+        blocks_in_epoch += inner_bitvec[block_index as usize].count_ones();
+    }
+
+    Ok((blocks_in_epoch, inner_bitvec))
 }

programs/validator-history/src/lib.rs

@@ -14,11 +14,11 @@ pub use state::*;
 
 cfg_if::cfg_if! {
     if #[cfg(feature = "mainnet-beta")] {
-        declare_id!("HistoryJTGbKQD2mRgLZ3XhqHnN811Qpez8X9kCcGHoa");
+        declare_id!("HisTBTgDnsdxfMp3m63fgKxCx9xVQE17MhA9BWRdrAP");
     } else if #[cfg(feature = "testnet")] {
         declare_id!("HisTBTgDnsdxfMp3m63fgKxCx9xVQE17MhA9BWRdrAP");
     } else {
-        declare_id!("HistoryJTGbKQD2mRgLZ3XhqHnN811Qpez8X9kCcGHoa");
+        declare_id!("HisTBTgDnsdxfMp3m63fgKxCx9xVQE17MhA9BWRdrAP");
     }
 }
 

programs/validator-history/src/state.rs

@@ -785,9 +785,23 @@ impl ClusterHistory {
     // Sets total blocks for the target epoch
     pub fn set_blocks(&mut self, epoch: u16, blocks_in_epoch: u32) -> Result<()> {
         if let Some(entry) = self.history.last_mut() {
-            if entry.epoch == epoch {
-                entry.total_blocks = blocks_in_epoch;
-                return Ok(());
+            match entry.epoch.cmp(&epoch) {
+                Ordering::Equal => {
+                    entry.total_blocks = blocks_in_epoch;
+                    return Ok(());
+                }
+                Ordering::Greater => {
+                    if let Some(entry) = self
+                        .history
+                        .arr_mut()
+                        .iter_mut()
+                        .find(|entry| entry.epoch == epoch)
+                    {
+                        entry.total_blocks = blocks_in_epoch;
+                    }
+                    return Ok(());
+                }
+                Ordering::Less => {}
             }
         }
         let entry = ClusterHistoryEntry {
@@ -799,11 +813,13 @@ impl ClusterHistory {
 
         Ok(())
     }
+
     pub fn set_epoch_start_timestamp(
         &mut self,
         epoch: u16,
         epoch_start_timestamp: u64,
     ) -> Result<()> {
+        // Always called after `set_blocks` so we can assume the entry for this epoch exists
         if let Some(entry) = self.history.last_mut() {
             if entry.epoch == epoch {
                 entry.epoch_start_timestamp = epoch_start_timestamp;

tests/Cargo.toml

@@ -9,6 +9,7 @@ authors = ["Jito Foundation <[email protected]>"]
 [dependencies]
 anchor-lang = "0.28.0"
 bincode = "1.3.3"
+bv = "0.11.1"
 bytemuck = { version = "1.13.1", features = ["derive", "min_const_generics"] }
 cfg-if = "1.0.0"
 ed25519-dalek = "1.0.1"
@@ -23,5 +24,5 @@ solana-gossip = "1.16"
 solana-program-test = "1.16"
 solana-sdk = "1.16"
 solana-version = "1.16"
-validator-history = { path = "../programs/validator-history" }
+validator-history = { features = ["no-entrypoint"], path = "../programs/validator-history" }
 validator-history-vote-state = { path = "../utils/vote-state" }

tests/tests/test_cluster_history.rs

@@ -6,11 +6,11 @@ use {
     },
     solana_program_test::*,
     solana_sdk::{
-        clock::Clock, compute_budget::ComputeBudgetInstruction, signer::Signer,
-        transaction::Transaction,
+        clock::Clock, compute_budget::ComputeBudgetInstruction, epoch_schedule::EpochSchedule,
+        signer::Signer, transaction::Transaction,
     },
     tests::fixtures::TestFixture,
-    validator_history::ClusterHistory,
+    validator_history::{confirmed_blocks_in_epoch, ClusterHistory},
 };
 
 const MS_PER_SLOT: u64 = 400;
@@ -27,7 +27,7 @@ fn create_copy_cluster_history_transaction(fixture: &TestFixture) -> Transaction
         .to_account_metas(None),
     };
     let heap_request_ix = ComputeBudgetInstruction::request_heap_frame(256 * 1024);
-    let compute_budget_ix = ComputeBudgetInstruction::set_compute_unit_limit(300_000);
+    let compute_budget_ix = ComputeBudgetInstruction::set_compute_unit_limit(1_000_000);
 
     Transaction::new_signed_with_payer(
         &[heap_request_ix, compute_budget_ix, instruction],
@@ -120,3 +120,70 @@ async fn test_start_epoch_timestamp() {
         account.history.arr[1].epoch_start_timestamp - dif_slots * MS_PER_SLOT
     );
 }
+
+#[tokio::test]
+async fn test_cluster_history_compute_limit() {
+    // Initialize
+    let fixture = TestFixture::new().await;
+    let ctx = &fixture.ctx;
+    fixture.initialize_config().await;
+    fixture.initialize_cluster_history_account().await;
+
+    // Set EpochSchedule with 432,000 slots
+    let epoch_schedule = EpochSchedule::default();
+    ctx.borrow_mut().set_sysvar(&epoch_schedule);
+
+    fixture
+        .advance_num_epochs(epoch_schedule.first_normal_epoch + 1)
+        .await;
+    fixture
+        .advance_clock(epoch_schedule.first_normal_epoch + 1, 400)
+        .await;
+
+    let clock: Clock = ctx
+        .borrow_mut()
+        .banks_client
+        .get_sysvar()
+        .await
+        .expect("clock");
+    println!("epoch: {}", clock.epoch);
+
+    // Set SlotHistory sysvar
+    let mut slot_history = SlotHistory::default();
+    for i in epoch_schedule.get_first_slot_in_epoch(clock.epoch - 1)
+        ..=epoch_schedule.get_last_slot_in_epoch(clock.epoch)
+    {
+        if i % 2 == 0 {
+            slot_history.add(i);
+        }
+    }
+
+    ctx.borrow_mut().set_sysvar(&slot_history);
+
+    // Submit instruction
+    let transaction = create_copy_cluster_history_transaction(&fixture);
+    fixture.submit_transaction_assert_success(transaction).await;
+
+    let account: ClusterHistory = fixture
+        .load_and_deserialize(&fixture.cluster_history_account)
+        .await;
+
+    assert!(account.history.arr[0].epoch as u64 == clock.epoch - 1);
+    assert!(account.history.arr[0].total_blocks == 216_000);
+    assert!(account.history.arr[1].epoch as u64 == clock.epoch);
+    assert!(account.history.arr[1].total_blocks == 216_000);
+}
+
+// Non-fixture test to ensure that the SlotHistory partial slot logic works
+#[test]
+fn test_confirmed_blocks_in_epoch_partial_blocks() {
+    let mut slot_history = SlotHistory::default();
+    for i in 50..=149 {
+        slot_history.add(i);
+    }
+    // First partial block: 50 -> 64
+    // Full block: 64 -> 127
+    // Last partial block: 128 -> 149
+    let (num_blocks, _) = confirmed_blocks_in_epoch(50, 149, Box::new(slot_history)).unwrap();
+    assert_eq!(num_blocks, 100);
+}