runtime: use leader schedule epoch to serve sol_get_epoch_stake

programs/sbf/Cargo.toml

@@ -64,6 +64,8 @@ solana-svm-transaction = { path = "../../svm-transaction", version = "=2.1.0" }
 solana-timings = { path = "../../timings", version = "=2.1.0" }
 solana-transaction-status = { path = "../../transaction-status", version = "=2.1.0" }
 solana-type-overrides = { path = "../../type-overrides", version = "=2.1.0" }
+solana-vote = { path = "../../vote", version = "=2.1.0" }
+solana-vote-program = { path = "../../programs/vote", version = "=2.1.0" }
 agave-validator = { path = "../../validator", version = "=2.1.0" }
 solana-zk-token-sdk = { path = "../../zk-token-sdk", version = "=2.1.0" }
 solana_rbpf = "=0.8.5"
@@ -128,6 +130,8 @@ solana-svm-transaction = { workspace = true }
 solana-timings = { workspace = true }
 solana-transaction-status = { workspace = true }
 solana-type-overrides = { workspace = true }
+solana-vote = { workspace = true }
+solana-vote-program = { workspace = true }
 solana_rbpf = { workspace = true }
 
 [[bench]]
@@ -186,6 +190,7 @@ members = [
     "rust/simulation",
     "rust/spoof1",
     "rust/spoof1_system",
+    "rust/syscall-get-epoch-stake",
     "rust/sysvar",
     "rust/upgradeable",
     "rust/upgraded",

programs/sbf/rust/syscall-get-epoch-stake/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "solana-sbf-syscall-get-epoch-stake"
+version = { workspace = true }
+description = { workspace = true }
+authors = { workspace = true }
+repository = { workspace = true }
+homepage = { workspace = true }
+license = { workspace = true }
+edition = { workspace = true }
+
+[dependencies]
+solana-program = { workspace = true }
+
+[lib]
+crate-type = ["cdylib"]

programs/sbf/rust/syscall-get-epoch-stake/src/lib.rs

@@ -0,0 +1,39 @@
+//! Example Rust-based SBF program that tests the `sol_get_epoch_stake`
+//! syscall.
+
+extern crate solana_program;
+use solana_program::{
+    account_info::AccountInfo,
+    entrypoint::ProgramResult,
+    epoch_stake::{get_epoch_stake_for_vote_account, get_epoch_total_stake},
+    msg,
+    program::set_return_data,
+    pubkey::Pubkey,
+};
+
+solana_program::entrypoint_no_alloc!(process_instruction);
+pub fn process_instruction(
+    _program_id: &Pubkey,
+    accounts: &[AccountInfo],
+    _instruction_data: &[u8],
+) -> ProgramResult {
+    // Total stake.
+    let total_stake = get_epoch_total_stake();
+    assert_ne!(total_stake, 0);
+    msg!("Total Stake: {}", total_stake);
+
+    // Vote accounts.
+    let check_vote_account_stake = |i: usize| {
+        let vote_address = accounts[i].key;
+        let vote_stake = get_epoch_stake_for_vote_account(vote_address);
+        assert_ne!(vote_stake, 0);
+        msg!("Vote Stake for account {}: {}", i, vote_stake);
+    };
+    check_vote_account_stake(0);
+    check_vote_account_stake(1);
+
+    // For good measure, set the return data to total stake.
+    set_return_data(&total_stake.to_le_bytes());
+
+    Ok(())
+}

programs/sbf/tests/syscall_get_epoch_stake.rs

@@ -0,0 +1,102 @@
+#![cfg(feature = "sbf_rust")]
+
+use {
+    solana_runtime::{
+        bank::Bank,
+        bank_client::BankClient,
+        epoch_stakes::EpochStakes,
+        genesis_utils::{
+            create_genesis_config_with_vote_accounts, GenesisConfigInfo, ValidatorVoteKeypairs,
+        },
+        loader_utils::load_upgradeable_program_and_advance_slot,
+    },
+    solana_sdk::{
+        instruction::{AccountMeta, Instruction},
+        message::Message,
+        signature::{Keypair, Signer},
+        transaction::{SanitizedTransaction, Transaction},
+    },
+    solana_vote::vote_account::VoteAccount,
+    solana_vote_program::vote_state::create_account_with_authorized,
+    std::collections::HashMap,
+};
+
+#[test]
+fn test_syscall_get_epoch_stake() {
+    solana_logger::setup();
+
+    // Two vote accounts with stake.
+    let stakes = vec![100_000_000, 500_000_000];
+    let voting_keypairs = vec![
+        ValidatorVoteKeypairs::new_rand(),
+        ValidatorVoteKeypairs::new_rand(),
+    ];
+    let total_stake: u64 = stakes.iter().sum();
+
+    let GenesisConfigInfo {
+        genesis_config,
+        mint_keypair,
+        ..
+    } = create_genesis_config_with_vote_accounts(1_000_000_000, &voting_keypairs, stakes.clone());
+
+    let mut bank = Bank::new_for_tests(&genesis_config);
+
+    // Intentionally overwrite the bank epoch with no stake, to ensure the
+    // syscall gets the _current_ epoch stake based on the leader schedule
+    // (N + 1).
+    let epoch_stakes_epoch_0 = EpochStakes::new_for_tests(
+        voting_keypairs
+            .iter()
+            .map(|keypair| {
+                let node_id = keypair.node_keypair.pubkey();
+                let authorized_voter = keypair.vote_keypair.pubkey();
+                let vote_account = VoteAccount::try_from(create_account_with_authorized(
+                    &node_id,
+                    &authorized_voter,
+                    &node_id,
+                    0,
+                    100,
+                ))
+                .unwrap();
+                (authorized_voter, (0, vote_account)) // No stake.
+            })
+            .collect::<HashMap<_, _>>(),
+        0, // Leader schedule epoch 0
+    );
+    bank.set_epoch_stakes_for_test(0, epoch_stakes_epoch_0);
+
+    let (bank, bank_forks) = bank.wrap_with_bank_forks_for_tests();
+    let mut bank_client = BankClient::new_shared(bank);
+
+    let authority_keypair = Keypair::new();
+    let (bank, program_id) = load_upgradeable_program_and_advance_slot(
+        &mut bank_client,
+        bank_forks.as_ref(),
+        &mint_keypair,
+        &authority_keypair,
+        "solana_sbf_syscall_get_epoch_stake",
+    );
+    bank.freeze();
+
+    let instruction = Instruction::new_with_bytes(
+        program_id,
+        &[],
+        vec![
+            AccountMeta::new_readonly(voting_keypairs[0].vote_keypair.pubkey(), false),
+            AccountMeta::new_readonly(voting_keypairs[1].vote_keypair.pubkey(), false),
+        ],
+    );
+
+    let blockhash = bank.last_blockhash();
+    let message = Message::new(&[instruction], Some(&mint_keypair.pubkey()));
+    let transaction = Transaction::new(&[&mint_keypair], message, blockhash);
+    let sanitized_tx = SanitizedTransaction::from_transaction_for_tests(transaction);
+
+    let result = bank.simulate_transaction(&sanitized_tx, false);
+
+    assert!(result.result.is_ok());
+
+    let return_data_le_bytes: [u8; 8] = result.return_data.unwrap().data[0..8].try_into().unwrap();
+    let total_stake_from_return_data = u64::from_le_bytes(return_data_le_bytes);
+    assert_eq!(total_stake_from_return_data, total_stake);
+}

runtime/src/bank.rs

@@ -3673,8 +3673,8 @@ impl Bank {
             RentCollectorWithMetrics::new(self.rent_collector.clone());
         let processing_environment = TransactionProcessingEnvironment {
             blockhash,
-            epoch_total_stake: self.epoch_total_stake(self.epoch()),
-            epoch_vote_accounts: self.epoch_vote_accounts(self.epoch()),
+            epoch_total_stake: Some(self.get_current_epoch_total_stake()),
+            epoch_vote_accounts: Some(self.get_current_epoch_vote_accounts()),
             feature_set: Arc::clone(&self.feature_set),
             fee_structure: Some(&self.fee_structure),
             lamports_per_signature,
@@ -6302,13 +6302,27 @@ impl Bank {
             .map(|epoch_stakes| epoch_stakes.total_stake())
     }
 
+    /// Get the total epoch stake for the current Bank::epoch
+    pub fn get_current_epoch_total_stake(&self) -> u64 {
+        self.current_epoch_stakes().total_stake()
+    }
+
     /// vote accounts for the specific epoch along with the stake
     ///   attributed to each account
     pub fn epoch_vote_accounts(&self, epoch: Epoch) -> Option<&VoteAccountsHashMap> {
         let epoch_stakes = self.epoch_stakes.get(&epoch)?.stakes();
         Some(epoch_stakes.vote_accounts().as_ref())
     }
 
+    /// Get the vote accounts along with the stake attributed to each account
+    /// for the current Bank::epoch
+    pub fn get_current_epoch_vote_accounts(&self) -> &VoteAccountsHashMap {
+        self.current_epoch_stakes()
+            .stakes()
+            .vote_accounts()
+            .as_ref()
+    }
+
     /// Get the fixed authorized voter for the given vote account for the
     /// current epoch
     pub fn epoch_authorized_voter(&self, vote_account: &Pubkey) -> Option<&Pubkey> {