Bound Direct relay buffers with an adaptive global memory envelope

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/config/defaults.rs

@@ -39,6 +39,8 @@ const DEFAULT_ME_QUOTA_SOFT_OVERSHOOT_BYTES: u64 = 64 * 1024;
 const DEFAULT_ME_D2C_FRAME_BUF_SHRINK_THRESHOLD_BYTES: usize = 256 * 1024;
 const DEFAULT_DIRECT_RELAY_COPY_BUF_C2S_BYTES: usize = 64 * 1024;
 const DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES: usize = 256 * 1024;
+pub(crate) const DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES: usize = 4 * 1024;
+const DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES: usize = 0;
 const DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE: u8 = 3;
 const DEFAULT_ME_HEALTH_INTERVAL_MS_UNHEALTHY: u64 = 1000;
 const DEFAULT_ME_HEALTH_INTERVAL_MS_HEALTHY: u64 = 3000;
@@ -515,6 +517,10 @@ pub(crate) fn default_direct_relay_copy_buf_s2c_bytes() -> usize {
     DEFAULT_DIRECT_RELAY_COPY_BUF_S2C_BYTES
 }
 
+pub(crate) fn default_direct_relay_buffer_budget_max_bytes() -> usize {
+    DEFAULT_DIRECT_RELAY_BUFFER_BUDGET_MAX_BYTES
+}
+
 pub(crate) fn default_me_writer_pick_sample_size() -> u8 {
     DEFAULT_ME_WRITER_PICK_SAMPLE_SIZE
 }

src/config/load.rs

@@ -42,6 +42,8 @@ const MAX_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 16_384;
 const MAX_ME_WRITER_BYTE_BUDGET_BYTES: usize = 256 * 1024 * 1024;
 const MAX_ME_ROUTE_CHANNEL_CAPACITY: usize = 8_192;
 const MAX_ME_C2ME_CHANNEL_CAPACITY: usize = 8_192;
+const MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 16 * 1024 * 1024;
+const MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES: usize = 2 * 1024 * 1024 * 1024;
 const MIN_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024;
 const MAX_MAX_CLIENT_FRAME_BYTES: usize = 16 * 1024 * 1024;
 const MAX_API_REQUEST_BODY_LIMIT_BYTES: usize = 1024 * 1024;
@@ -616,6 +618,24 @@ impl ProxyConfig {
             ));
         }
 
