Rewrite entire schedule without holding lock + fix metrics (solana-la…

…bs#23)
jito-labs · Aug 15, 2022 · 4db699c · 4db699c
1 parent a2c1db3
commit 4db699c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 14 deletions.
diff --git a/relayer/src/relayer.rs b/relayer/src/relayer.rs
@@ -381,7 +381,12 @@ impl RelayerImpl {
                         proto_packet_batch.clone(),
                     )),
                 })) {
-                    Ok(_) => None,
+                    Ok(_) => {
+                        router_metrics.num_batches_forwarded += 1;
+                        router_metrics.num_packets_forwarded +=
+                            proto_packet_batch.packets.len() as u64;
+                        None
+                    }
                     Err(TrySendError::Full(_)) => {
                         error!("packet channel is full for pubkey: {:?}", pubkey);
                         router_metrics.num_try_send_channel_full += 1;
@@ -394,9 +399,6 @@ impl RelayerImpl {
                 }
             })
             .collect();
-
-        router_metrics.num_batches_forwarded += 1;
-        router_metrics.num_packets_forwarded += proto_packet_batch.packets.len() as u64;
         Ok(failed_forwards)
     }
 

diff --git a/relayer/src/schedule_cache.rs b/relayer/src/schedule_cache.rs
@@ -11,9 +11,12 @@ use std::{
 };
 
 use jito_rpc::load_balancer::LoadBalancer;
-use log::{debug, error};
+use log::{debug, error, Log};
 use solana_metrics::datapoint_info;
-use solana_sdk::{clock::Slot, pubkey::Pubkey};
+use solana_sdk::{
+    clock::{Slot, DEFAULT_SLOTS_PER_EPOCH},
+    pubkey::Pubkey,
+};
 
 pub struct LeaderScheduleCacheUpdater {
     /// Maps slots to scheduled pubkey
@@ -123,21 +126,18 @@ impl LeaderScheduleCacheUpdater {
             if let Ok(Some(leader_schedule)) = rpc_client.get_leader_schedule(None) {
                 let epoch_offset = epoch_info.absolute_slot - epoch_info.slot_index;
 
-                debug!("Got Leader Schedule. Length = {}", leader_schedule.len());
-
-                let mut schedule = schedule.write().unwrap();
+                debug!("read leader schedule of length: {}", leader_schedule.len());
 
-                // Remove Old Slots
-                schedule.retain(|s, _| *s >= epoch_info.absolute_slot);
-
-                // Add New Slots
+                let mut new_schedule = HashMap::with_capacity(DEFAULT_SLOTS_PER_EPOCH as usize);
                 for (pk_str, slots) in leader_schedule.iter() {
                     for slot in slots.iter() {
                         if let Ok(pubkey) = Pubkey::from_str(pk_str) {
-                            schedule.insert(*slot as u64 + epoch_offset, pubkey);
+                            new_schedule.insert(*slot as u64 + epoch_offset, pubkey);
                         }
                     }
                 }
+                *schedule.write().unwrap() = new_schedule;
+
                 return true;
             } else {
                 error!("Couldn't Get Leader Schedule Update from RPC!!!")