add token bucket limiter on peer conn recv (#1842)

We should limit peer conn recv to make sure we don't recv too much from peers.
This commit is contained in:
KKRainbow
2026-01-29 16:12:26 +08:00
committed by GitHub
parent ccc684a9ab
commit ffe5644ddc
5 changed files with 54 additions and 8 deletions
+25 -3
View File
@@ -40,7 +40,7 @@ use crate::{
peers::peer_session::{PeerSessionStore, SessionKey, UpsertResponderSessionReturn},
proto::{
api::instance::{PeerConnInfo, PeerConnStats},
common::{SecureModeConfig, TunnelInfo},
common::{LimiterConfig, SecureModeConfig, TunnelInfo},
peer_rpc::{
HandshakeRequest, PeerConnNoiseMsg1Pb, PeerConnNoiseMsg2Pb, PeerConnNoiseMsg3Pb,
PeerConnSessionActionPb, SecureAuthLevel,
@@ -1171,6 +1171,23 @@ impl PeerConn {
};
self.counters.store(Some(Arc::new(counters)));
let is_foreign_network = conn_info_for_instrument.network_name
!= self.global_ctx.get_network_identity().network_name;
let recv_limiter = if is_foreign_network {
let relay_network_bps_limit = self.global_ctx.get_flags().foreign_relay_bps_limit;
let limiter_config = LimiterConfig {
burst_rate: None,
bps: Some(relay_network_bps_limit),
fill_duration_ms: None,
};
Some(self.global_ctx.token_bucket_manager().get_or_create(
&format!("{}:recv", conn_info_for_instrument.network_name),
limiter_config.into(),
))
} else {
None
};
let counters = self.counters.load_full().unwrap();
self.tasks.spawn(
@@ -1185,8 +1202,9 @@ impl PeerConn {
}
let mut zc_packet = ret.unwrap();
let buf_len = zc_packet.buf_len() as u64;
counters.traffic_rx_bytes.add(zc_packet.buf_len() as u64);
counters.traffic_rx_bytes.add(buf_len);
counters.traffic_rx_packets.inc();
let Some(peer_mgr_hdr) = zc_packet.mut_peer_manager_header() else {
@@ -1194,7 +1212,7 @@ impl PeerConn {
"unexpected packet: {:?}, cannot decode peer manager hdr",
zc_packet
);
continue;
break;
};
if peer_mgr_hdr.packet_type == PacketType::Ping as u8 {
@@ -1209,6 +1227,10 @@ impl PeerConn {
} else if sender.send(zc_packet).await.is_err() {
break;
}
if let Some(limiter) = recv_limiter.as_ref() {
limiter.consume(buf_len).await;
}
}
tracing::info!("end recving peer conn packet");