From fd6e7020eb590441a95017fee805c0cad32d1f97 Mon Sep 17 00:00:00 2001
From: carllin <wumu727@gmail.com>
Date: Fri, 20 Sep 2019 19:37:40 -0700
Subject: [PATCH] Fix bank overlapping another bank's broadcast (#6012)

---
 core/src/broadcast_stage/broadcast_utils.rs | 97 ++++++++++++++++++++-
 1 file changed, 93 insertions(+), 4 deletions(-)

diff --git a/core/src/broadcast_stage/broadcast_utils.rs b/core/src/broadcast_stage/broadcast_utils.rs
index f0fe41874..df7e4bf52 100644
--- a/core/src/broadcast_stage/broadcast_utils.rs
+++ b/core/src/broadcast_stage/broadcast_utils.rs
@@ -22,12 +22,12 @@ const RECEIVE_ENTRY_COUNT_THRESHOLD: usize = 8;
 
 pub(super) fn recv_slot_entries(receiver: &Receiver<WorkingBankEntry>) -> Result<ReceiveResults> {
     let timer = Duration::new(1, 0);
-    let (bank, (entry, mut last_tick)) = receiver.recv_timeout(timer)?;
+    let (mut current_bank, (entry, mut last_tick)) = receiver.recv_timeout(timer)?;
     let recv_start = Instant::now();
 
     let mut entries = vec![entry];
-    let mut slot = bank.slot();
-    let mut max_tick_height = bank.max_tick_height();
+    let mut slot = current_bank.slot();
+    let mut max_tick_height = current_bank.max_tick_height();
 
     assert!(last_tick <= max_tick_height);
 
@@ -39,6 +39,7 @@ pub(super) fn recv_slot_entries(receiver: &Receiver<WorkingBankEntry>) -> Result
                 entries.clear();
                 slot = bank.slot();
                 max_tick_height = bank.max_tick_height();
+                current_bank = bank;
             }
             last_tick = tick_height;
             entries.push(entry);
@@ -58,7 +59,7 @@ pub(super) fn recv_slot_entries(receiver: &Receiver<WorkingBankEntry>) -> Result
     Ok(ReceiveResults {
         entries,
         time_elapsed,
-        bank,
+        bank: current_bank,
         last_tick,
     })
 }
@@ -96,3 +97,91 @@ pub(super) fn entries_to_shreds(
 
     (shred_infos, u64::from(shredder.index))
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::genesis_utils::{create_genesis_block, GenesisBlockInfo};
+    use solana_sdk::genesis_block::GenesisBlock;
+    use solana_sdk::pubkey::Pubkey;
+    use solana_sdk::system_transaction;
+    use solana_sdk::transaction::Transaction;
+    use std::sync::mpsc::channel;
+
+    fn setup_test() -> (GenesisBlock, Arc<Bank>, Transaction) {
+        let GenesisBlockInfo {
+            genesis_block,
+            mint_keypair,
+            ..
+        } = create_genesis_block(2);
+        let bank0 = Arc::new(Bank::new(&genesis_block));
+        let tx = system_transaction::create_user_account(
+            &mint_keypair,
+            &Pubkey::new_rand(),
+            1,
+            genesis_block.hash(),
+        );
+
+        (genesis_block, bank0, tx)
+    }
+
+    #[test]
+    fn test_recv_slot_entries_1() {
+        let (genesis_block, bank0, tx) = setup_test();
+
+        let bank1 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 1));
+        let (s, r) = channel();
+        let mut last_hash = genesis_block.hash();
+
+        assert!(bank1.max_tick_height() > 1);
+        let entries: Vec<_> = (0..bank1.max_tick_height() + 1)
+            .map(|i| {
+                let entry = Entry::new(&last_hash, 1, vec![tx.clone()]);
+                last_hash = entry.hash;
+                s.send((bank1.clone(), (entry.clone(), i))).unwrap();
+                entry
+            })
+            .collect();
+
+        let result = recv_slot_entries(&r).unwrap();
+
+        assert_eq!(result.bank.slot(), bank1.slot());
+        assert_eq!(result.last_tick, bank1.max_tick_height());
+        assert_eq!(result.entries, entries);
+    }
+
+    #[test]
+    fn test_recv_slot_entries_2() {
+        let (genesis_block, bank0, tx) = setup_test();
+
+        let bank1 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 1));
+        let bank2 = Arc::new(Bank::new_from_parent(&bank1, &Pubkey::default(), 2));
+        let (s, r) = channel();
+
+        let mut last_hash = genesis_block.hash();
+        assert!(bank1.max_tick_height() > 1);
+        // Simulate slot 2 interrupting slot 1's transmission
+        let expected_last_index = bank1.max_tick_height() - 1;
+        let last_entry = (0..bank1.max_tick_height())
+            .map(|i| {
+                let entry = Entry::new(&last_hash, 1, vec![tx.clone()]);
+                last_hash = entry.hash;
+                // Interrupt slot 1 right before the last tick
+                if i == expected_last_index {
+                    s.send((bank2.clone(), (entry.clone(), i))).unwrap();
+                    Some(entry)
+                } else {
+                    s.send((bank1.clone(), (entry, i))).unwrap();
+                    None
+                }
+            })
+            .last()
+            .unwrap()
+            .unwrap();
+
+        let result = recv_slot_entries(&r).unwrap();
+        assert_eq!(result.bank.slot(), bank2.slot());
+        assert_eq!(result.last_tick, expected_last_index);
+        assert_eq!(result.entries, vec![last_entry]);
+    }
+}