geyser: add num_partitions to block info (#2158)

* geyser: add num_partitions to block info * fix comment
anza-xyz · Jul 18, 2024 · d8791a6 · d8791a6
1 parent c9a42d1
commit d8791a6
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 25 deletions.
diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
@@ -3236,7 +3236,7 @@ impl ReplayStage {
                         &parent_blockhash.to_string(),
                         bank.slot(),
                         &bank.last_blockhash().to_string(),
-                        &bank.rewards,
+                        &bank.get_rewards_and_num_partitions(),
                         Some(bank.clock().unix_timestamp),
                         Some(bank.block_height()),
                         bank.executed_transaction_count(),

diff --git a/geyser-plugin-interface/src/geyser_plugin_interface.rs b/geyser-plugin-interface/src/geyser_plugin_interface.rs
@@ -8,7 +8,7 @@ use {
         signature::Signature,
         transaction::SanitizedTransaction,
     },
-    solana_transaction_status::{Reward, TransactionStatusMeta},
+    solana_transaction_status::{Reward, RewardsAndNumPartitions, TransactionStatusMeta},
     std::{any::Any, error, io},
     thiserror::Error,
 };
@@ -251,11 +251,27 @@ pub struct ReplicaBlockInfoV3<'a> {
     pub entry_count: u64,
 }
 
+/// Extending ReplicaBlockInfo by sending RewardsAndNumPartitions.
+#[derive(Clone, Debug)]
+#[repr(C)]
+pub struct ReplicaBlockInfoV4<'a> {
+    pub parent_slot: Slot,
+    pub parent_blockhash: &'a str,
+    pub slot: Slot,
+    pub blockhash: &'a str,
+    pub rewards: &'a RewardsAndNumPartitions,
+    pub block_time: Option<UnixTimestamp>,
+    pub block_height: Option<u64>,
+    pub executed_transaction_count: u64,
+    pub entry_count: u64,
+}
+
 #[repr(u32)]
 pub enum ReplicaBlockInfoVersions<'a> {
     V0_0_1(&'a ReplicaBlockInfo<'a>),
     V0_0_2(&'a ReplicaBlockInfoV2<'a>),
     V0_0_3(&'a ReplicaBlockInfoV3<'a>),
+    V0_0_4(&'a ReplicaBlockInfoV4<'a>),
 }
 
 /// Errors returned by plugin calls

diff --git a/geyser-plugin-manager/src/block_metadata_notifier.rs b/geyser-plugin-manager/src/block_metadata_notifier.rs
@@ -4,13 +4,14 @@ use {
         geyser_plugin_manager::GeyserPluginManager,
     },
     agave_geyser_plugin_interface::geyser_plugin_interface::{
-        ReplicaBlockInfoV3, ReplicaBlockInfoVersions,
+        ReplicaBlockInfoV4, ReplicaBlockInfoVersions,
     },
     log::*,
     solana_measure::measure::Measure,
     solana_metrics::*,
-    solana_sdk::{clock::UnixTimestamp, pubkey::Pubkey, reward_info::RewardInfo},
-    solana_transaction_status::{Reward, Rewards},
+    solana_runtime::bank::KeyedRewardsAndNumPartitions,
+    solana_sdk::clock::UnixTimestamp,
+    solana_transaction_status::{Reward, RewardsAndNumPartitions},
     std::sync::{Arc, RwLock},
 };
 
@@ -26,7 +27,7 @@ impl BlockMetadataNotifier for BlockMetadataNotifierImpl {
         parent_blockhash: &str,
         slot: u64,
         blockhash: &str,
-        rewards: &RwLock<Vec<(Pubkey, RewardInfo)>>,
+        rewards: &KeyedRewardsAndNumPartitions,
         block_time: Option<UnixTimestamp>,
         block_height: Option<u64>,
         executed_transaction_count: u64,
@@ -51,7 +52,7 @@ impl BlockMetadataNotifier for BlockMetadataNotifierImpl {
                 executed_transaction_count,
                 entry_count,
             );
-            let block_info = ReplicaBlockInfoVersions::V0_0_3(&block_info);
+            let block_info = ReplicaBlockInfoVersions::V0_0_4(&block_info);
             match plugin.notify_block_metadata(block_info) {
                 Err(err) => {
                     error!(
@@ -81,32 +82,35 @@ impl BlockMetadataNotifier for BlockMetadataNotifierImpl {
 }
 
 impl BlockMetadataNotifierImpl {
-    fn build_rewards(rewards: &RwLock<Vec<(Pubkey, RewardInfo)>>) -> Rewards {
-        let rewards = rewards.read().unwrap();
-        rewards
-            .iter()
-            .map(|(pubkey, reward)| Reward {
-                pubkey: pubkey.to_string(),
-                lamports: reward.lamports,
-                post_balance: reward.post_balance,
-                reward_type: Some(reward.reward_type),
-                commission: reward.commission,
-            })
-            .collect()
+    fn build_rewards(rewards: &KeyedRewardsAndNumPartitions) -> RewardsAndNumPartitions {
+        RewardsAndNumPartitions {
+            rewards: rewards
+                .keyed_rewards
+                .iter()
+                .map(|(pubkey, reward)| Reward {
+                    pubkey: pubkey.to_string(),
+                    lamports: reward.lamports,
+                    post_balance: reward.post_balance,
+                    reward_type: Some(reward.reward_type),
+                    commission: reward.commission,
+                })
+                .collect(),
+            num_partitions: rewards.num_partitions,
+        }
     }
 
     fn build_replica_block_info<'a>(
         parent_slot: u64,
         parent_blockhash: &'a str,
         slot: u64,
         blockhash: &'a str,
-        rewards: &'a [Reward],
+        rewards: &'a RewardsAndNumPartitions,
         block_time: Option<UnixTimestamp>,
         block_height: Option<u64>,
         executed_transaction_count: u64,
         entry_count: u64,
-    ) -> ReplicaBlockInfoV3<'a> {
-        ReplicaBlockInfoV3 {
+    ) -> ReplicaBlockInfoV4<'a> {
+        ReplicaBlockInfoV4 {
             parent_slot,
             parent_blockhash,
             slot,

diff --git a/geyser-plugin-manager/src/block_metadata_notifier_interface.rs b/geyser-plugin-manager/src/block_metadata_notifier_interface.rs
@@ -1,6 +1,6 @@
 use {
-    solana_sdk::{clock::UnixTimestamp, pubkey::Pubkey, reward_info::RewardInfo},
-    std::sync::{Arc, RwLock},
+    solana_runtime::bank::KeyedRewardsAndNumPartitions, solana_sdk::clock::UnixTimestamp,
+    std::sync::Arc,
 };
 
 /// Interface for notifying block metadata changes
@@ -13,7 +13,7 @@ pub trait BlockMetadataNotifier {
         parent_blockhash: &str,
         slot: u64,
         blockhash: &str,
-        rewards: &RwLock<Vec<(Pubkey, RewardInfo)>>,
+        rewards: &KeyedRewardsAndNumPartitions,
         block_time: Option<UnixTimestamp>,
         block_height: Option<u64>,
         executed_transaction_count: u64,

diff --git a/transaction-status/src/lib.rs b/transaction-status/src/lib.rs
@@ -628,6 +628,7 @@ pub struct Reward {
 
 pub type Rewards = Vec<Reward>;
 
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct RewardsAndNumPartitions {
     pub rewards: Rewards,
     pub num_partitions: Option<u64>,