From 51cef294156e8cdc2f41b728ff558172120c517a Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Mon, 27 Mar 2023 23:28:28 +0000
Subject: [PATCH] Forwarder: separate get_leader_and_addr

---
 core/src/banking_stage/forwarder.rs | 31 ++++++++++++++++-------------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/core/src/banking_stage/forwarder.rs b/core/src/banking_stage/forwarder.rs
index 4699ee001e9edc..9460dc69c6a219 100644
--- a/core/src/banking_stage/forwarder.rs
+++ b/core/src/banking_stage/forwarder.rs
@@ -17,7 +17,7 @@ use {
     solana_streamer::sendmmsg::batch_send,
     std::{
         iter::repeat,
-        net::UdpSocket,
+        net::{SocketAddr, UdpSocket},
         sync::{atomic::Ordering, Arc, RwLock},
     },
 };
@@ -147,19 +147,8 @@ impl Forwarder {
         usize,
         Option<Pubkey>,
     ) {
-        let leader_and_addr = match forward_option {
-            ForwardOption::NotForward => return (Ok(()), 0, None),
-            ForwardOption::ForwardTransaction => {
-                next_leader_tpu_forwards(&self.cluster_info, &self.poh_recorder)
-            }
-
-            ForwardOption::ForwardTpuVote => {
-                next_leader_tpu_vote(&self.cluster_info, &self.poh_recorder)
-            }
-        };
-        let (leader_pubkey, addr) = match leader_and_addr {
-            Some(leader_and_addr) => leader_and_addr,
-            None => return (Ok(()), 0, None),
+        let Some((leader_pubkey, addr)) = self.get_leader_and_addr(forward_option) else {
+            return (Ok(()), 0, None);
         };
 
         self.update_data_budget();
@@ -218,6 +207,20 @@ impl Forwarder {
         (Ok(()), packet_vec_len, Some(leader_pubkey))
     }
 
+    /// Get the pubkey and socket address for the leader to forward to
+    fn get_leader_and_addr(&self, forward_option: &ForwardOption) -> Option<(Pubkey, SocketAddr)> {
+        match forward_option {
+            ForwardOption::NotForward => None,
+            ForwardOption::ForwardTransaction => {
+                next_leader_tpu_forwards(&self.cluster_info, &self.poh_recorder)
+            }
+
+            ForwardOption::ForwardTpuVote => {
+                next_leader_tpu_vote(&self.cluster_info, &self.poh_recorder)
+            }
+        }
+    }
+
     /// Re-fill the data budget if enough time has passed
     fn update_data_budget(&self) {
         const INTERVAL_MS: u64 = 100;