Send votes to next leader's TPU instead of our TPU

solana-labs · Apr 20, 2021 · c8b474c · c8b474c
1 parent a19c3ba
commit c8b474c
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 12 deletions.
diff --git a/core/src/banking_stage.rs b/core/src/banking_stage.rs
@@ -566,15 +566,8 @@ impl BankingStage {
             }
             return;
         }
-        let next_leader = match poh_recorder
-            .lock()
-            .unwrap()
-            .leader_after_n_slots(FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET)
-        {
-            Some(pubkey) => pubkey,
-            None => return,
-        };
-        let addr = match cluster_info.lookup_contact_info(&next_leader, |ci| ci.tpu_forwards) {
+
+        let addr = match next_leader_tpu_forwards(cluster_info, poh_recorder) {
             Some(addr) => addr,
             None => return,
         };
@@ -1368,6 +1361,36 @@ impl BankingStage {
     }
 }
 
+pub(crate) fn next_leader_tpu(
+    cluster_info: &ClusterInfo,
+    poh_recorder: &Arc<Mutex<PohRecorder>>,
+) -> Option<std::net::SocketAddr> {
+    if let Some(leader_pubkey) = poh_recorder
+        .lock()
+        .unwrap()
+        .leader_after_n_slots(FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET)
+    {
+        cluster_info.lookup_contact_info(&leader_pubkey, |leader| leader.tpu)
+    } else {
+        None
+    }
+}
+
+fn next_leader_tpu_forwards(
+    cluster_info: &ClusterInfo,
+    poh_recorder: &Arc<Mutex<PohRecorder>>,
+) -> Option<std::net::SocketAddr> {
+    if let Some(leader_pubkey) = poh_recorder
+        .lock()
+        .unwrap()
+        .leader_after_n_slots(FORWARD_TRANSACTIONS_TO_LEADER_AT_SLOT_OFFSET)
+    {
+        cluster_info.lookup_contact_info(&leader_pubkey, |leader| leader.tpu_forwards)
+    } else {
+        None
+    }
+}
+
 pub fn create_test_recorder(
     bank: &Arc<Bank>,
     blockstore: &Arc<Blockstore>,

diff --git a/core/src/cluster_info.rs b/core/src/cluster_info.rs
@@ -1182,8 +1182,8 @@ impl ClusterInfo {
             .process_push_message(&self_pubkey, vec![vote], now);
     }
 
-    pub fn send_vote(&self, vote: &Transaction) -> Result<()> {
-        let tpu = self.my_contact_info().tpu;
+    pub fn send_vote(&self, vote: &Transaction, tpu: Option<SocketAddr>) -> Result<()> {
+        let tpu = tpu.unwrap_or_else(|| self.my_contact_info().tpu);
         let buf = serialize(vote)?;
         self.socket.send_to(&buf, &tpu)?;
         Ok(())

diff --git a/core/src/replay_stage.rs b/core/src/replay_stage.rs
@@ -526,6 +526,7 @@ impl ReplayStage {
 
                         Self::handle_votable_bank(
                             &vote_bank,
+                            &poh_recorder,
                             switch_fork_decision,
                             &bank_forks,
                             &mut tower,
@@ -1250,6 +1251,7 @@ impl ReplayStage {
     #[allow(clippy::too_many_arguments)]
     fn handle_votable_bank(
         bank: &Arc<Bank>,
+        poh_recorder: &Arc<Mutex<PohRecorder>>,
         switch_fork_decision: &SwitchForkDecision,
         bank_forks: &Arc<RwLock<BankForks>>,
         tower: &mut Tower,
@@ -1350,6 +1352,7 @@ impl ReplayStage {
         Self::push_vote(
             cluster_info,
             bank,
+            poh_recorder,
             vote_account_pubkey,
             authorized_voter_keypairs,
             last_vote,
@@ -1360,9 +1363,11 @@ impl ReplayStage {
         );
     }
 
+    #[allow(clippy::too_many_arguments)]
     fn push_vote(
         cluster_info: &ClusterInfo,
         bank: &Arc<Bank>,
+        poh_recorder: &Arc<Mutex<PohRecorder>>,
         vote_account_pubkey: &Pubkey,
         authorized_voter_keypairs: &[Arc<Keypair>],
         vote: Vote,
@@ -1452,7 +1457,10 @@ impl ReplayStage {
             vote_signatures.clear();
         }
 
-        let _ = cluster_info.send_vote(&vote_tx);
+        let _ = cluster_info.send_vote(
+            &vote_tx,
+            crate::banking_stage::next_leader_tpu(cluster_info, poh_recorder),
+        );
         cluster_info.push_vote(tower, vote_tx);
     }