Include post balance information for rewards (#12598)

* Include post balance information for rewards

* Add post-balance to stored Reward struct

* Handle extended Reward in bigtable

Co-authored-by: Michael Vines <mvines@gmail.com>

core/src/rewards_recorder_service.rs

@@ -1,5 +1,6 @@
 use crossbeam_channel::{Receiver, RecvTimeoutError, Sender};
 use solana_ledger::blockstore::Blockstore;
+use solana_runtime::bank::RewardInfo;
 use solana_sdk::{clock::Slot, pubkey::Pubkey};
 use solana_transaction_status::Reward;
 use std::{
@@ -11,8 +12,8 @@ use std::{
     time::Duration,
 };
 
-pub type RewardsRecorderReceiver = Receiver<(Slot, Vec<(Pubkey, i64)>)>;
+pub type RewardsRecorderReceiver = Receiver<(Slot, Vec<(Pubkey, RewardInfo)>)>;
-pub type RewardsRecorderSender = Sender<(Slot, Vec<(Pubkey, i64)>)>;
+pub type RewardsRecorderSender = Sender<(Slot, Vec<(Pubkey, RewardInfo)>)>;
 
 pub struct RewardsRecorderService {
     thread_hdl: JoinHandle<()>,
@@ -49,9 +50,10 @@ impl RewardsRecorderService {
         let (slot, rewards) = rewards_receiver.recv_timeout(Duration::from_secs(1))?;
         let rpc_rewards = rewards
             .into_iter()
-            .map(|(pubkey, lamports)| Reward {
+            .map(|(pubkey, reward_info)| Reward {
                 pubkey: pubkey.to_string(),
-                lamports,
+                lamports: reward_info.lamports,
+                post_balance: reward_info.post_balance,
             })
             .collect();
 

runtime/src/bank.rs

@@ -477,6 +477,12 @@ pub(crate) struct BankFieldsToSerialize<'a> {
     pub(crate) is_delta: bool,
 }
 
+#[derive(Debug, PartialEq, Serialize, Deserialize, AbiExample, Default, Clone, Copy)]
+pub struct RewardInfo {
+    pub lamports: i64,     // Reward amount
+    pub post_balance: u64, // Account balance in lamports after `lamports` was applied
+}
+
 /// Manager for the state of all accounts and programs after processing its entries.
 /// AbiExample is needed even without Serialize/Deserialize; actual (de-)serialization
 /// are implemented elsewhere for versioning
@@ -592,7 +598,7 @@ pub struct Bank {
     pub last_vote_sync: AtomicU64,
 
     /// Rewards that were paid out immediately after this bank was created
-    pub rewards: Option<Vec<(Pubkey, i64)>>,
+    pub rewards: Option<Vec<(Pubkey, RewardInfo)>>,
 
     pub skip_drop: AtomicBool,
 
@@ -1142,7 +1148,13 @@ impl Bank {
         if let Some(rewards) = self.rewards.as_ref() {
             assert_eq!(
                 validator_rewards_paid,
-                u64::try_from(rewards.iter().map(|(_pubkey, reward)| reward).sum::<i64>()).unwrap()
+                u64::try_from(
+                    rewards
+                        .iter()
+                        .map(|(_pubkey, reward_info)| reward_info.lamports)
+                        .sum::<i64>()
+                )
+                .unwrap()
             );
         }
 
@@ -1247,11 +1259,19 @@ impl Bank {
                     vote_account_changed = true;
 
                     if voters_reward > 0 {
-                        *rewards.entry(*vote_pubkey).or_insert(0i64) += voters_reward as i64;
+                        let reward_info = rewards
+                            .entry(*vote_pubkey)
+                            .or_insert_with(RewardInfo::default);
+                        reward_info.lamports += voters_reward as i64;
+                        reward_info.post_balance = vote_account.lamports;
                     }
 
                     if stakers_reward > 0 {
-                        *rewards.entry(*stake_pubkey).or_insert(0i64) += stakers_reward as i64;
+                        let reward_info = rewards
+                            .entry(*stake_pubkey)
+                            .or_insert_with(RewardInfo::default);
+                        reward_info.lamports += stakers_reward as i64;
+                        reward_info.post_balance = stake_account.lamports;
                     }
                 } else {
                     debug!(
@@ -5513,7 +5533,10 @@ mod tests {
             bank1.rewards,
             Some(vec![(
                 stake_id,
-                (rewards.validator_point_value * validator_points as f64) as i64
+                RewardInfo {
+                    lamports: (rewards.validator_point_value * validator_points as f64) as i64,
+                    post_balance: bank1.get_balance(&stake_id),
+                }
             )])
         );
         bank1.freeze();

storage-bigtable/proto/solana.bigtable.confirmed_block.rs

@@ -87,6 +87,8 @@ pub struct Reward {
     pub pubkey: std::string::String,
     #[prost(int64, tag = "2")]
     pub lamports: i64,
+    #[prost(uint64, tag = "3")]
+    pub post_balance: u64,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct UnixTimestamp {

storage-bigtable/src/confirmed_block.proto

@@ -60,6 +60,7 @@ message CompiledInstruction {
 message Reward {
     string pubkey = 1;
     int64 lamports = 2;
+    uint64 post_balance = 3;
 }
 
 message UnixTimestamp {

storage-bigtable/src/convert.rs

@@ -23,6 +23,7 @@ impl From<Reward> for generated::Reward {
         Self {
             pubkey: reward.pubkey,
             lamports: reward.lamports,
+            post_balance: reward.post_balance,
         }
     }
 }
@@ -32,6 +33,7 @@ impl From<generated::Reward> for Reward {
         Self {
             pubkey: reward.pubkey,
             lamports: reward.lamports,
+            post_balance: reward.post_balance,
         }
     }
 }

storage-bigtable/src/lib.rs

@@ -8,7 +8,7 @@ use solana_sdk::{
     transaction::{Transaction, TransactionError},
 };
 use solana_transaction_status::{
-    ConfirmedBlock, ConfirmedTransaction, ConfirmedTransactionStatusWithSignature, Rewards,
+    ConfirmedBlock, ConfirmedTransaction, ConfirmedTransactionStatusWithSignature, Reward,
     TransactionStatus, TransactionStatusMeta, TransactionWithStatusMeta,
 };
 use std::{collections::HashMap, convert::TryInto};
@@ -86,7 +86,7 @@ struct StoredConfirmedBlock {
     blockhash: String,
     parent_slot: Slot,
     transactions: Vec<StoredConfirmedBlockTransaction>,
-    rewards: Rewards,
+    rewards: StoredConfirmedBlockRewards,
     block_time: Option<UnixTimestamp>,
 }
 
@@ -106,7 +106,7 @@ impl From<ConfirmedBlock> for StoredConfirmedBlock {
             blockhash,
             parent_slot,
             transactions: transactions.into_iter().map(|tx| tx.into()).collect(),
-            rewards,
+            rewards: rewards.into_iter().map(|reward| reward.into()).collect(),
             block_time,
         }
     }
@@ -128,7 +128,7 @@ impl From<StoredConfirmedBlock> for ConfirmedBlock {
             blockhash,
             parent_slot,
             transactions: transactions.into_iter().map(|tx| tx.into()).collect(),
-            rewards,
+            rewards: rewards.into_iter().map(|reward| reward.into()).collect(),
             block_time,
         }
     }
@@ -206,6 +206,34 @@ impl From<TransactionStatusMeta> for StoredConfirmedBlockTransactionStatusMeta {
     }
 }
 
+type StoredConfirmedBlockRewards = Vec<StoredConfirmedBlockReward>;
+
+#[derive(Serialize, Deserialize)]
+struct StoredConfirmedBlockReward {
+    pubkey: String,
+    lamports: i64,
+}
+
+impl From<StoredConfirmedBlockReward> for Reward {
+    fn from(value: StoredConfirmedBlockReward) -> Self {
+        let StoredConfirmedBlockReward { pubkey, lamports } = value;
+        Self {
+            pubkey,
+            lamports,
+            post_balance: 0,
+        }
+    }
+}
+
+impl From<Reward> for StoredConfirmedBlockReward {
+    fn from(value: Reward) -> Self {
+        let Reward {
+            pubkey, lamports, ..
+        } = value;
+        Self { pubkey, lamports }
+    }
+}
+
 // A serialized `TransactionInfo` is stored in the `tx` table
 #[derive(Serialize, Deserialize)]
 struct TransactionInfo {

transaction-status/src/lib.rs

@@ -228,9 +228,12 @@ pub struct ConfirmedTransactionStatusWithSignature {
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
 pub struct Reward {
     pub pubkey: String,
     pub lamports: i64,
+    #[serde(deserialize_with = "default_on_eof")]
+    pub post_balance: u64, // Account balance in lamports after `lamports` was applied
 }
 
 pub type Rewards = Vec<Reward>;