add hash_rewards_into_partitions (#31795)

* add sort_and_shuffle_partitioned_rewards

* break stake accounts into partitions by hashing

* fmt

* Update runtime/src/epoch_rewards_hasher.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* Update runtime/src/bank/tests.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* Update runtime/src/bank/tests.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* Update runtime/src/bank.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* fix build

* more build fixes

* use spihasher13 for reward partition

* update comments

* revert calc reward struct change to avoid abi changes for now

* cippy

* sort deps

* saturating math

* use copy traits on hasher
use parent block hash to seed rewards partition

* fmt

* use concurrent map to speed up reward partition

* clippy

* Update runtime/src/epoch_rewards_hasher.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* Update runtime/Cargo.toml

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* Update runtime/src/bank.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* rename

* address review comments

* cargo.lock

* Update runtime/src/bank.rs

Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>

* update test

* clippy

* comments

* move hash_address_to_partition to epoch_rewards_hasher.rs

* tests to nail down width of partitions

* rename parameter to match with comments

* hash_address takes &self

* clippy

* EpochRewardHasher -> EpochRewardsHasher

* revert "hash_address takes &self" and refactor

* stake_rewards takes by value

---------

Co-authored-by: HaoranYi <haoran.yi@gmail.com>
Co-authored-by: Trent Nelson <trent.a.b.nelson@gmail.com>
Co-authored-by: HaoranYi <haoran.yi@solana.com>

Cargo.lock

@@ -4838,6 +4838,12 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2a30f10c911c0355f80f1c2faa8096efc4a58cdf8590b954d5b395efa071c711"
 
+[[package]]
+name = "siphasher"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de"
+
 [[package]]
 name = "sized-chunks"
 version = "0.6.5"
@@ -6560,6 +6566,7 @@ dependencies = [
  "rustc_version 0.4.0",
  "serde",
  "serde_derive",
+ "siphasher",
  "solana-address-lookup-table-program",
  "solana-bpf-loader-program",
  "solana-bucket-map",

Cargo.toml

@@ -285,6 +285,7 @@ serde_with = { version = "2.3.3", default-features = false }
 sha2 = "0.10.6"
 sha3 = "0.10.4"
 signal-hook = "0.3.15"
+siphasher = "0.3.10"
 smpl_jwt = "0.7.1"
 socket2 = "0.4.9"
 soketto = "0.7"

programs/sbf/Cargo.lock

@@ -4368,6 +4368,12 @@ version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "2a30f10c911c0355f80f1c2faa8096efc4a58cdf8590b954d5b395efa071c711"
 
+[[package]]
+name = "siphasher"
+version = "0.3.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7bd3e3206899af3f8b12af284fafc038cc1dc2b41d1b89dd17297221c5d225de"
+
 [[package]]
 name = "sized-chunks"
 version = "0.6.5"
@@ -5439,6 +5445,7 @@ dependencies = [
  "rustc_version",
  "serde",
  "serde_derive",
+ "siphasher",
  "solana-address-lookup-table-program",
  "solana-bpf-loader-program",
  "solana-bucket-map",

runtime/Cargo.toml

@@ -43,6 +43,7 @@ rayon = { workspace = true }
 regex = { workspace = true }
 serde = { workspace = true, features = ["rc"] }
 serde_derive = { workspace = true }
+siphasher = { workspace = true }
 solana-address-lookup-table-program = { workspace = true }
 solana-bpf-loader-program = { workspace = true }
 solana-bucket-map = { workspace = true }

runtime/src/bank/metrics.rs

@@ -23,6 +23,7 @@ pub(crate) struct RewardsMetrics {
     pub(crate) invalid_cached_stake_accounts: usize,
     pub(crate) invalid_cached_stake_accounts_rent_epoch: usize,
     pub(crate) vote_accounts_cache_miss_count: usize,
+    pub(crate) hash_partition_rewards_us: u64,
 }
 
 pub(crate) struct NewBankTimings {
@@ -115,6 +116,11 @@ pub(crate) fn report_new_epoch_metrics(
             metrics.vote_accounts_cache_miss_count,
             i64
         ),
+        (
+            "hash_partition_rewards_us",
+            metrics.hash_partition_rewards_us,
+            i64
+        ),
     );
 }
 

runtime/src/bank/tests.rs

@@ -119,6 +119,46 @@ use {
     test_case::test_case,
 };
 
+impl StakeReward {
+    pub fn new_random() -> Self {
+        let mut rng = rand::thread_rng();
+
+        let rent = Rent::free();
+
+        let validator_pubkey = solana_sdk::pubkey::new_rand();
+        let validator_stake_lamports = 20;
+        let validator_staking_keypair = Keypair::new();
+        let validator_voting_keypair = Keypair::new();
+
+        let validator_vote_account = vote_state::create_account(
+            &validator_voting_keypair.pubkey(),
+            &validator_pubkey,
+            10,
+            validator_stake_lamports,
+        );
+
+        let validator_stake_account = stake_state::create_account(
+            &validator_staking_keypair.pubkey(),
+            &validator_voting_keypair.pubkey(),
+            &validator_vote_account,
+            &rent,
+            validator_stake_lamports,
+        );
+
+        Self {
+            stake_pubkey: Pubkey::new_unique(),
+            stake_reward_info: RewardInfo {
+                reward_type: RewardType::Staking,
+                lamports: rng.gen_range(1, 200),
+                post_balance: 0,  /* unused atm */
+                commission: None, /* unused atm */
+            },
+
+            stake_account: validator_stake_account,
+        }
+    }
+}
+
 #[test]
 fn test_race_register_tick_freeze() {
     solana_logger::setup();

runtime/src/epoch_rewards_hasher.rs

@@ -0,0 +1,278 @@
+use {
+    crate::bank::StakeRewards,
+    siphasher::sip::SipHasher13,
+    solana_sdk::{hash::Hash, pubkey::Pubkey},
+    std::hash::Hasher,
+};
+
+#[derive(Debug, Clone)]
+pub(crate) struct EpochRewardsHasher {
+    hasher: SipHasher13,
+    partitions: usize,
+}
+
+impl EpochRewardsHasher {
+    /// Use SipHasher13 keyed on the `seed` for calculating epoch reward partition
+    pub(crate) fn new(partitions: usize, seed: &Hash) -> Self {
+        let mut hasher = SipHasher13::new();
+        hasher.write(seed.as_ref());
+        Self { hasher, partitions }
+    }
+
+    /// Return partition index (0..partitions) by hashing `address` with the `hasher`
+    pub(crate) fn hash_address_to_partition(self, address: &Pubkey) -> usize {
+        let Self {
+            mut hasher,
+            partitions,
+        } = self;
+        hasher.write(address.as_ref());
+        let hash64 = hasher.finish();
+
+        hash_to_partition(hash64, partitions)
+    }
+}
+
+/// Compute the partition index by modulo the address hash to number of partitions w.o bias.
+/// (rand_int * DESIRED_RANGE_MAX) / (RAND_MAX + 1)
+fn hash_to_partition(hash: u64, partitions: usize) -> usize {
+    ((partitions as u128)
+        .saturating_mul(u128::from(hash))
+        .saturating_div(u128::from(u64::MAX).saturating_add(1))) as usize
+}
+
+#[allow(dead_code)]
+fn hash_rewards_into_partitions(
+    stake_rewards: StakeRewards,
+    parent_block_hash: &Hash,
+    num_partitions: usize,
+) -> Vec<StakeRewards> {
+    let hasher = EpochRewardsHasher::new(num_partitions, parent_block_hash);
+    let mut result = vec![vec![]; num_partitions];
+
+    for reward in stake_rewards {
+        // clone here so the hasher's state is re-used on each call to `hash_address_to_partition`.
+        // This prevents us from re-hashing the seed each time.
+        // The clone is explicit (as opposed to an implicit copy) so it is clear this is intended.
+        let partition_index = hasher
+            .clone()
+            .hash_address_to_partition(&reward.stake_pubkey);
+        result[partition_index].push(reward);
+    }
+    result
+}
+
+#[cfg(test)]
+mod tests {
+    use {
+        super::*,
+        crate::stake_rewards::StakeReward,
+        std::{collections::HashMap, ops::RangeInclusive},
+    };
+
+    #[test]
+    fn test_get_equal_partition_range() {
+        // show how 2 equal partition ranges are 0..=(max/2), (max/2+1)..=max
+        // the inclusive is tricky to think about
+        let range = get_equal_partition_range(0, 2);
+        assert_eq!(*range.start(), 0);
+        assert_eq!(*range.end(), u64::MAX / 2);
+        let range = get_equal_partition_range(1, 2);
+        assert_eq!(*range.start(), u64::MAX / 2 + 1);
+        assert_eq!(*range.end(), u64::MAX);
+    }
+
+    #[test]
+    fn test_hash_to_partitions() {
+        let partitions = 16;
+        assert_eq!(hash_to_partition(0, partitions), 0);
+        assert_eq!(hash_to_partition(u64::MAX / 16, partitions), 0);
+        assert_eq!(hash_to_partition(u64::MAX / 16 + 1, partitions), 1);
+        assert_eq!(hash_to_partition(u64::MAX / 16 * 2, partitions), 1);
+        assert_eq!(hash_to_partition(u64::MAX / 16 * 2 + 1, partitions), 1);
+        assert_eq!(hash_to_partition(u64::MAX - 1, partitions), partitions - 1);
+        assert_eq!(hash_to_partition(u64::MAX, partitions), partitions - 1);
+    }
+
+    fn test_partitions(partition: usize, partitions: usize) {
+        let partition = partition.min(partitions - 1);
+        let range = get_equal_partition_range(partition, partitions);
+        // beginning and end of this partition
+        assert_eq!(hash_to_partition(*range.start(), partitions), partition);
+        assert_eq!(hash_to_partition(*range.end(), partitions), partition);
+        if partition < partitions - 1 {
+            // first index in next partition
+            assert_eq!(
+                hash_to_partition(*range.end() + 1, partitions),
+                partition + 1
+            );
+        } else {
+            assert_eq!(*range.end(), u64::MAX);
+        }
+        if partition > 0 {
+            // last index in previous partition
+            assert_eq!(
+                hash_to_partition(*range.start() - 1, partitions),
+                partition - 1
+            );
+        } else {
+            assert_eq!(*range.start(), 0);
+        }
+    }
+
+    #[test]
+    fn test_hash_to_partitions_equal_ranges() {
+        for partitions in [2, 4, 8, 16, 4096] {
+            assert_eq!(hash_to_partition(0, partitions), 0);
+            for partition in [0, 1, 2, partitions - 1] {
+                test_partitions(partition, partitions);
+            }
+
+            let range = get_equal_partition_range(0, partitions);
+            for partition in 1..partitions {
+                let this_range = get_equal_partition_range(partition, partitions);
+                assert_eq!(
+                    this_range.end() - this_range.start(),
+                    range.end() - range.start()
+                );
+            }
+        }
+        // verify non-evenly divisible partitions (partitions will be different sizes by at most 1 from any other partition)
+        for partitions in [3, 19, 1019, 4095] {
+            for partition in [0, 1, 2, partitions - 1] {
+                test_partitions(partition, partitions);
+            }
+            let expected_len_of_partition =
+                ((u128::from(u64::MAX) + 1) / partitions as u128) as u64;
+            for partition in 0..partitions {
+                let this_range = get_equal_partition_range(partition, partitions);
+                let len = this_range.end() - this_range.start();
+                // size is same or 1 less
+                assert!(
+                    len == expected_len_of_partition || len + 1 == expected_len_of_partition,
+                    "{}, {}, {}, {}",
+                    expected_len_of_partition,
+                    len,
+                    partition,
+                    partitions
+                );
+            }
+        }
+    }
+
+    /// return start and end_inclusive of `partition` indexes out of from u64::MAX+1 elements in equal `partitions`
+    /// These will be equal as long as (u64::MAX + 1) divides by `partitions` evenly
+    fn get_equal_partition_range(partition: usize, partitions: usize) -> RangeInclusive<u64> {
+        let max_inclusive = u128::from(u64::MAX);
+        let max_plus_1 = max_inclusive + 1;
+        let partition = partition as u128;
+        let partitions = partitions as u128;
+        let mut start = max_plus_1 * partition / partitions;
+        if partition > 0 && start * partitions / max_plus_1 == partition - 1 {
+            // partitions don't evenly divide and the start of this partition needs to be 1 greater
+            start += 1;
+        }
+
+        let mut end_inclusive = start + max_plus_1 / partitions - 1;
+        if partition < partitions.saturating_sub(1) {
+            let next = end_inclusive + 1;
+            if next * partitions / max_plus_1 == partition {
+                // this partition is far enough into partitions such that the len of this partition is 1 larger than expected
+                end_inclusive += 1;
+            }
+        } else {
+            end_inclusive = max_inclusive;
+        }
+        RangeInclusive::new(start as u64, end_inclusive as u64)
+    }
+
+    /// Make sure that each time hash_address_to_partition is called, it uses the initial seed state and that clone correctly copies the initial hasher state.
+    #[test]
+    fn test_hasher_copy() {
+        let seed = Hash::new_unique();
+        let partitions = 10;
+        let hasher = EpochRewardsHasher::new(partitions, &seed);
+
+        let pk = Pubkey::new_unique();
+
+        let b1 = hasher.clone().hash_address_to_partition(&pk);
+        let b2 = hasher.hash_address_to_partition(&pk);
+        assert_eq!(b1, b2);
+
+        // make sure b1 includes the seed's hash
+        let mut hasher = SipHasher13::new();
+        hasher.write(seed.as_ref());
+        hasher.write(pk.as_ref());
+        let partition = hash_to_partition(hasher.finish(), partitions);
+        assert_eq!(partition, b1);
+    }
+
+    #[test]
+    fn test_hash_rewards_into_partitions() {
+        // setup the expected number of stake rewards
+        let expected_num = 12345;
+
+        let stake_rewards = (0..expected_num)
+            .map(|_| StakeReward::new_random())
+            .collect::<Vec<_>>();
+
+        let total = stake_rewards
+            .iter()
+            .map(|stake_reward| stake_reward.stake_reward_info.lamports)
+            .sum::<i64>();
+
+        let stake_rewards_in_bucket =
+            hash_rewards_into_partitions(stake_rewards.clone(), &Hash::default(), 5);
+
+        let stake_rewards_in_bucket_clone =
+            stake_rewards_in_bucket.iter().flatten().cloned().collect();
+        compare(&stake_rewards, &stake_rewards_in_bucket_clone);
+
+        let total_after_hash_partition = stake_rewards_in_bucket
+            .iter()
+            .flatten()
+            .map(|stake_reward| stake_reward.stake_reward_info.lamports)
+            .sum::<i64>();
+
+        let total_num_after_hash_partition: usize =
+            stake_rewards_in_bucket.iter().map(|x| x.len()).sum();
+
+        // assert total is same, so nothing is dropped or duplicated
+        assert_eq!(total, total_after_hash_partition);
+        assert_eq!(expected_num, total_num_after_hash_partition);
+    }
+
+    #[test]
+    fn test_hash_rewards_into_partitions_empty() {
+        let stake_rewards = vec![];
+        let total = 0;
+
+        let num_partitions = 5;
+        let stake_rewards_in_bucket =
+            hash_rewards_into_partitions(stake_rewards, &Hash::default(), num_partitions);
+
+        let total_after_hash_partition = stake_rewards_in_bucket
+            .iter()
+            .flatten()
+            .map(|stake_reward| stake_reward.stake_reward_info.lamports)
+            .sum::<i64>();
+
+        assert_eq!(total, total_after_hash_partition);
+
+        assert_eq!(stake_rewards_in_bucket.len(), num_partitions);
+        for bucket in stake_rewards_in_bucket.iter().take(num_partitions) {
+            assert!(bucket.is_empty());
+        }
+    }
+
+    fn compare(a: &StakeRewards, b: &StakeRewards) {
+        let mut a = a
+            .iter()
+            .map(|stake_reward| (stake_reward.stake_pubkey, stake_reward.clone()))
+            .collect::<HashMap<_, _>>();
+        b.iter().for_each(|stake_reward| {
+            let reward = a.remove(&stake_reward.stake_pubkey).unwrap();
+            assert_eq!(&reward, stake_reward);
+        });
+        assert!(a.is_empty());
+    }
+}

runtime/src/lib.rs

@@ -39,6 +39,7 @@ pub mod contains;
 pub mod cost_model;
 pub mod cost_tracker;
 pub mod epoch_accounts_hash;
+mod epoch_rewards_hasher;
 pub mod epoch_stakes;
 pub mod genesis_utils;
 pub mod hardened_unpack;