From c38d0d15cd386cec16fc9e5348bb9667e066ded3 Mon Sep 17 00:00:00 2001
From: Ben | mrgn <96081648+SOELTH@users.noreply.github.com>
Date: Sat, 11 May 2024 11:35:06 -0600
Subject: [PATCH] Fix misnomer and warn when StakedNodes fails to update. 
 (#130)

Co-authored-by: ben | mrgn <hyponyms@proton.me>
---
 core/src/staked_nodes_updater_service.rs | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/core/src/staked_nodes_updater_service.rs b/core/src/staked_nodes_updater_service.rs
index 7dfd155c289798..355763e5e5defc 100644
--- a/core/src/staked_nodes_updater_service.rs
+++ b/core/src/staked_nodes_updater_service.rs
@@ -10,11 +10,12 @@ use std::{
 };
 
 use jito_rpc::load_balancer::LoadBalancer;
+use log::warn;
 use solana_client::client_error;
 use solana_sdk::pubkey::Pubkey;
 use solana_streamer::streamer::StakedNodes;
 
-const IP_TO_STAKE_REFRESH_DURATION: Duration = Duration::from_secs(5);
+const PK_TO_STAKE_REFRESH_DURATION: Duration = Duration::from_secs(5);
 
 pub struct StakedNodesUpdaterService {
     thread_hdl: JoinHandle<()>,
@@ -33,13 +34,21 @@ impl StakedNodesUpdaterService {
                 let mut last_stakes = Instant::now();
                 while !exit.load(Ordering::Relaxed) {
                     let mut stake_map = Arc::new(HashMap::new());
-                    if let Ok(true) = Self::try_refresh_ip_to_stake(
+                    match Self::try_refresh_pk_to_stake(
                         &mut last_stakes,
                         &mut stake_map,
                         &rpc_load_balancer,
                     ) {
-                        let shared = StakedNodes::new(stake_map, staked_nodes_overrides.clone());
-                        *shared_staked_nodes.write().unwrap() = shared;
+                        Ok(true) => {
+                            let shared =
+                                StakedNodes::new(stake_map, staked_nodes_overrides.clone());
+                            *shared_staked_nodes.write().unwrap() = shared;
+                        }
+                        Err(err) => {
+                            warn!("Failed to refresh pk to stake map! Error: {:?}", err);
+                            sleep(PK_TO_STAKE_REFRESH_DURATION);
+                        }
+                        _ => {}
                     }
                 }
             })
@@ -48,12 +57,12 @@ impl StakedNodesUpdaterService {
         Self { thread_hdl }
     }
 
-    fn try_refresh_ip_to_stake(
+    fn try_refresh_pk_to_stake(
         last_stakes: &mut Instant,
         pubkey_stake_map: &mut Arc<HashMap<Pubkey, u64>>,
         rpc_load_balancer: &Arc<LoadBalancer>,
     ) -> client_error::Result<bool> {
-        if last_stakes.elapsed() > IP_TO_STAKE_REFRESH_DURATION {
+        if last_stakes.elapsed() > PK_TO_STAKE_REFRESH_DURATION {
             let client = rpc_load_balancer.rpc_client();
             let vote_accounts = client.get_vote_accounts()?;