extract identity stakes logic

quic-forward-proxy-integration-test/tests/quic_proxy_tpu_integrationtest.rs

@@ -120,7 +120,6 @@ pub fn with_1000_transactions_direct() {
 }
 
 // note: this tests are flakes on CI ond also local (see https://mangolana.atlassian.net/browse/MAN-59)
-#[ignore]
 #[test]
 pub fn bench_proxy() {
     configure_logging(true);

quic-forward-proxy/src/lib.rs

@@ -6,7 +6,6 @@ mod outbound;
 pub mod proxy;
 pub mod proxy_request_format;
 mod quic_util;
-mod quinn_auto_reconnect;
 pub mod tls_config_provider_client;
 pub mod tls_config_provider_server;
 pub mod tls_self_signed_pair_generator;

quic-forward-proxy/src/main.rs

@@ -16,7 +16,6 @@ mod outbound;
 pub mod proxy;
 pub mod proxy_request_format;
 pub mod quic_util;
-mod quinn_auto_reconnect;
 pub mod tls_config_provider_client;
 pub mod tls_config_provider_server;
 pub mod tls_self_signed_pair_generator;

quic-forward-proxy/src/outbound/ng_forward.rs

@@ -1,6 +1,6 @@
 use crate::validator_identity::ValidatorIdentity;
 use anyhow::{bail};
-use solana_streamer::nonblocking::quic::{ConnectionPeerType};
+use solana_streamer::nonblocking::quic::{compute_max_allowed_uni_streams, ConnectionPeerType};
 use solana_streamer::tls_certificates::new_self_signed_tls_certificate;
 use std::collections::HashMap;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
@@ -78,10 +78,16 @@ pub async fn ng_forwarder(
             requested_connections.insert(tpu_node.identity_tpunode, tpu_node.tpu_socket_addr);
         }
 
+        let max_uni_stream_connections = compute_max_allowed_uni_streams(
+            identity_stakes.peer_type,
+            identity_stakes.stakes,
+            identity_stakes.total_stakes,
+        );
+
         tpu_connection_manager
             .update_connections(
                 &requested_connections,
-                identity_stakes,
+                max_uni_stream_connections,
                 DataCache::new_for_tests(),
                 QUIC_CONNECTION_PARAMS, // TODO improve
             )

quic-forward-proxy/src/outbound/tpu_connection_manager.rs

@@ -148,17 +148,12 @@ impl ActiveConnection {
     pub fn start_listening(
         &self,
         broadcast_receiver: Receiver<BroadcastMessage>,
-        identity_stakes: IdentityStakesData,
+        max_uni_stream_connections: usize,
     ) {
         let addr = self.tpu_address;
 
         let max_number_of_connections = self.connection_parameters.max_number_of_connections;
 
-        let max_uni_stream_connections = compute_max_allowed_uni_streams(
-            identity_stakes.peer_type,
-            identity_stakes.stakes,
-            identity_stakes.total_stakes,
-        );
         let exit_signal = self.exit_signal.clone();
 
         let connection_pool = QuicConnectionPool::new(
@@ -232,7 +227,7 @@ impl TpuConnectionManager {
     pub async fn update_connections(
         &self,
         requested_connections: &HashMap<Pubkey, SocketAddr>,
-        identity_stakes: IdentityStakesData,
+        max_uni_stream_connections: usize,
         data_cache: DataCache,
         connection_parameters: QuicConnectionParameters,
     ) {
@@ -252,7 +247,7 @@ impl TpuConnectionManager {
 
             let broadcast_receiver = self.broadcast_sender.subscribe();
 
-            active_connection.start_listening(broadcast_receiver, identity_stakes);
+            active_connection.start_listening(broadcast_receiver, max_uni_stream_connections);
             self.identity_to_active_connection.insert(
                 *tpu_identity,
                 Arc::new(active_connection),