Return receiver in scheduler::run

Author: KirillLykov

tls-utils/src/quic_client_certificate.rs

@@ -10,8 +10,13 @@ pub struct QuicClientCertificate {
 }
 
 impl QuicClientCertificate {
-    pub fn new(keypair: &Keypair) -> Self {
-        let (certificate, key) = new_dummy_x509_certificate(keypair);
-        Self { certificate, key }
+    pub fn new(keypair: Option<&Keypair>) -> Self {
+        if let Some(keypair) = keypair {
+            let (certificate, key) = new_dummy_x509_certificate(keypair);
+            Self { certificate, key }
+        } else {
+            let (certificate, key) = new_dummy_x509_certificate(&Keypair::new());
+            Self { certificate, key }
+        }
     }
 }

tpu-client-next/src/connection_workers_scheduler.rs

@@ -70,7 +70,7 @@ pub struct ConnectionWorkersSchedulerConfig {
 
     /// Optional stake identity keypair used in the endpoint certificate for
     /// identifying the sender.
-    pub stake_identity: Option<Keypair>,
+    pub identity: Option<Keypair>,
 
     /// The number of connections to be maintained by the scheduler.
     pub num_connections: usize,
@@ -90,6 +90,11 @@ pub struct ConnectionWorkersSchedulerConfig {
     pub leaders_fanout: Fanout,
 }
 
+pub type TransactionStatsAndReceiver = (
+    SendTransactionStatsPerAddr,
+    mpsc::Receiver<TransactionBatch>,
+);
+
 impl ConnectionWorkersScheduler {
     /// Starts the scheduler, which manages the distribution of transactions to
     /// the network's upcoming leaders.
@@ -103,7 +108,7 @@ impl ConnectionWorkersScheduler {
     pub async fn run(
         ConnectionWorkersSchedulerConfig {
             bind,
-            stake_identity: validator_identity,
+            identity,
             num_connections,
             skip_check_transaction_age,
             worker_channel_size,
@@ -113,14 +118,14 @@ impl ConnectionWorkersScheduler {
         mut leader_updater: Box<dyn LeaderUpdater>,
         mut transaction_receiver: mpsc::Receiver<TransactionBatch>,
         cancel: CancellationToken,
-    ) -> Result<SendTransactionStatsPerAddr, ConnectionWorkersSchedulerError> {
-        let endpoint = Self::setup_endpoint(bind, validator_identity)?;
+    ) -> Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError> {
+        let endpoint = Self::setup_endpoint(bind, identity.as_ref())?;
         debug!("Client endpoint bind address: {:?}", endpoint.local_addr());
         let mut workers = WorkersCache::new(num_connections, cancel.clone());
         let mut send_stats_per_addr = SendTransactionStatsPerAddr::new();
 
         loop {
-            let transaction_batch = tokio::select! {
+            let transaction_batch: TransactionBatch = tokio::select! {
                 recv_res = transaction_receiver.recv() => match recv_res {
                     Some(txs) => txs,
                     None => {
@@ -184,19 +189,15 @@ impl ConnectionWorkersScheduler {
 
         endpoint.close(0u32.into(), b"Closing connection");
         leader_updater.stop().await;
-        Ok(send_stats_per_addr)
+        Ok((send_stats_per_addr, transaction_receiver))
     }
 
     /// Sets up the QUIC endpoint for the scheduler to handle connections.
     fn setup_endpoint(
         bind: SocketAddr,
-        validator_identity: Option<Keypair>,
+        identity: Option<&Keypair>,
     ) -> Result<Endpoint, ConnectionWorkersSchedulerError> {
-        let client_certificate = if let Some(validator_identity) = validator_identity {
-            Arc::new(QuicClientCertificate::new(&validator_identity))
-        } else {
-            Arc::new(QuicClientCertificate::new(&Keypair::new()))
-        };
+        let client_certificate = QuicClientCertificate::new(identity);
         let client_config = create_client_config(client_certificate);
         let endpoint = create_client_endpoint(bind, client_config)?;
         Ok(endpoint)

tpu-client-next/src/quic_networking.rs

@@ -18,7 +18,7 @@ pub use {
     solana_tls_utils::QuicClientCertificate,
 };
 
-pub(crate) fn create_client_config(client_certificate: Arc<QuicClientCertificate>) -> ClientConfig {
+pub(crate) fn create_client_config(client_certificate: QuicClientCertificate) -> ClientConfig {
     // adapted from QuicLazyInitializedEndpoint::create_endpoint
     let mut crypto = tls_client_config_builder()
         .with_client_auth_cert(

tpu-client-next/tests/connection_workers_scheduler_test.rs

@@ -15,11 +15,13 @@ use {
         streamer::StakedNodes,
     },
     solana_tpu_client_next::{
-        connection_workers_scheduler::{ConnectionWorkersSchedulerConfig, Fanout},
+        connection_workers_scheduler::{
+            ConnectionWorkersSchedulerConfig, Fanout, TransactionStatsAndReceiver,
+        },
         leader_updater::create_leader_updater,
         send_transaction_stats::SendTransactionStatsNonAtomic,
         transaction_batch::TransactionBatch,
-        ConnectionWorkersScheduler, ConnectionWorkersSchedulerError, SendTransactionStatsPerAddr,
+        ConnectionWorkersScheduler, ConnectionWorkersSchedulerError,
     },
     std::{
         collections::HashMap,
@@ -40,10 +42,10 @@ use {
     tokio_util::sync::CancellationToken,
 };
 
-fn test_config(validator_identity: Option<Keypair>) -> ConnectionWorkersSchedulerConfig {
+fn test_config(identity: Option<Keypair>) -> ConnectionWorkersSchedulerConfig {
     ConnectionWorkersSchedulerConfig {
         bind: SocketAddr::new(Ipv4Addr::new(127, 0, 0, 1).into(), 0),
-        stake_identity: validator_identity,
+        identity,
         num_connections: 1,
         skip_check_transaction_age: false,
         // At the moment we have only one strategy to send transactions: we try
@@ -63,9 +65,9 @@ fn test_config(validator_identity: Option<Keypair>) -> ConnectionWorkersSchedule
 async fn setup_connection_worker_scheduler(
     tpu_address: SocketAddr,
     transaction_receiver: Receiver<TransactionBatch>,
-    validator_identity: Option<Keypair>,
+    identity: Option<Keypair>,
 ) -> (
-    JoinHandle<Result<SendTransactionStatsPerAddr, ConnectionWorkersSchedulerError>>,
+    JoinHandle<Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError>>,
     CancellationToken,
 ) {
     let json_rpc_url = "http://127.0.0.1:8899";
@@ -82,7 +84,7 @@ async fn setup_connection_worker_scheduler(
         .expect("Leader updates was successfully created");
 
     let cancel = CancellationToken::new();
-    let config = test_config(validator_identity);
+    let config = test_config(identity);
     let scheduler = tokio::spawn(ConnectionWorkersScheduler::run(
         config,
         leader_updater,
@@ -95,10 +97,10 @@ async fn setup_connection_worker_scheduler(
 
 async fn join_scheduler(
     scheduler_handle: JoinHandle<
-        Result<SendTransactionStatsPerAddr, ConnectionWorkersSchedulerError>,
+        Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError>,
     >,
 ) -> SendTransactionStatsNonAtomic {
-    let stats_per_ip = scheduler_handle
+    let (stats_per_ip, _) = scheduler_handle
         .await
         .unwrap()
         .expect("Scheduler should stop successfully.");
@@ -400,8 +402,8 @@ async fn test_connection_pruned_and_reopened() {
 /// connection and verify that all the txs has been received.
 #[tokio::test]
 async fn test_staked_connection() {
-    let validator_identity = Keypair::new();
-    let stakes = HashMap::from([(validator_identity.pubkey(), 100_000)]);
+    let identity = Keypair::new();
+    let stakes = HashMap::from([(identity.pubkey(), 100_000)]);
     let staked_nodes = StakedNodes::new(Arc::new(stakes), HashMap::<Pubkey, u64>::default());
 
     let SpawnTestServerResult {
@@ -432,8 +434,7 @@ async fn test_staked_connection() {
     } = spawn_tx_sender(tx_size, expected_num_txs, Duration::from_millis(100));
 
     let (scheduler_handle, _scheduler_cancel) =
-        setup_connection_worker_scheduler(server_address, tx_receiver, Some(validator_identity))
-            .await;
+        setup_connection_worker_scheduler(server_address, tx_receiver, Some(identity)).await;
 
     // Check results
     let actual_num_packets = count_received_packets_for(receiver, tx_size, TEST_MAX_TIME).await;
@@ -533,7 +534,7 @@ async fn test_no_host() {
 
     // While attempting to establish a connection with a nonexistent host, we fill the worker's
     // channel.
-    let stats = scheduler_handle
+    let (stats, _) = scheduler_handle
         .await
         .expect("Scheduler should stop successfully")
         .expect("Scheduler execution was successful");