From c8d9016ec8dde1d3f31b5aeb28692b358d57b1a5 Mon Sep 17 00:00:00 2001
From: Lijun Wang <83639177+lijunwangs@users.noreply.github.com>
Date: Wed, 9 Oct 2024 18:15:58 -0700
Subject: [PATCH] Support notifying bank created slot status in geyser

---
 core/src/replay_stage.rs                         | 16 +++++++++++++---
 .../src/geyser_plugin_interface.rs               |  4 ++++
 .../src/slot_status_notifier.rs                  |  4 ++++
 rpc/src/slot_status_notifier.rs                  |  2 +-
 4 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
index 3884b39f4d113c..c9a32b77bbf5b6 100644
--- a/core/src/replay_stage.rs
+++ b/core/src/replay_stage.rs
@@ -54,7 +54,7 @@ use {
     solana_rpc::{
         optimistically_confirmed_bank_tracker::{BankNotification, BankNotificationSenderConfig},
         rpc_subscriptions::RpcSubscriptions,
-        slot_status_notifier::{self, SlotStatusNotifier},
+        slot_status_notifier::SlotStatusNotifier,
     },
     solana_rpc_client_api::response::SlotUpdate,
     solana_runtime::{
@@ -1126,6 +1126,7 @@ impl ReplayStage {
                         &poh_recorder,
                         &leader_schedule_cache,
                         &rpc_subscriptions,
+                        &slot_status_notifier,
                         &mut progress,
                         &retransmit_slots_sender,
                         &mut skipped_slots_info,
@@ -2056,7 +2057,7 @@ impl ReplayStage {
         poh_recorder: &Arc<RwLock<PohRecorder>>,
         leader_schedule_cache: &Arc<LeaderScheduleCache>,
         rpc_subscriptions: &Arc<RpcSubscriptions>,
-        slot_status_notifier: Option<SlotStatusNotifier>,
+        slot_status_notifier: &Option<SlotStatusNotifier>,
         progress_map: &mut ProgressMap,
         retransmit_slots_sender: &Sender<Slot>,
         skipped_slots_info: &mut SkippedSlotsInfo,
@@ -4077,7 +4078,7 @@ impl ReplayStage {
             slot_status_notifier
                 .read()
                 .unwrap()
-                .notify_bank_created(slot, parent);
+                .notify_created_bank(slot, parent.slot());
         }
         Bank::new_from_parent_with_options(parent, leader, slot, new_bank_options)
     }
@@ -4425,6 +4426,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -4453,6 +4455,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -6322,6 +6325,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -6351,6 +6355,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -6381,6 +6386,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -6410,6 +6416,7 @@ pub(crate) mod tests {
             &bank_forks,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &mut replay_timing,
         );
@@ -8344,6 +8351,7 @@ pub(crate) mod tests {
             &poh_recorder,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &retransmit_slots_sender,
             &mut SkippedSlotsInfo::default(),
@@ -9012,6 +9020,7 @@ pub(crate) mod tests {
             &poh_recorder,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &retransmit_slots_sender,
             &mut SkippedSlotsInfo::default(),
@@ -9038,6 +9047,7 @@ pub(crate) mod tests {
             &poh_recorder,
             &leader_schedule_cache,
             &rpc_subscriptions,
+            &None,
             &mut progress,
             &retransmit_slots_sender,
             &mut SkippedSlotsInfo::default(),
diff --git a/geyser-plugin-interface/src/geyser_plugin_interface.rs b/geyser-plugin-interface/src/geyser_plugin_interface.rs
index 97271310a99f5f..2fdcb750f3cfa4 100644
--- a/geyser-plugin-interface/src/geyser_plugin_interface.rs
+++ b/geyser-plugin-interface/src/geyser_plugin_interface.rs
@@ -325,6 +325,9 @@ pub enum SlotStatus {
 
     /// All shreds for the slot have been received.
     Completed,
+
+    /// A new bank is created with the slot
+    CreatedBank,
 }
 
 impl SlotStatus {
@@ -335,6 +338,7 @@ impl SlotStatus {
             SlotStatus::Rooted => "rooted",
             SlotStatus::FirstShredReceived => "first_shread_received",
             SlotStatus::Completed => "completed",
+            Self::CreatedBank => "created_bank",
         }
     }
 }
diff --git a/geyser-plugin-manager/src/slot_status_notifier.rs b/geyser-plugin-manager/src/slot_status_notifier.rs
index 573ed97d7787af..27dcd03dac28f3 100644
--- a/geyser-plugin-manager/src/slot_status_notifier.rs
+++ b/geyser-plugin-manager/src/slot_status_notifier.rs
@@ -33,6 +33,10 @@ impl SlotStatusNotifierInterface for SlotStatusNotifierImpl {
     fn notify_completed(&self, slot: Slot) {
         self.notify_slot_status(slot, None, SlotStatus::Completed);
     }
+
+    fn notify_created_bank(&self, slot: Slot, parent: Slot) {
+        self.notify_slot_status(slot, Some(parent), SlotStatus::CreatedBank);
+    }
 }
 
 impl SlotStatusNotifierImpl {
diff --git a/rpc/src/slot_status_notifier.rs b/rpc/src/slot_status_notifier.rs
index a71711d96a5f76..38e9bf60a6e091 100644
--- a/rpc/src/slot_status_notifier.rs
+++ b/rpc/src/slot_status_notifier.rs
@@ -20,7 +20,7 @@ pub trait SlotStatusNotifierInterface {
     fn notify_completed(&self, slot: Slot);
 
     /// Notified when the slot has bank created.
-    fn notify_bank_created(&self, slot: Slot, parent: Slot);
+    fn notify_created_bank(&self, slot: Slot, parent: Slot);
 }
 
 pub type SlotStatusNotifier = Arc<RwLock<dyn SlotStatusNotifierInterface + Sync + Send>>;