Enable partitioned epoch reward by feature id (solana-labs#32174)

enable partitioned epoch reward feature by feature id
update stake rewards tests with partitioned epoch rewards feature enable

Co-authored-by: HaoranYi <haoran.yi@solana.com>

cli/tests/stake.rs

@@ -157,8 +157,8 @@ fn test_stake_redelegation() {
     };
     process_command(&config).unwrap();
 
-    // wait for new epoch
-    wait_for_next_epoch_plus_n_slots(&rpc_client, 0);
+    // wait for new epoch plus one additional slot for rewards payout
+    wait_for_next_epoch_plus_n_slots(&rpc_client, 1);
 
     // `stake_keypair` should now be delegated to `vote_keypair` and fully activated
     let stake_account = rpc_client.get_account(&stake_keypair.pubkey()).unwrap();
@@ -197,10 +197,11 @@ fn test_stake_redelegation() {
     )
     .unwrap();
 
-    // wait for a new epoch to ensure the `Redelegate` happens as soon as possible in the epoch
-    // to reduce the risk of a race condition when checking the stake account correctly enters the
-    // deactivating state for the remainder of the current epoch
-    wait_for_next_epoch_plus_n_slots(&rpc_client, 0);
+    // wait for a new epoch to ensure the `Redelegate` happens as soon as possible (i.e. till the
+    // last reward distribution block in the new epoch) to reduce the risk of a race condition
+    // when checking the stake account correctly enters the deactivating state for the
+    // remainder of the current epoch.
+    wait_for_next_epoch_plus_n_slots(&rpc_client, 1);
 
     // Redelegate to `vote2_keypair` via `stake2_keypair
     config.signers = vec![&default_signer, &stake2_keypair];
@@ -250,8 +251,8 @@ fn test_stake_redelegation() {
     check_balance!(rent_exempt_reserve, &rpc_client, &stake_keypair.pubkey());
     check_balance!(50_000_000_000, &rpc_client, &stake2_keypair.pubkey());
 
-    // wait for new epoch
-    wait_for_next_epoch_plus_n_slots(&rpc_client, 0);
+    // wait for new epoch plus reward blocks
+    wait_for_next_epoch_plus_n_slots(&rpc_client, 1);
 
     // `stake_keypair` should now be deactivated
     assert_eq!(

program-test/src/lib.rs

@@ -1167,6 +1167,43 @@ impl ProgramTestContext {
         Ok(())
     }
 
+    /// warp forward one more slot and force reward interval end
+    pub fn warp_forward_force_reward_interval_end(&mut self) -> Result<(), ProgramTestError> {
+        let mut bank_forks = self.bank_forks.write().unwrap();
+        let bank = bank_forks.working_bank();
+
+        // Fill ticks until a new blockhash is recorded, otherwise retried transactions will have
+        // the same signature
+        bank.fill_bank_with_ticks_for_tests();
+        let pre_warp_slot = bank.slot();
+
+        bank_forks.set_root(
+            pre_warp_slot,
+            &solana_runtime::accounts_background_service::AbsRequestSender::default(),
+            Some(pre_warp_slot),
+        );
+
+        // warp_bank is frozen so go forward to get unfrozen bank at warp_slot
+        let warp_slot = pre_warp_slot + 1;
+        let mut warp_bank = Bank::new_from_parent(&bank, &Pubkey::default(), warp_slot);
+
+        warp_bank.force_reward_interval_end_for_tests();
+        bank_forks.insert(warp_bank);
+
+        // Update block commitment cache, otherwise banks server will poll at
+        // the wrong slot
+        let mut w_block_commitment_cache = self.block_commitment_cache.write().unwrap();
+        // HACK: The root set here should be `pre_warp_slot`, but since we're
+        // in a testing environment, the root bank never updates after a warp.
+        // The ticking thread only updates the working bank, and never the root
+        // bank.
+        w_block_commitment_cache.set_all_slots(warp_slot, warp_slot);
+
+        let bank = bank_forks.working_bank();
+        self.last_blockhash = bank.last_blockhash();
+        Ok(())
+    }
+
     /// Get a new latest blockhash, similar in spirit to RpcClient::get_latest_blockhash()
     pub async fn get_new_latest_blockhash(&mut self) -> io::Result<Hash> {
         let blockhash = self

program-test/tests/warp.rs

@@ -245,7 +245,7 @@ async fn stake_rewards_from_warp() {
     // go forward and see that rewards have been distributed
     let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
     context
-        .warp_to_slot(first_normal_slot + slots_per_epoch)
+        .warp_to_slot(first_normal_slot + slots_per_epoch + 1) // when partitioned rewards are enabled, the rewards are paid at 1 slot after the first slot of the epoch
         .unwrap();
 
     let account = context
@@ -350,7 +350,7 @@ async fn stake_rewards_filter_bench_core(num_stake_accounts: u64) {
     // go forward and see that rewards have been distributed
     let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
     context
-        .warp_to_slot(first_normal_slot + slots_per_epoch)
+        .warp_to_slot(first_normal_slot + slots_per_epoch + 1) // when partitioned rewards are enabled, the rewards are paid at 1 slot after the first slot of the epoch
         .unwrap();
 
     let account = context
@@ -427,6 +427,7 @@ async fn stake_merge_immediately_after_activation() {
     let slots_per_epoch = context.genesis_config().epoch_schedule.slots_per_epoch;
     let mut current_slot = first_normal_slot + slots_per_epoch;
     context.warp_to_slot(current_slot).unwrap();
+    context.warp_forward_force_reward_interval_end().unwrap();
 
     // this is annoying, but if no stake has earned rewards, the bank won't
     // iterate through the stakes at all, which means we can only test the
@@ -442,6 +443,7 @@ async fn stake_merge_immediately_after_activation() {
 
     current_slot += slots_per_epoch;
     context.warp_to_slot(current_slot).unwrap();
+    context.warp_forward_force_reward_interval_end().unwrap();
 
     // make another stake which will just have its credits observed advanced
     let absorbed_stake_address =
@@ -454,6 +456,7 @@ async fn stake_merge_immediately_after_activation() {
     context.increment_vote_account_credits(&vote_address, 100);
     current_slot += slots_per_epoch;
     context.warp_to_slot(current_slot).unwrap();
+    context.warp_forward_force_reward_interval_end().unwrap();
 
     // check that base stake has earned rewards and credits moved forward
     let stake_account = context

runtime/src/bank.rs

@@ -1493,7 +1493,8 @@ impl Bank {
 
     #[allow(dead_code)]
     fn is_partitioned_rewards_feature_enabled(&self) -> bool {
-        false // Will be feature later. It is convenient to have a constant fn at the moment.
+        self.feature_set
+            .is_active(&feature_set::enable_partitioned_epoch_reward::id())
     }
 
     #[allow(dead_code)]

runtime/src/bank/tests.rs

@@ -12604,6 +12604,16 @@ fn test_rewards_point_calculation_empty() {
     assert!(point_value.is_none());
 }
 
+#[test]
+fn test_is_partitioned_reward_feature_enable() {
+    let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);
+
+    let mut bank = Bank::new_for_tests(&genesis_config);
+    assert!(!bank.is_partitioned_rewards_feature_enabled());
+    bank.activate_feature(&feature_set::enable_partitioned_epoch_reward::id());
+    assert!(bank.is_partitioned_rewards_feature_enabled());
+}
+
 #[test]
 fn test_deactivate_epoch_reward_status() {
     let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000 * LAMPORTS_PER_SOL);

runtime/tests/stake.rs

@@ -420,8 +420,8 @@ fn test_stake_account_lifetime() {
     let pre_staked = get_staked(&bank, &stake_pubkey);
     let pre_balance = bank.get_balance(&stake_pubkey);
 
-    // next epoch bank should pay rewards
-    bank = next_epoch_and_n_slots(&bank, 0);
+    // next epoch bank plus one additional slot should pay rewards
+    bank = next_epoch_and_n_slots(&bank, 1);
 
     // Test that balance increased, and that the balance got staked
     let staked = get_staked(&bank, &stake_pubkey);
@@ -489,7 +489,7 @@ fn test_stake_account_lifetime() {
         .send_and_confirm_message(&[&mint_keypair, &stake_keypair], message)
         .is_err());
 
-    let mut bank = next_epoch_and_n_slots(&bank, 0);
+    let mut bank = next_epoch_and_n_slots(&bank, 1);
 
     let bank_client = BankClient::new_shared(&bank);
 
@@ -535,7 +535,7 @@ fn test_stake_account_lifetime() {
         if get_staked(&bank, &split_stake_pubkey) == 0 {
             break;
         }
-        bank = next_epoch_and_n_slots(&bank, 0);
+        bank = next_epoch_and_n_slots(&bank, 1);
     }
     let bank_client = BankClient::new_shared(&bank);