Rl flow

Cargo.toml

@@ -53,6 +53,7 @@ exponential-backoff = "1"
 fdlimit = "0.3.0"
 futures-util = "0.3"
 getip = "0.1"
+governor = { version = "0.10.0" }
 hex = "0.4"
 humantime = "2"
 indexmap = "2.2"

network/Cargo.toml

@@ -27,6 +27,7 @@ ed25519 = { workspace = true, features = ["alloc", "pkcs8"] }
 everscale-crypto = { workspace = true }
 exponential-backoff = { workspace = true }
 futures-util = { workspace = true, features = ["sink"] }
+governor = { workspace = true, features = ["std", "jitter"] }
 hex = { workspace = true }
 indexmap = { workspace = true }
 metrics = { workspace = true }
@@ -48,7 +49,7 @@ tokio-util = { workspace = true, features = ["time"] }
 tracing = { workspace = true }
 
 # local deps
-tycho-util = { workspace = true }
+tycho-util = { workspace = true, features = ["governor"] }
 
 [dev-dependencies]
 clap = { workspace = true, features = ["derive"] }

network/src/network/config.rs

@@ -73,6 +73,10 @@ pub struct NetworkConfig {
 
     /// Default: disabled.
     pub connection_metrics: Option<ConnectionMetricsLevel>,
+
+    /// Default: 250 mbit/s
+    /// Limits all outgoing connections to the same peer to this value.
+    pub bandwidth_cap: bytesize::ByteSize,
 }
 
 impl Default for NetworkConfig {
@@ -93,6 +97,7 @@ impl Default for NetworkConfig {
             shutdown_idle_timeout: Duration::from_secs(60),
             enable_0rtt: false,
             connection_metrics: None,
+            bandwidth_cap: bytesize::ByteSize::mb(250 / 8), // in megabits per second
         }
     }
 }

network/src/network/peer.rs

@@ -1,7 +1,10 @@
 use std::sync::Arc;
 
 use anyhow::Result;
+use governor::clock::QuantaClock;
+use governor::state::{InMemoryState, NotKeyed};
 use tokio_util::codec::{FramedRead, FramedWrite};
+use tycho_util::io::ratelimit::{RatelimitAsyncWriteExt, RatelimitedWriter};
 use tycho_util::metrics::{GaugeGuard, HistogramGuard};
 
 use crate::network::config::NetworkConfig;
@@ -25,11 +28,18 @@ const METRIC_OUT_MESSAGES: &str = "tycho_net_out_messages";
 pub struct Peer {
     connection: Connection,
     config: Arc<NetworkConfig>,
+    rate_limiter: Arc<governor::RateLimiter<NotKeyed, InMemoryState, QuantaClock>>,
 }
 
 impl Peer {
     pub(crate) fn new(connection: Connection, config: Arc<NetworkConfig>) -> Self {
-        Self { connection, config }
+        Self {
+            connection,
+            rate_limiter: Arc::new(
+                tycho_util::io::ratelimit::rate_limiter(config.bandwidth_cap).unwrap(),
+            ),
+            config,
+        }
     }
 
     pub fn peer_id(&self) -> &PeerId {
@@ -42,12 +52,13 @@ impl Peer {
         let _histogram = HistogramGuard::begin(METRIC_OUT_QUERIES_TIME);
 
         let (send_stream, recv_stream) = self.connection.open_bi().await?;
+        let send_stream = send_stream.ratelimit_write(&self.rate_limiter);
+        let send_stream = std::pin::pin!(send_stream);
         let mut send_stream = FramedWrite::new(send_stream, make_codec(&self.config));
         let mut recv_stream = FramedRead::new(recv_stream, make_codec(&self.config));
 
         send_request(&mut send_stream, request).await?;
-        send_stream.get_mut().finish()?;
-
+        RatelimitedWriter::get_mut(send_stream.into_inner()).finish()?;
         recv_response(&mut recv_stream).await.map_err(Into::into)
     }
 

util/Cargo.toml

@@ -15,10 +15,12 @@ ahash = { workspace = true }
 anyhow = { workspace = true }
 base64 = { workspace = true }
 bytes = { workspace = true }
+bytesize = { workspace = true, optional = true }
 castaway = { workspace = true }
 dashmap = { workspace = true }
 futures-util = { workspace = true }
 getip = { workspace = true, optional = true }
+governor = { workspace = true, optional = true, features = ["std", "jitter"] }
 humantime = { workspace = true }
 libc = { workspace = true, optional = true }
 metrics = { workspace = true }
@@ -39,11 +41,12 @@ zstd-sys = { workspace = true }
 
 metrics-exporter-prometheus = { workspace = true, optional = true }
 tikv-jemalloc-ctl = { workspace = true, optional = true }
+pin-project-lite = { workspace = true, optional = true }
 
 [dev-dependencies]
 criterion = "0.5.1"
 tempfile = { workspace = true }
-tokio = { workspace = true, features = ["time", "sync", "rt-multi-thread", "macros"] }
+tokio = { workspace = true, features = ["time", "sync", "rt-multi-thread", "macros", "io-util"] }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
 
 [features]
@@ -58,6 +61,11 @@ cli = [
     "metrics-exporter-prometheus",
     "tikv-jemalloc-ctl",
 ]
+governor = [
+    "dep:governor",
+    "dep:pin-project-lite",
+    "dep:bytesize",
+]
 
 [[bench]]
 name = "p2"

util/src/io.rs → util/src/io/mod.rs

@@ -1,3 +1,5 @@
+#[cfg(feature = "governor")]
+pub mod ratelimit;
 // TODO: Extend with required methods.
 pub trait ByteOrderRead {
     fn read_be_uint(&mut self, bytes: usize) -> std::io::Result<u64>;