+        if config.general.direct_relay_buffer_budget_max_bytes != 0 {
+            if config.general.direct_relay_buffer_budget_max_bytes
+                % DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES
+                != 0
+            {
+                return Err(ProxyError::Config(format!(
+                    "general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of {DIRECT_RELAY_BUFFER_BUDGET_UNIT_BYTES}"
+                )));
+            }
+            if !(MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES..=MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES)
+                .contains(&config.general.direct_relay_buffer_budget_max_bytes)
+            {
+                return Err(ProxyError::Config(format!(
+                    "general.direct_relay_buffer_budget_max_bytes must be 0 or within [{MIN_DIRECT_RELAY_BUFFER_BUDGET_BYTES}, {MAX_DIRECT_RELAY_BUFFER_BUDGET_BYTES}]"
+                )));
+            }
+        }
+
         if config.general.me_health_interval_ms_unhealthy == 0 {
             return Err(ProxyError::Config(
                 "general.me_health_interval_ms_unhealthy must be > 0".to_string(),

src/config/load/strict_keys.rs

@@ -65,6 +65,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[
     "me_d2c_frame_buf_shrink_threshold_bytes",
     "direct_relay_copy_buf_c2s_bytes",
     "direct_relay_copy_buf_s2c_bytes",
+    "direct_relay_buffer_budget_max_bytes",
     "crypto_pending_buffer",
     "max_client_frame",
     "desync_all_full",

src/config/tests/load_memory_envelope_tests.rs

@@ -156,6 +156,42 @@ me_writer_byte_budget_bytes = 268451840
     remove_temp_config(&path);
 }
 
+#[test]
+fn load_rejects_unaligned_direct_relay_buffer_budget() {
+    let path = write_temp_config(
+        r#"
+[general]
+direct_relay_buffer_budget_max_bytes = 16777217
+"#,
+    );
+
+    let err = ProxyConfig::load(&path)
+        .expect_err("unaligned direct relay buffer budget must fail");
+    assert!(
+        err.to_string().contains(
+            "general.direct_relay_buffer_budget_max_bytes must be 0 or a multiple of 4096"
+        )
+    );
+    remove_temp_config(&path);
+}
+
+#[test]
+fn load_rejects_direct_relay_buffer_budget_above_hard_cap() {
+    let path = write_temp_config(
+        r#"
+[general]
+direct_relay_buffer_budget_max_bytes = 2147487744
+"#,
+    );
+
+    let err = ProxyConfig::load(&path)
+        .expect_err("direct relay buffer budget above hard cap must fail");
+    assert!(err.to_string().contains(
+        "general.direct_relay_buffer_budget_max_bytes must be 0 or within [16777216, 2147483648]"
+    ));
+    remove_temp_config(&path);
+}
+
 #[test]
 fn load_rejects_listen_backlog_above_i32_upper_bound() {
     let path = write_temp_config(
@@ -203,6 +239,7 @@ me_writer_cmd_channel_capacity = 16384
 me_writer_byte_budget_bytes = 268435456
 me_route_channel_capacity = 8192
 me_c2me_channel_capacity = 8192
+direct_relay_buffer_budget_max_bytes = 2147483648
 max_client_frame = 16777216
 "#,
     );
@@ -212,6 +249,10 @@ max_client_frame = 16777216
     assert_eq!(cfg.general.me_writer_byte_budget_bytes, 256 * 1024 * 1024);
     assert_eq!(cfg.general.me_route_channel_capacity, 8192);
     assert_eq!(cfg.general.me_c2me_channel_capacity, 8192);
+    assert_eq!(
+        cfg.general.direct_relay_buffer_budget_max_bytes,
+        2 * 1024 * 1024 * 1024
+    );
     assert_eq!(cfg.general.max_client_frame, 16 * 1024 * 1024);
 
     remove_temp_config(&path);

src/config/types.rs

@@ -626,7 +626,7 @@ pub struct GeneralConfig {
     #[serde(default = "default_me_d2c_frame_buf_shrink_threshold_bytes")]
     pub me_d2c_frame_buf_shrink_threshold_bytes: usize,
 
-    /// Copy buffer size for client->DC direction in direct relay.
+    /// Copy buffer ceiling for client->DC direction in direct relay.
     ///
     /// This is also the upper bound for one amortized upload rate-limit burst:
     /// upload debt is settled before the next relay read instead of blocking
@@ -634,13 +634,18 @@ pub struct GeneralConfig {
     #[serde(default = "default_direct_relay_copy_buf_c2s_bytes")]
     pub direct_relay_copy_buf_c2s_bytes: usize,
 
-    /// Copy buffer size for DC->client direction in direct relay.
+    /// Copy buffer ceiling for DC->client direction in direct relay.
     ///
     /// This bounds one direct download rate-limit grant because writes are
     /// clipped to the currently available shaper budget.
     #[serde(default = "default_direct_relay_copy_buf_s2c_bytes")]
     pub direct_relay_copy_buf_s2c_bytes: usize,
 
+    /// Process-wide hard ceiling for Direct relay copy buffers.
+    /// `0` derives the ceiling from host and cgroup memory limits.
+    #[serde(default = "default_direct_relay_buffer_budget_max_bytes")]
+    pub direct_relay_buffer_budget_max_bytes: usize,
+
     /// Max pending ciphertext buffer per client writer (bytes).
     /// Controls FakeTLS backpressure vs throughput.
     #[serde(default = "default_crypto_pending_buffer")]
@@ -1121,6 +1126,8 @@ impl Default for GeneralConfig {
                 default_me_d2c_frame_buf_shrink_threshold_bytes(),
             direct_relay_copy_buf_c2s_bytes: default_direct_relay_copy_buf_c2s_bytes(),
             direct_relay_copy_buf_s2c_bytes: default_direct_relay_copy_buf_s2c_bytes(),
+            direct_relay_buffer_budget_max_bytes:
+                default_direct_relay_buffer_budget_max_bytes(),
             me_warmup_stagger_enabled: default_true(),
             me_warmup_step_delay_ms: default_warmup_step_delay_ms(),
             me_warmup_step_jitter_ms: default_warmup_step_jitter_ms(),