corrected to not use hardcoded connections count for unstaked

streamer/src/nonblocking/quic.rs

@@ -172,8 +172,8 @@ async fn run_server(
     let unstaked_connection_table: Arc<Mutex<ConnectionTable>> =
         Arc::new(Mutex::new(ConnectionTable::new()));
     let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
-        max_unstaked_connections > 0,
         stats.clone(),
+        max_unstaked_connections,
     ));
     let staked_connection_table: Arc<Mutex<ConnectionTable>> =
         Arc::new(Mutex::new(ConnectionTable::new()));

streamer/src/nonblocking/stream_throttle.rs

@@ -1,8 +1,5 @@
 use {
-    crate::{
-        nonblocking::quic::ConnectionPeerType,
-        quic::{StreamStats, MAX_UNSTAKED_CONNECTIONS},
-    },
+    crate::{nonblocking::quic::ConnectionPeerType, quic::StreamStats},
     percentage::Percentage,
     std::{
         cmp,
@@ -38,7 +35,8 @@ pub(crate) struct StakedStreamLoadEMA {
 }
 
 impl StakedStreamLoadEMA {
-    pub(crate) fn new(allow_unstaked_streams: bool, stats: Arc<StreamStats>) -> Self {
+    pub(crate) fn new(stats: Arc<StreamStats>, max_unstaked_connections: usize) -> Self {
+        let allow_unstaked_streams = max_unstaked_connections > 0;
         let max_staked_load_in_ema_window = if allow_unstaked_streams {
             (MAX_STREAMS_PER_MS
                 - Percentage::from(MAX_UNSTAKED_STREAMS_PERCENT).apply_to(MAX_STREAMS_PER_MS))
@@ -48,17 +46,21 @@ impl StakedStreamLoadEMA {
         };
 
         let max_num_unstaked_connections =
-            u64::try_from(MAX_UNSTAKED_CONNECTIONS).unwrap_or_else(|_| {
+            u64::try_from(max_unstaked_connections).unwrap_or_else(|_| {
                 error!(
                     "Failed to convert maximum number of unstaked connections {} to u64.",
-                    MAX_UNSTAKED_CONNECTIONS
+                    max_unstaked_connections
                 );
                 500
             });
 
-        let max_unstaked_load_in_throttling_window = Percentage::from(MAX_UNSTAKED_STREAMS_PERCENT)
-            .apply_to(MAX_STREAMS_PER_MS * STREAM_THROTTLING_INTERVAL_MS)
-            .saturating_div(max_num_unstaked_connections);
+        let max_unstaked_load_in_throttling_window = if allow_unstaked_streams {
+            Percentage::from(MAX_UNSTAKED_STREAMS_PERCENT)
+                .apply_to(MAX_STREAMS_PER_MS * STREAM_THROTTLING_INTERVAL_MS)
+                .saturating_div(max_num_unstaked_connections)
+        } else {
+            0
+        };
 
         Self {
             current_load_ema: AtomicU64::default(),
@@ -225,7 +227,10 @@ impl ConnectionStreamCounter {
 pub mod test {
     use {
         super::*,
-        crate::{nonblocking::stream_throttle::STREAM_LOAD_EMA_INTERVAL_MS, quic::StreamStats},
+        crate::{
+            nonblocking::stream_throttle::STREAM_LOAD_EMA_INTERVAL_MS,
+            quic::{StreamStats, MAX_UNSTAKED_CONNECTIONS},
+        },
         std::{
             sync::{atomic::Ordering, Arc},
             time::{Duration, Instant},
@@ -235,8 +240,8 @@ pub mod test {
     #[test]
     fn test_max_streams_for_unstaked_connection() {
         let load_ema = Arc::new(StakedStreamLoadEMA::new(
-            true,
             Arc::new(StreamStats::default()),
+            MAX_UNSTAKED_CONNECTIONS,
         ));
         // 25K packets per ms * 20% / 500 max unstaked connections
         assert_eq!(
@@ -251,8 +256,8 @@ pub mod test {
     #[test]
     fn test_max_streams_for_staked_connection() {
         let load_ema = Arc::new(StakedStreamLoadEMA::new(
-            true,
             Arc::new(StreamStats::default()),
+            MAX_UNSTAKED_CONNECTIONS,
         ));
 
         // EMA load is used for staked connections to calculate max number of allowed streams.
@@ -342,8 +347,8 @@ pub mod test {
     #[test]
     fn test_max_streams_for_staked_connection_with_no_unstaked_connections() {
         let load_ema = Arc::new(StakedStreamLoadEMA::new(
-            false,
             Arc::new(StreamStats::default()),
+            0,
         ));
 
         // EMA load is used for staked connections to calculate max number of allowed streams.
@@ -413,12 +418,12 @@ pub mod test {
             10000
         );
 
-        // At 1/40000 stake weight, and minimum load, it should still allow
+        // At 1/400000 stake weight, and minimum load, it should still allow
         // max_unstaked_load_in_throttling_window + 1 streams.
         assert_eq!(
             load_ema.available_load_capacity_in_throttling_duration(
                 ConnectionPeerType::Staked(1),
-                40000
+                400000
             ),
             load_ema
                 .max_unstaked_load_in_throttling_window
@@ -429,8 +434,8 @@ pub mod test {
     #[test]
     fn test_update_ema() {
         let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
-            true,
             Arc::new(StreamStats::default()),
+            MAX_UNSTAKED_CONNECTIONS,
         ));
         stream_load_ema
             .load_in_recent_interval
@@ -457,8 +462,8 @@ pub mod test {
     #[test]
     fn test_update_ema_missing_interval() {
         let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
-            true,
             Arc::new(StreamStats::default()),
+            MAX_UNSTAKED_CONNECTIONS,
         ));
         stream_load_ema
             .load_in_recent_interval
@@ -476,8 +481,8 @@ pub mod test {
     #[test]
     fn test_update_ema_if_needed() {
         let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
-            true,
             Arc::new(StreamStats::default()),
+            MAX_UNSTAKED_CONNECTIONS,
         ));
         stream_load_ema
             .load_in_recent_interval