From da7d94d0f07bff8a58ae54f6f290b85744d0e9a7 Mon Sep 17 00:00:00 2001
From: sakridge <stephen@solana.com>
Date: Sun, 6 Oct 2019 12:56:17 -0700
Subject: [PATCH] Retransmit stage benchmark (#6249)

---
 core/Cargo.toml                  |  6 +++
 core/benches/retransmit_stage.rs | 65 ++++++++++++++++++++++++++++++++
 core/src/retransmit_stage.rs     |  2 +-
 3 files changed, 72 insertions(+), 1 deletion(-)
 create mode 100644 core/benches/retransmit_stage.rs

diff --git a/core/Cargo.toml b/core/Cargo.toml
index 51c840f090..8deb32620b 100644
--- a/core/Cargo.toml
+++ b/core/Cargo.toml
@@ -111,6 +111,12 @@ name = "sigverify_stage"
 [[bench]]
 name = "poh"
 
+[[bench]]
+name = "retransmit_stage"
+
+[[bench]]
+name = "cluster_info"
+
 [[bench]]
 name = "chacha"
 required-features = ["chacha"]
diff --git a/core/benches/retransmit_stage.rs b/core/benches/retransmit_stage.rs
new file mode 100644
index 0000000000..afdedbf239
--- /dev/null
+++ b/core/benches/retransmit_stage.rs
@@ -0,0 +1,65 @@
+#![feature(test)]
+
+extern crate solana_core;
+extern crate test;
+
+use log::*;
+use solana_core::bank_forks::BankForks;
+use solana_core::cluster_info::{ClusterInfo, Node};
+use solana_core::contact_info::ContactInfo;
+use solana_core::genesis_utils::{create_genesis_block, GenesisBlockInfo};
+use solana_core::leader_schedule_cache::LeaderScheduleCache;
+use solana_core::packet::to_packets_chunked;
+use solana_core::retransmit_stage::retransmit;
+use solana_core::test_tx::test_tx;
+use solana_runtime::bank::Bank;
+use solana_sdk::pubkey::Pubkey;
+use solana_sdk::timing::timestamp;
+use std::net::UdpSocket;
+use std::sync::mpsc::channel;
+use std::sync::{Arc, RwLock};
+use test::Bencher;
+
+#[bench]
+fn bench_retransmit(bencher: &mut Bencher) {
+    solana_logger::setup();
+    let mut cluster_info = ClusterInfo::new_with_invalid_keypair(Node::new_localhost().info);
+    const NUM_PEERS: usize = 2;
+    for _ in 0..NUM_PEERS {
+        let id = Pubkey::new_rand();
+        let contact_info = ContactInfo::new_localhost(&id, timestamp());
+        cluster_info.insert_info(contact_info);
+    }
+    let cluster_info = Arc::new(RwLock::new(cluster_info));
+
+    let GenesisBlockInfo { genesis_block, .. } = create_genesis_block(100_000);
+    let bank0 = Bank::new(&genesis_block);
+    let bank_forks = BankForks::new(0, bank0);
+    let bank = bank_forks.working_bank();
+    let bank_forks = Arc::new(RwLock::new(bank_forks));
+    let (packet_sender, packet_receiver) = channel();
+    let socket = UdpSocket::bind("0.0.0.0:0").unwrap();
+    socket.set_nonblocking(true).unwrap();
+
+    let leader_schedule_cache = Arc::new(LeaderScheduleCache::new_from_bank(&bank));
+
+    let tx = test_tx();
+    let len = 4096;
+    let chunk_size = 1024;
+    let batches = to_packets_chunked(&vec![tx; len], chunk_size);
+
+    bencher.iter(move || {
+        for packets in batches.clone() {
+            packet_sender.send(packets).unwrap();
+        }
+        info!("sent...");
+        retransmit(
+            &bank_forks,
+            &leader_schedule_cache,
+            &cluster_info,
+            &packet_receiver,
+            &socket,
+        )
+        .unwrap();
+    });
+}
diff --git a/core/src/retransmit_stage.rs b/core/src/retransmit_stage.rs
index 154260cc3c..559a238785 100644
--- a/core/src/retransmit_stage.rs
+++ b/core/src/retransmit_stage.rs
@@ -23,7 +23,7 @@ use std::sync::{Arc, RwLock};
 use std::thread::{self, Builder, JoinHandle};
 use std::time::Duration;
 
-fn retransmit(
+pub fn retransmit(
     bank_forks: &Arc<RwLock<BankForks>>,
     leader_schedule_cache: &Arc<LeaderScheduleCache>,
     cluster_info: &Arc<RwLock<ClusterInfo>>,