From e7a676dc012976e69e25ce20e432a434cbe1e517 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Thu, 15 Jun 2023 11:32:44 -0500
Subject: [PATCH] add distribute_partitioned_epoch_rewards (#32124)

* add distribute_partitioned_epoch_rewards

* reviews

* reanme

---------

Co-authored-by: HaoranYi <haoran.yi@solana.com>
---
 runtime/src/bank.rs       | 45 +++++++++++++++++++++++++++++++++++++--
 runtime/src/bank/tests.rs | 30 +++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 3 deletions(-)

diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index bf30bd7311..5d5a600eef 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1730,6 +1730,9 @@ impl Bank {
                 let leader_schedule_epoch = epoch_schedule.get_leader_schedule_epoch(slot);
                 new.update_epoch_stakes(leader_schedule_epoch);
             }
+            if new.is_partitioned_rewards_code_enabled() {
+                new.distribute_partitioned_epoch_rewards();
+            }
         });
 
         // Update sysvars before processing transactions
@@ -1892,15 +1895,53 @@ impl Bank {
         self.set_epoch_reward_status_active(stake_rewards_by_partition);
 
         datapoint_info!(
-            "epoch-reward-status-update",
+            "epoch-rewards-status-update",
             ("start_slot", slot, i64),
             ("start_block_height", self.block_height(), i64),
-            ("activate", 1, i64),
+            ("active", 1, i64),
             ("parent_slot", parent_slot, i64),
             ("parent_block_height", parent_block_height, i64),
         );
     }
 
+    /// Process reward distribution for the block if it is inside reward interval.
+    fn distribute_partitioned_epoch_rewards(&mut self) {
+        let EpochRewardStatus::Active(status) = &self.epoch_reward_status
+            else {
+                return;
+            };
+
+        assert!(
+            self.epoch_schedule.get_slots_in_epoch(self.epoch)
+                > self.get_reward_total_num_blocks(status.calculated_epoch_stake_rewards.len())
+        );
+        let height = self.block_height();
+        let start_block_height = status.start_block_height;
+        let credit_start = start_block_height + self.get_reward_calculation_num_blocks();
+        let credit_end_exclusive = credit_start
+            + self.get_reward_distribution_num_blocks(status.calculated_epoch_stake_rewards.len());
+
+        if height >= credit_start && height < credit_end_exclusive {
+            let partition_index = height - credit_start;
+            self.distribute_epoch_rewards_in_partition(
+                &status.calculated_epoch_stake_rewards,
+                partition_index,
+            );
+        }
+
+        if height.saturating_add(1) >= credit_end_exclusive {
+            datapoint_info!(
+                "epoch-rewards-status-update",
+                ("slot", self.slot(), i64),
+                ("block_height", height, i64),
+                ("active", 0, i64),
+                ("start_block_height", start_block_height, i64),
+            );
+
+            self.deactivate_epoch_reward_status();
+        }
+    }
+
     pub fn byte_limit_for_scans(&self) -> Option<usize> {
         self.rc
             .accounts
diff --git a/runtime/src/bank/tests.rs b/runtime/src/bank/tests.rs
index 95303b39d6..3bb0813672 100644
--- a/runtime/src/bank/tests.rs
+++ b/runtime/src/bank/tests.rs
@@ -12622,8 +12622,36 @@ fn test_deactivate_epoch_reward_status() {
     assert!(bank.get_reward_interval() == RewardInterval::OutsideInterval);
 }
 
-/// Test rewards compuation and partitioned rewards distribution at the epoch boundary
 #[test]
+fn test_distribute_partitioned_epoch_rewards() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    let mut bank = Bank::new_for_tests(&genesis_config);
+
+    let expected_num = 100;
+
+    let stake_rewards = (0..expected_num)
+        .map(|_| StakeReward::new_random())
+        .collect::<Vec<_>>();
+
+    let stake_rewards = hash_rewards_into_partitions(stake_rewards, &Hash::new(&[1; 32]), 100);
+
+    bank.set_epoch_reward_status_active(stake_rewards);
+
+    bank.distribute_partitioned_epoch_rewards();
+}
+
+#[test]
+fn test_distribute_partitioned_epoch_rewards_empty() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+    let mut bank = Bank::new_for_tests(&genesis_config);
+
+    bank.set_epoch_reward_status_active(vec![]);
+
+    bank.distribute_partitioned_epoch_rewards();
+}
+
+#[test]
+/// Test rewards compuation and partitioned rewards distribution at the epoch boundary
 fn test_rewards_computation() {
     solana_logger::setup();