Add num_partitions to Blockstore rewards (anza-xyz#1601)

* Add num_partitions field to Rewards proto definition * Add type to hold rewards plus num_partitions * Add Bank method to get rewards plus num_partitions for recording * Update Blockstore::write_rewards to use num_partitions * Update RewardsRecorderService to handle num_partitions * Populate num_partitions in ReplayStage::record_rewards * Write num_partitions to Bigtable * Reword KeyedRewardsAndNumPartitions method * Clone immediately * Determine epoch boundary by checking parent epoch * Rename UiConfirmedBlock field * nit: fix comment typo * Add test_get_rewards_and_partitions * Add pre-activation test * Add should_record unit test
pgarg66 · Jun 17, 2024 · 034cd73 · 034cd73
1 parent ebd063e
commit 034cd73
Show file tree

Hide file tree

Showing 11 changed files with 365 additions and 16 deletions.
diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
@@ -4393,10 +4393,10 @@ impl ReplayStage {
 
     fn record_rewards(bank: &Bank, rewards_recorder_sender: &Option<RewardsRecorderSender>) {
         if let Some(rewards_recorder_sender) = rewards_recorder_sender {
-            let rewards = bank.rewards.read().unwrap();
-            if !rewards.is_empty() {
+            let rewards = bank.get_rewards_and_num_partitions();
+            if rewards.should_record() {
                 rewards_recorder_sender
-                    .send(RewardsMessage::Batch((bank.slot(), rewards.clone())))
+                    .send(RewardsMessage::Batch((bank.slot(), rewards)))
                     .unwrap_or_else(|err| warn!("rewards_recorder_sender failed: {:?}", err));
             }
             rewards_recorder_sender

diff --git a/core/src/rewards_recorder_service.rs b/core/src/rewards_recorder_service.rs
@@ -1,8 +1,9 @@
 use {
     crossbeam_channel::{Receiver, RecvTimeoutError, Sender},
     solana_ledger::blockstore::Blockstore,
-    solana_sdk::{clock::Slot, pubkey::Pubkey, reward_info::RewardInfo},
-    solana_transaction_status::Reward,
+    solana_runtime::bank::KeyedRewardsAndNumPartitions,
+    solana_sdk::clock::Slot,
+    solana_transaction_status::{Reward, RewardsAndNumPartitions},
     std::{
         sync::{
             atomic::{AtomicBool, AtomicU64, Ordering},
@@ -13,7 +14,7 @@ use {
     },
 };
 
-pub type RewardsBatch = (Slot, Vec<(Pubkey, RewardInfo)>);
+pub type RewardsBatch = (Slot, KeyedRewardsAndNumPartitions);
 pub type RewardsRecorderReceiver = Receiver<RewardsMessage>;
 pub type RewardsRecorderSender = Sender<RewardsMessage>;
 
@@ -55,7 +56,13 @@ impl RewardsRecorderService {
         blockstore: &Blockstore,
     ) -> Result<(), RecvTimeoutError> {
         match rewards_receiver.recv_timeout(Duration::from_secs(1))? {
-            RewardsMessage::Batch((slot, rewards)) => {
+            RewardsMessage::Batch((
+                slot,
+                KeyedRewardsAndNumPartitions {
+                    keyed_rewards: rewards,
+                    num_partitions,
+                },
+            )) => {
                 let rpc_rewards = rewards
                     .into_iter()
                     .map(|(pubkey, reward_info)| Reward {
@@ -68,7 +75,13 @@ impl RewardsRecorderService {
                     .collect();
 
                 blockstore
-                    .write_rewards(slot, rpc_rewards)
+                    .write_rewards(
+                        slot,
+                        RewardsAndNumPartitions {
+                            rewards: rpc_rewards,
+                            num_partitions,
+                        },
+                    )
                     .expect("Expect database write to succeed");
             }
             RewardsMessage::Complete(slot) => {

diff --git a/ledger/src/blockstore.rs b/ledger/src/blockstore.rs
@@ -57,8 +57,9 @@ use {
     solana_storage_proto::{StoredExtendedRewards, StoredTransactionStatusMeta},
     solana_transaction_status::{
         ConfirmedTransactionStatusWithSignature, ConfirmedTransactionWithStatusMeta, Rewards,
-        TransactionStatusMeta, TransactionWithStatusMeta, VersionedConfirmedBlock,
-        VersionedConfirmedBlockWithEntries, VersionedTransactionWithStatusMeta,
+        RewardsAndNumPartitions, TransactionStatusMeta, TransactionWithStatusMeta,
+        VersionedConfirmedBlock, VersionedConfirmedBlockWithEntries,
+        VersionedTransactionWithStatusMeta,
     },
     std::{
         borrow::Cow,
@@ -2678,7 +2679,7 @@ impl Blockstore {
                     Hash::default()
                 };
 
-                let rewards = self
+                let (rewards, num_partitions) = self
                     .rewards_cf
                     .get_protobuf_or_bincode::<StoredExtendedRewards>(slot)?
                     .unwrap_or_default()
@@ -2699,6 +2700,7 @@ impl Blockstore {
                     transactions: self
                         .map_transactions_to_statuses(slot, slot_transaction_iterator)?,
                     rewards,
+                    num_partitions,
                     block_time,
                     block_height,
                 };
@@ -3371,7 +3373,7 @@ impl Blockstore {
             .map(|result| result.map(|option| option.into()))
     }
 
-    pub fn write_rewards(&self, index: Slot, rewards: Rewards) -> Result<()> {
+    pub fn write_rewards(&self, index: Slot, rewards: RewardsAndNumPartitions) -> Result<()> {
         let rewards = rewards.into();
         self.rewards_cf.put_protobuf(index, &rewards)
     }
@@ -8302,6 +8304,7 @@ pub mod tests {
             blockhash: blockhash.to_string(),
             previous_blockhash: Hash::default().to_string(),
             rewards: vec![],
+            num_partitions: None,
             block_time: None,
             block_height: None,
         };
@@ -8316,6 +8319,7 @@ pub mod tests {
             blockhash: blockhash.to_string(),
             previous_blockhash: blockhash.to_string(),
             rewards: vec![],
+            num_partitions: None,
             block_time: None,
             block_height: None,
         };
@@ -8333,6 +8337,7 @@ pub mod tests {
             blockhash: blockhash.to_string(),
             previous_blockhash: blockhash.to_string(),
             rewards: vec![],
+            num_partitions: None,
             block_time: None,
             block_height: None,
         };

diff --git a/rpc-client/src/mock_sender.rs b/rpc-client/src/mock_sender.rs
@@ -406,6 +406,7 @@ impl RpcSender for MockSender {
                     version: Some(TransactionVersion::LEGACY),
                 }],
                 rewards: Rewards::new(),
+                num_partitions: None,
                 block_time: None,
                 block_height: Some(428),
             })?,

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
@@ -35,7 +35,6 @@
 //! already been signed and verified.
 #[allow(deprecated)]
 use solana_sdk::recent_blockhashes_account;
-pub use solana_sdk::reward_type::RewardType;
 use {
     crate::{
         bank::{
@@ -205,6 +204,9 @@ use {
         time::{Duration, Instant},
     },
 };
+pub use {
+    partitioned_epoch_rewards::KeyedRewardsAndNumPartitions, solana_sdk::reward_type::RewardType,
+};
 #[cfg(feature = "dev-context-only-utils")]
 use {
     solana_accounts_db::accounts_db::{