Use num_partitions to find specific stake rewards in partitions

rpc-client-api/src/custom_error.rs

@@ -24,6 +24,7 @@ pub const JSON_RPC_SERVER_ERROR_TRANSACTION_SIGNATURE_LEN_MISMATCH: i64 = -32013
 pub const JSON_RPC_SERVER_ERROR_BLOCK_STATUS_NOT_AVAILABLE_YET: i64 = -32014;
 pub const JSON_RPC_SERVER_ERROR_UNSUPPORTED_TRANSACTION_VERSION: i64 = -32015;
 pub const JSON_RPC_SERVER_ERROR_MIN_CONTEXT_SLOT_NOT_REACHED: i64 = -32016;
+pub const JSON_RPC_SERVER_ERROR_EPOCH_REWARDS_PERIOD_ACTIVE: i64 = -32017;
 
 #[derive(Error, Debug)]
 pub enum RpcCustomError {
@@ -65,6 +66,12 @@ pub enum RpcCustomError {
     UnsupportedTransactionVersion(u8),
     #[error("MinContextSlotNotReached")]
     MinContextSlotNotReached { context_slot: Slot },
+    #[error("EpochRewardsPeriodActive")]
+    EpochRewardsPeriodActive {
+        slot: Slot,
+        current_block_height: u64,
+        rewards_complete_block_height: u64,
+    },
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -79,6 +86,13 @@ pub struct MinContextSlotNotReachedErrorData {
     pub context_slot: Slot,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct EpochRewardsPeriodActiveErrorData {
+    pub current_block_height: u64,
+    pub rewards_complete_block_height: u64,
+}
+
 impl From<EncodeError> for RpcCustomError {
     fn from(err: EncodeError) -> Self {
         match err {
@@ -206,6 +220,14 @@ impl From<RpcCustomError> for Error {
                     context_slot,
                 })),
             },
+            RpcCustomError::EpochRewardsPeriodActive { slot, current_block_height, rewards_complete_block_height } => Self {
+                code: ErrorCode::ServerError(JSON_RPC_SERVER_ERROR_EPOCH_REWARDS_PERIOD_ACTIVE),
+                message: format!("Epoch rewards period still active at slot {slot}"),
+                data: Some(serde_json::json!(EpochRewardsPeriodActiveErrorData {
+                    current_block_height,
+                    rewards_complete_block_height,
+                })),
+            },
         }
     }
 }

rpc/src/rpc.rs

@@ -704,9 +704,29 @@ impl JsonRpcRequestProcessor {
                 .await?;
 
             for (partition_index, addresses) in partition_index_addresses.iter() {
-                let slot = *block_list
-                    .get(*partition_index)
-                    .ok_or_else(Error::internal_error)?;
+                let slot = *block_list.get(*partition_index).ok_or_else(|| {
+                    // If block_list.len() too short to contain
+                    // partition_index, the epoch rewards period must be
+                    // currently active.
+                    let rewards_complete_block_height = epoch_boundary_block
+                        .block_height
+                        .map(|block_height| {
+                            block_height
+                                .saturating_add(num_partitions as u64)
+                                .saturating_add(1)
+                        })
+                        .expect(
+                            "every block after partitioned_epoch_reward_enabled should have a \
+                             populated block_height",
+                        );
+                    {
+                        RpcCustomError::EpochRewardsPeriodActive {
+                            slot: bank.slot(),
+                            current_block_height: bank.block_height(),
+                            rewards_complete_block_height,
+                        }
+                    }
+                })?;
 
                 let Ok(Some(block)) = self
                     .get_block(