From 92ebf0f80c14acda373e7580a84b3a248339f6bf Mon Sep 17 00:00:00 2001
From: Lijun Wang <83639177+lijunwangs@users.noreply.github.com>
Date: Wed, 10 Apr 2024 16:13:41 -0700
Subject: [PATCH] Treat super low staked as unstaked in streamer QOS (#701)

* Treat super low staked with QOS of unstaked

* simplify

* address some comment from Pankaj
---
 streamer/src/nonblocking/quic.rs            | 15 +++++++++++----
 streamer/src/nonblocking/stream_throttle.rs |  4 ++--
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/streamer/src/nonblocking/quic.rs b/streamer/src/nonblocking/quic.rs
index ddf4005c432e3d..55e2b65d2c76d4 100644
--- a/streamer/src/nonblocking/quic.rs
+++ b/streamer/src/nonblocking/quic.rs
@@ -1,7 +1,8 @@
 use {
     crate::{
         nonblocking::stream_throttle::{
-            ConnectionStreamCounter, StakedStreamLoadEMA, STREAM_STOP_CODE_THROTTLING,
+            ConnectionStreamCounter, StakedStreamLoadEMA, MAX_STREAMS_PER_MS,
+            STREAM_STOP_CODE_THROTTLING, STREAM_THROTTLING_INTERVAL_MS,
         },
         quic::{configure_server, QuicServerError, StreamStats},
         streamer::StakedNodes,
@@ -500,10 +501,16 @@ async fn setup_connection(
                         stats.clone(),
                     ),
                     |(pubkey, stake, total_stake, max_stake, min_stake)| {
-                        let peer_type = if stake > 0 {
-                            ConnectionPeerType::Staked(stake)
-                        } else {
+                        // The heuristic is that the stake should be large engouh to have 1 stream pass throuh within one throttle
+                        // interval during which we allow max (MAX_STREAMS_PER_MS * STREAM_THROTTLING_INTERVAL_MS) streams.
+                        let min_stake_ratio =
+                            1_f64 / (MAX_STREAMS_PER_MS * STREAM_THROTTLING_INTERVAL_MS) as f64;
+                        let stake_ratio = stake as f64 / total_stake as f64;
+                        let peer_type = if stake_ratio < min_stake_ratio {
+                            // If it is a staked connection with ultra low stake ratio, treat it as unstaked.
                             ConnectionPeerType::Unstaked
+                        } else {
+                            ConnectionPeerType::Staked(stake)
                         };
                         NewConnectionHandlerParams {
                             packet_sender,
diff --git a/streamer/src/nonblocking/stream_throttle.rs b/streamer/src/nonblocking/stream_throttle.rs
index 8fda63271816e6..e3da2be90ddfdc 100644
--- a/streamer/src/nonblocking/stream_throttle.rs
+++ b/streamer/src/nonblocking/stream_throttle.rs
@@ -12,9 +12,9 @@ use {
 };
 
 /// Limit to 250K PPS
-const MAX_STREAMS_PER_MS: u64 = 250;
+pub const MAX_STREAMS_PER_MS: u64 = 250;
 const MAX_UNSTAKED_STREAMS_PERCENT: u64 = 20;
-const STREAM_THROTTLING_INTERVAL_MS: u64 = 100;
+pub const STREAM_THROTTLING_INTERVAL_MS: u64 = 100;
 pub const STREAM_STOP_CODE_THROTTLING: u32 = 15;
 const STREAM_LOAD_EMA_INTERVAL_MS: u64 = 5;
 const STREAM_LOAD_EMA_INTERVAL_COUNT: u64 = 10;