Do not reserve QUIC stream capacity for unstaked client on forward port (#34779)

This commit is contained in:
Pankaj Garg 2024-01-16 16:35:29 -08:00 committed by GitHub
parent 3fa44e6fbe
commit 9edf65b877
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 204 additions and 113 deletions

View File

@ -1,7 +1,7 @@
use { use {
crate::{ crate::{
nonblocking::stream_throttle::{ nonblocking::stream_throttle::{
self, ConnectionStreamCounter, StakedStreamLoadEMA, STREAM_STOP_CODE_THROTTLING, ConnectionStreamCounter, StakedStreamLoadEMA, STREAM_STOP_CODE_THROTTLING,
}, },
quic::{configure_server, QuicServerError, StreamStats}, quic::{configure_server, QuicServerError, StreamStats},
streamer::StakedNodes, streamer::StakedNodes,
@ -90,7 +90,7 @@ pub enum ConnectionPeerType {
} }
impl ConnectionPeerType { impl ConnectionPeerType {
fn is_staked(&self) -> bool { pub(crate) fn is_staked(&self) -> bool {
matches!(self, ConnectionPeerType::Staked(_)) matches!(self, ConnectionPeerType::Staked(_))
} }
} }
@ -156,7 +156,10 @@ async fn run_server(
let mut last_datapoint = Instant::now(); let mut last_datapoint = Instant::now();
let unstaked_connection_table: Arc<Mutex<ConnectionTable>> = let unstaked_connection_table: Arc<Mutex<ConnectionTable>> =
Arc::new(Mutex::new(ConnectionTable::new())); Arc::new(Mutex::new(ConnectionTable::new()));
let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(stats.clone())); let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
max_unstaked_connections > 0,
stats.clone(),
));
let staked_connection_table: Arc<Mutex<ConnectionTable>> = let staked_connection_table: Arc<Mutex<ConnectionTable>> =
Arc::new(Mutex::new(ConnectionTable::new())); Arc::new(Mutex::new(ConnectionTable::new()));
let (sender, receiver) = async_unbounded(); let (sender, receiver) = async_unbounded();
@ -718,25 +721,17 @@ async fn handle_connection(
); );
let stable_id = connection.stable_id(); let stable_id = connection.stable_id();
stats.total_connections.fetch_add(1, Ordering::Relaxed); stats.total_connections.fetch_add(1, Ordering::Relaxed);
let mut max_streams_per_throttling_interval =
stream_throttle::max_streams_for_connection_in_throttling_duration(
params.peer_type,
params.total_stake,
stream_load_ema.clone(),
);
while !stream_exit.load(Ordering::Relaxed) { while !stream_exit.load(Ordering::Relaxed) {
if let Ok(stream) = if let Ok(stream) =
tokio::time::timeout(WAIT_FOR_STREAM_TIMEOUT, connection.accept_uni()).await tokio::time::timeout(WAIT_FOR_STREAM_TIMEOUT, connection.accept_uni()).await
{ {
match stream { match stream {
Ok(mut stream) => { Ok(mut stream) => {
if let ConnectionPeerType::Staked(peer_stake) = params.peer_type { let max_streams_per_throttling_interval = stream_load_ema
max_streams_per_throttling_interval = stream_load_ema .available_load_capacity_in_throttling_duration(
.available_load_capacity_in_throttling_duration( params.peer_type,
peer_stake, params.total_stake,
params.total_stake, );
);
}
stream_counter.reset_throttling_params_if_needed(); stream_counter.reset_throttling_params_if_needed();
if stream_counter.stream_count.load(Ordering::Relaxed) if stream_counter.stream_count.load(Ordering::Relaxed)
@ -746,9 +741,7 @@ async fn handle_connection(
let _ = stream.stop(VarInt::from_u32(STREAM_STOP_CODE_THROTTLING)); let _ = stream.stop(VarInt::from_u32(STREAM_STOP_CODE_THROTTLING));
continue; continue;
} }
if params.peer_type.is_staked() { stream_load_ema.increment_load(params.peer_type);
stream_load_ema.increment_load();
}
stream_counter.stream_count.fetch_add(1, Ordering::Relaxed); stream_counter.stream_count.fetch_add(1, Ordering::Relaxed);
stats.total_streams.fetch_add(1, Ordering::Relaxed); stats.total_streams.fetch_add(1, Ordering::Relaxed);
stats.total_new_streams.fetch_add(1, Ordering::Relaxed); stats.total_new_streams.fetch_add(1, Ordering::Relaxed);
@ -797,9 +790,7 @@ async fn handle_connection(
} }
} }
stats.total_streams.fetch_sub(1, Ordering::Relaxed); stats.total_streams.fetch_sub(1, Ordering::Relaxed);
if params.peer_type.is_staked() { stream_load_ema.update_ema_if_needed();
stream_load_ema.update_ema_if_needed();
}
}); });
} }
Err(e) => { Err(e) => {

View File

@ -21,22 +21,52 @@ const STREAM_THROTTLING_INTERVAL_MS: u64 = 100;
pub const STREAM_STOP_CODE_THROTTLING: u32 = 15; pub const STREAM_STOP_CODE_THROTTLING: u32 = 15;
const STREAM_LOAD_EMA_INTERVAL_MS: u64 = 5; const STREAM_LOAD_EMA_INTERVAL_MS: u64 = 5;
const STREAM_LOAD_EMA_INTERVAL_COUNT: u64 = 10; const STREAM_LOAD_EMA_INTERVAL_COUNT: u64 = 10;
const MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION: u64 = 8; const EMA_WINDOW_MS: u64 = STREAM_LOAD_EMA_INTERVAL_MS * STREAM_LOAD_EMA_INTERVAL_COUNT;
pub(crate) struct StakedStreamLoadEMA { pub(crate) struct StakedStreamLoadEMA {
current_load_ema: AtomicU64, current_load_ema: AtomicU64,
load_in_recent_interval: AtomicU64, load_in_recent_interval: AtomicU64,
last_update: RwLock<Instant>, last_update: RwLock<Instant>,
stats: Arc<StreamStats>, stats: Arc<StreamStats>,
// Maximum number of streams for a staked connection in EMA window
// Note: EMA window can be different than stream throttling window. EMA is being calculated
// specifically for staked connections. Unstaked connections have fixed limit on
// stream load, which is tracked by `max_unstaked_load_in_throttling_window` field.
max_staked_load_in_ema_window: u64,
// Maximum number of streams for an unstaked connection in stream throttling window
max_unstaked_load_in_throttling_window: u64,
} }
impl StakedStreamLoadEMA { impl StakedStreamLoadEMA {
pub(crate) fn new(stats: Arc<StreamStats>) -> Self { pub(crate) fn new(allow_unstaked_streams: bool, stats: Arc<StreamStats>) -> Self {
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))
* EMA_WINDOW_MS
} else {
MAX_STREAMS_PER_MS * EMA_WINDOW_MS
};
let max_num_unstaked_connections =
u64::try_from(MAX_UNSTAKED_CONNECTIONS).unwrap_or_else(|_| {
error!(
"Failed to convert maximum number of unstaked connections {} to u64.",
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);
Self { Self {
current_load_ema: AtomicU64::default(), current_load_ema: AtomicU64::default(),
load_in_recent_interval: AtomicU64::default(), load_in_recent_interval: AtomicU64::default(),
last_update: RwLock::new(Instant::now()), last_update: RwLock::new(Instant::now()),
stats, stats,
max_staked_load_in_ema_window,
max_unstaked_load_in_throttling_window,
} }
} }
@ -106,76 +136,56 @@ impl StakedStreamLoadEMA {
} }
} }
pub(crate) fn increment_load(&self) { pub(crate) fn increment_load(&self, peer_type: ConnectionPeerType) {
self.load_in_recent_interval.fetch_add(1, Ordering::Relaxed); if peer_type.is_staked() {
self.load_in_recent_interval.fetch_add(1, Ordering::Relaxed);
}
self.update_ema_if_needed(); self.update_ema_if_needed();
} }
pub(crate) fn available_load_capacity_in_throttling_duration( pub(crate) fn available_load_capacity_in_throttling_duration(
&self, &self,
stake: u64, peer_type: ConnectionPeerType,
total_stake: u64, total_stake: u64,
) -> u64 { ) -> u64 {
let ema_window_ms = STREAM_LOAD_EMA_INTERVAL_MS * STREAM_LOAD_EMA_INTERVAL_COUNT; match peer_type {
let max_load_in_ema_window = u128::from( ConnectionPeerType::Unstaked => self.max_unstaked_load_in_throttling_window,
(MAX_STREAMS_PER_MS ConnectionPeerType::Staked(stake) => {
- Percentage::from(MAX_UNSTAKED_STREAMS_PERCENT).apply_to(MAX_STREAMS_PER_MS)) // If the current load is low, cap it to 25% of max_load.
* ema_window_ms, let current_load = u128::from(cmp::max(
); self.current_load_ema.load(Ordering::Relaxed),
self.max_staked_load_in_ema_window / 4,
));
// If the current load is low, cap it to 25% of max_load. // Formula is (max_load ^ 2 / current_load) * (stake / total_stake)
let current_load = cmp::max( let capacity_in_ema_window = (u128::from(self.max_staked_load_in_ema_window)
u128::from(self.current_load_ema.load(Ordering::Relaxed)), * u128::from(self.max_staked_load_in_ema_window)
max_load_in_ema_window / 4, * u128::from(stake))
); / (current_load * u128::from(total_stake));
// Formula is (max_load ^ 2 / current_load) * (stake / total_stake) let calculated_capacity = capacity_in_ema_window
let capacity_in_ema_window = * u128::from(STREAM_THROTTLING_INTERVAL_MS)
(max_load_in_ema_window * max_load_in_ema_window * u128::from(stake)) / u128::from(EMA_WINDOW_MS);
/ (current_load * u128::from(total_stake)); let calculated_capacity = u64::try_from(calculated_capacity).unwrap_or_else(|_| {
let calculated_capacity = capacity_in_ema_window
* u128::from(STREAM_THROTTLING_INTERVAL_MS)
/ u128::from(ema_window_ms);
let calculated_capacity = u64::try_from(calculated_capacity).unwrap_or_else(|_| {
error!(
"Failed to convert stream capacity {} to u64. Using minimum load capacity",
calculated_capacity
);
self.stats
.stream_load_capacity_overflow
.fetch_add(1, Ordering::Relaxed);
MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION
});
cmp::max(
calculated_capacity,
MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION,
)
}
}
pub(crate) fn max_streams_for_connection_in_throttling_duration(
peer_type: ConnectionPeerType,
total_stake: u64,
ema_load: Arc<StakedStreamLoadEMA>,
) -> u64 {
match peer_type {
ConnectionPeerType::Unstaked => {
let max_num_connections =
u64::try_from(MAX_UNSTAKED_CONNECTIONS).unwrap_or_else(|_| {
error!( error!(
"Failed to convert maximum number of unstaked connections {} to u64.", "Failed to convert stream capacity {} to u64. Using minimum load capacity",
MAX_UNSTAKED_CONNECTIONS calculated_capacity
); );
500 self.stats
.stream_load_capacity_overflow
.fetch_add(1, Ordering::Relaxed);
self.max_unstaked_load_in_throttling_window
.saturating_add(1)
}); });
Percentage::from(MAX_UNSTAKED_STREAMS_PERCENT)
.apply_to(MAX_STREAMS_PER_MS * STREAM_THROTTLING_INTERVAL_MS) // 1 is added to `max_unstaked_load_in_throttling_window` to guarantee that staked
.saturating_div(max_num_connections) // clients get at least 1 more number of streams than unstaked connections.
} cmp::max(
ConnectionPeerType::Staked(stake) => { calculated_capacity,
ema_load.available_load_capacity_in_throttling_duration(stake, total_stake) self.max_unstaked_load_in_throttling_window
.saturating_add(1),
)
}
} }
} }
} }
@ -215,14 +225,7 @@ impl ConnectionStreamCounter {
pub mod test { pub mod test {
use { use {
super::*, super::*,
crate::{ crate::{nonblocking::stream_throttle::STREAM_LOAD_EMA_INTERVAL_MS, quic::StreamStats},
nonblocking::stream_throttle::{
max_streams_for_connection_in_throttling_duration,
MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION,
STREAM_LOAD_EMA_INTERVAL_MS,
},
quic::StreamStats,
},
std::{ std::{
sync::{atomic::Ordering, Arc}, sync::{atomic::Ordering, Arc},
time::{Duration, Instant}, time::{Duration, Instant},
@ -231,20 +234,26 @@ pub mod test {
#[test] #[test]
fn test_max_streams_for_unstaked_connection() { fn test_max_streams_for_unstaked_connection() {
let load_ema = Arc::new(StakedStreamLoadEMA::new(Arc::new(StreamStats::default()))); let load_ema = Arc::new(StakedStreamLoadEMA::new(
true,
Arc::new(StreamStats::default()),
));
// 25K packets per ms * 20% / 500 max unstaked connections // 25K packets per ms * 20% / 500 max unstaked connections
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Unstaked, ConnectionPeerType::Unstaked,
10000, 10000,
load_ema.clone(),
), ),
10 10
); );
} }
#[test] #[test]
fn test_max_streams_for_staked_connection() { fn test_max_streams_for_staked_connection() {
let load_ema = Arc::new(StakedStreamLoadEMA::new(Arc::new(StreamStats::default()))); let load_ema = Arc::new(StakedStreamLoadEMA::new(
true,
Arc::new(StreamStats::default()),
));
// EMA load is used for staked connections to calculate max number of allowed streams. // EMA load is used for staked connections to calculate max number of allowed streams.
// EMA window = 5ms interval * 10 intervals = 50ms // EMA window = 5ms interval * 10 intervals = 50ms
@ -258,10 +267,9 @@ pub mod test {
// ema_load = 10K, stake = 15, total_stake = 10K // ema_load = 10K, stake = 15, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 10K) * 15 / 10K = 30 // max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 10K) * 15 / 10K = 30
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15), ConnectionPeerType::Staked(15),
10000, 10000,
load_ema.clone(),
), ),
30 30
); );
@ -269,10 +277,9 @@ pub mod test {
// ema_load = 10K, stake = 1K, total_stake = 10K // ema_load = 10K, stake = 1K, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 10K) * 1K / 10K = 2K // max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 10K) * 1K / 10K = 2K
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000), ConnectionPeerType::Staked(1000),
10000, 10000,
load_ema.clone(),
), ),
2000 2000
); );
@ -281,10 +288,9 @@ pub mod test {
// ema_load = 2.5K, stake = 15, total_stake = 10K // ema_load = 2.5K, stake = 15, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 2.5K) * 15 / 10K = 120 // max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 2.5K) * 15 / 10K = 120
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15), ConnectionPeerType::Staked(15),
10000, 10000,
load_ema.clone(),
), ),
120 120
); );
@ -292,10 +298,9 @@ pub mod test {
// ema_load = 2.5K, stake = 1K, total_stake = 10K // ema_load = 2.5K, stake = 1K, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 2.5K) * 1K / 10K = 8000 // max_streams in 100ms (throttling window) = 2 * ((10K * 10K) / 2.5K) * 1K / 10K = 8000
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000), ConnectionPeerType::Staked(1000),
10000, 10000,
load_ema.clone(),
), ),
8000 8000
); );
@ -305,39 +310,128 @@ pub mod test {
load_ema.current_load_ema.store(2000, Ordering::Relaxed); load_ema.current_load_ema.store(2000, Ordering::Relaxed);
// function = ((10K * 10K) / 25% of 10K) * stake / total_stake // function = ((10K * 10K) / 25% of 10K) * stake / total_stake
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15), ConnectionPeerType::Staked(15),
10000, 10000,
load_ema.clone(),
), ),
120 120
); );
// function = ((10K * 10K) / 25% of 10K) * stake / total_stake // function = ((10K * 10K) / 25% of 10K) * stake / total_stake
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000), ConnectionPeerType::Staked(1000),
10000, 10000,
load_ema.clone(),
), ),
8000 8000
); );
// At 1/40000 stake weight, and minimum load, it should still allow // At 1/40000 stake weight, and minimum load, it should still allow
// MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION of streams. // max_unstaked_load_in_throttling_window + 1 streams.
assert_eq!( assert_eq!(
max_streams_for_connection_in_throttling_duration( load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1), ConnectionPeerType::Staked(1),
40000, 40000,
load_ema.clone(),
), ),
MIN_STREAMS_PER_THROTTLING_INTERVAL_FOR_STAKED_CONNECTION load_ema
.max_unstaked_load_in_throttling_window
.saturating_add(1)
);
}
#[test]
fn test_max_streams_for_staked_connection_with_no_unstaked_connections() {
let load_ema = Arc::new(StakedStreamLoadEMA::new(
false,
Arc::new(StreamStats::default()),
));
// EMA load is used for staked connections to calculate max number of allowed streams.
// EMA window = 5ms interval * 10 intervals = 50ms
// max streams per window = 250K streams/sec = 12.5K per 50ms
// max_streams in 50ms = ((12.5K * 12.5K) / ema_load) * stake / total_stake
//
// Stream throttling window is 100ms. So it'll double the amount of max streams.
// max_streams in 100ms (throttling window) = 2 * ((12.5K * 12.5K) / ema_load) * stake / total_stake
load_ema.current_load_ema.store(10000, Ordering::Relaxed);
// ema_load = 10K, stake = 15, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((12.5K * 12.5K) / 10K) * 15 / 10K = 46.875
assert!(
(46u64..=47).contains(&load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15),
10000
))
);
// ema_load = 10K, stake = 1K, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((12.5K * 12.5K) / 10K) * 1K / 10K = 3125
assert!((3124u64..=3125).contains(
&load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000),
10000
)
));
load_ema.current_load_ema.store(5000, Ordering::Relaxed);
// ema_load = 5K, stake = 15, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((12.5K * 12.5K) / 5K) * 15 / 10K = 93.75
assert!(
(92u64..=94).contains(&load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15),
10000
))
);
// ema_load = 5K, stake = 1K, total_stake = 10K
// max_streams in 100ms (throttling window) = 2 * ((12.5K * 12.5K) / 5K) * 1K / 10K = 6250
assert!((6248u64..=6250).contains(
&load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000),
10000
)
));
// At 2000, the load is less than 25% of max_load (12.5K).
// Test that we cap it to 25%, yielding the same result as if load was 12.5K/4.
load_ema.current_load_ema.store(2000, Ordering::Relaxed);
// function = ((10K * 10K) / 25% of 12.5K) * stake / total_stake
assert_eq!(
load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(15),
10000
),
150
);
// function = ((12.5K * 12.5K) / 25% of 12.5K) * stake / total_stake
assert_eq!(
load_ema.available_load_capacity_in_throttling_duration(
ConnectionPeerType::Staked(1000),
10000
),
10000
);
// At 1/40000 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
),
load_ema
.max_unstaked_load_in_throttling_window
.saturating_add(1)
); );
} }
#[test] #[test]
fn test_update_ema() { fn test_update_ema() {
let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(Arc::new(StreamStats::default()))); let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
true,
Arc::new(StreamStats::default()),
));
stream_load_ema stream_load_ema
.load_in_recent_interval .load_in_recent_interval
.store(2500, Ordering::Relaxed); .store(2500, Ordering::Relaxed);
@ -362,7 +456,10 @@ pub mod test {
#[test] #[test]
fn test_update_ema_missing_interval() { fn test_update_ema_missing_interval() {
let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(Arc::new(StreamStats::default()))); let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
true,
Arc::new(StreamStats::default()),
));
stream_load_ema stream_load_ema
.load_in_recent_interval .load_in_recent_interval
.store(2500, Ordering::Relaxed); .store(2500, Ordering::Relaxed);
@ -378,7 +475,10 @@ pub mod test {
#[test] #[test]
fn test_update_ema_if_needed() { fn test_update_ema_if_needed() {
let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(Arc::new(StreamStats::default()))); let stream_load_ema = Arc::new(StakedStreamLoadEMA::new(
true,
Arc::new(StreamStats::default()),
));
stream_load_ema stream_load_ema
.load_in_recent_interval .load_in_recent_interval
.store(2500, Ordering::Relaxed); .store(2500, Ordering::Relaxed);