From d4c4980e5a6e32db02cb51274e0c1a634d092345 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Sat, 11 Jul 2026 18:43:42 +0300 Subject: [PATCH] Bound ME writer queues by resident payload bytes Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/config/defaults.rs | 16 + src/config/load.rs | 17 + src/config/load/strict_keys.rs | 1 + .../tests/load_memory_envelope_tests.rs | 63 ++++ src/config/types.rs | 5 + src/maestro/me_startup.rs | 1 + src/metrics.rs | 79 ++++ src/stats/me_getters.rs | 29 ++ src/stats/mod.rs | 6 + src/stats/writer_counters.rs | 50 +++ src/transport/middle_proxy/codec.rs | 93 ++++- src/transport/middle_proxy/health.rs | 13 +- src/transport/middle_proxy/pool.rs | 17 +- src/transport/middle_proxy/pool_writer.rs | 16 +- src/transport/middle_proxy/registry.rs | 10 +- src/transport/middle_proxy/registry/tests.rs | 38 +- src/transport/middle_proxy/registry/writer.rs | 23 +- src/transport/middle_proxy/send.rs | 349 +++++++++++++++--- src/transport/middle_proxy/send/close.rs | 18 +- .../tests/health_adversarial_tests.rs | 7 +- .../tests/health_integration_tests.rs | 7 +- .../tests/health_regression_tests.rs | 7 +- .../tests/pool_refill_security_tests.rs | 1 + .../tests/pool_writer_security_tests.rs | 7 +- .../tests/send_adversarial_tests.rs | 7 +- 25 files changed, 798 insertions(+), 82 deletions(-) diff --git a/src/config/defaults.rs b/src/config/defaults.rs index fca0425..3b8ae55 100644 --- a/src/config/defaults.rs +++ b/src/config/defaults.rs @@ -24,6 +24,10 @@ const DEFAULT_ME_ADAPTIVE_FLOOR_MAX_WARM_WRITERS_GLOBAL: u32 = 256; const DEFAULT_ME_ROUTE_BACKPRESSURE_ENABLED: bool = false; const DEFAULT_ME_ROUTE_FAIRSHARE_ENABLED: bool = false; const DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY: usize = 4096; +pub(crate) const ME_WRITER_BYTE_PERMIT_UNIT_BYTES: usize = 16 * 1024; +pub(crate) const ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES: usize = 256; +const DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES: usize = + 32 * 1024 * 1024 + ME_WRITER_BYTE_PERMIT_UNIT_BYTES; const DEFAULT_ME_ROUTE_CHANNEL_CAPACITY: usize = 768; const DEFAULT_ME_C2ME_CHANNEL_CAPACITY: usize = 1024; const DEFAULT_ME_READER_ROUTE_DATA_WAIT_MS: u64 = 2; @@ -455,6 +459,18 @@ pub(crate) fn default_me_writer_cmd_channel_capacity() -> usize { DEFAULT_ME_WRITER_CMD_CHANNEL_CAPACITY } +pub(crate) fn default_me_writer_byte_budget_bytes() -> usize { + DEFAULT_ME_WRITER_BYTE_BUDGET_BYTES +} + +pub(crate) fn minimum_me_writer_byte_budget_bytes(max_client_frame: usize) -> usize { + max_client_frame + .saturating_mul(2) + .saturating_add(ME_WRITER_FRAME_OVERHEAD_RESERVE_BYTES) + .div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES) + .saturating_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES) +} + pub(crate) fn default_me_route_channel_capacity() -> usize { DEFAULT_ME_ROUTE_CHANNEL_CAPACITY } diff --git a/src/config/load.rs b/src/config/load.rs index 4506d57..cf13ba9 100644 --- a/src/config/load.rs +++ b/src/config/load.rs @@ -39,6 +39,7 @@ use self::validation::{ }; 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_MAX_CLIENT_FRAME_BYTES: usize = 4 * 1024; @@ -540,6 +541,22 @@ impl ProxyConfig { ))); } + let min_writer_byte_budget = + minimum_me_writer_byte_budget_bytes(config.general.max_client_frame); + if config.general.me_writer_byte_budget_bytes % ME_WRITER_BYTE_PERMIT_UNIT_BYTES != 0 { + return Err(ProxyError::Config(format!( + "general.me_writer_byte_budget_bytes must be a multiple of {ME_WRITER_BYTE_PERMIT_UNIT_BYTES}" + ))); + } + if !(min_writer_byte_budget..=MAX_ME_WRITER_BYTE_BUDGET_BYTES) + .contains(&config.general.me_writer_byte_budget_bytes) + { + return Err(ProxyError::Config(format!( + "general.me_writer_byte_budget_bytes must be within [{min_writer_byte_budget}, {MAX_ME_WRITER_BYTE_BUDGET_BYTES}] for general.max_client_frame={}", + config.general.max_client_frame + ))); + } + if config.general.me_c2me_send_timeout_ms > 60_000 { return Err(ProxyError::Config( "general.me_c2me_send_timeout_ms must be within [0, 60000]".to_string(), diff --git a/src/config/load/strict_keys.rs b/src/config/load/strict_keys.rs index 9d832e4..47bc3e0 100644 --- a/src/config/load/strict_keys.rs +++ b/src/config/load/strict_keys.rs @@ -52,6 +52,7 @@ const GENERAL_CONFIG_KEYS: &[&str] = &[ "me_keepalive_payload_random", "rpc_proxy_req_every", "me_writer_cmd_channel_capacity", + "me_writer_byte_budget_bytes", "me_route_channel_capacity", "me_c2me_channel_capacity", "me_c2me_send_timeout_ms", diff --git a/src/config/tests/load_memory_envelope_tests.rs b/src/config/tests/load_memory_envelope_tests.rs index 06681d3..0702c5c 100644 --- a/src/config/tests/load_memory_envelope_tests.rs +++ b/src/config/tests/load_memory_envelope_tests.rs @@ -95,6 +95,67 @@ max_client_frame = 16777217 remove_temp_config(&path); } +#[test] +fn load_rejects_writer_byte_budget_below_frame_residency_minimum() { + let path = write_temp_config( + r#" +[general] +max_client_frame = 16777216 +me_writer_byte_budget_bytes = 33554432 +"#, + ); + + let err = ProxyConfig::load(&path) + .expect_err("writer byte budget below frame residency minimum must fail"); + let msg = err.to_string(); + assert!( + msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"), + "error must explain writer byte budget minimum, got: {msg}" + ); + + remove_temp_config(&path); +} + +#[test] +fn load_rejects_unaligned_writer_byte_budget() { + let path = write_temp_config( + r#" +[general] +me_writer_byte_budget_bytes = 33570817 +"#, + ); + + let err = ProxyConfig::load(&path) + .expect_err("writer byte budget outside permit granularity must fail"); + let msg = err.to_string(); + assert!( + msg.contains("general.me_writer_byte_budget_bytes must be a multiple of 16384"), + "error must explain writer byte budget alignment, got: {msg}" + ); + + remove_temp_config(&path); +} + +#[test] +fn load_rejects_writer_byte_budget_above_hard_cap() { + let path = write_temp_config( + r#" +[general] +me_writer_byte_budget_bytes = 268451840 +"#, + ); + + let err = + ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail"); + let msg = err.to_string(); + assert!( + msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"), + "error must explain writer byte budget hard cap, got: {msg}" + ); + + remove_temp_config(&path); +} + #[test] fn load_rejects_listen_backlog_above_i32_upper_bound() { let path = write_temp_config( @@ -139,6 +200,7 @@ fn load_accepts_memory_limits_at_hard_upper_bounds() { r#" [general] me_writer_cmd_channel_capacity = 16384 +me_writer_byte_budget_bytes = 268435456 me_route_channel_capacity = 8192 me_c2me_channel_capacity = 8192 max_client_frame = 16777216 @@ -147,6 +209,7 @@ max_client_frame = 16777216 let cfg = ProxyConfig::load(&path).expect("hard upper bound values must be accepted"); assert_eq!(cfg.general.me_writer_cmd_channel_capacity, 16384); + 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.max_client_frame, 16 * 1024 * 1024); diff --git a/src/config/types.rs b/src/config/types.rs index b89dbe2..0d4d8ab 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -579,6 +579,10 @@ pub struct GeneralConfig { #[serde(default = "default_me_writer_cmd_channel_capacity")] pub me_writer_cmd_channel_capacity: usize, + /// Resident-memory budget in bytes for each ME writer data queue. + #[serde(default = "default_me_writer_byte_budget_bytes")] + pub me_writer_byte_budget_bytes: usize, + /// Capacity of per-connection ME response route channel. #[serde(default = "default_me_route_channel_capacity")] pub me_route_channel_capacity: usize, @@ -1103,6 +1107,7 @@ impl Default for GeneralConfig { me_keepalive_payload_random: default_true(), rpc_proxy_req_every: default_rpc_proxy_req_every(), me_writer_cmd_channel_capacity: default_me_writer_cmd_channel_capacity(), + me_writer_byte_budget_bytes: default_me_writer_byte_budget_bytes(), me_route_channel_capacity: default_me_route_channel_capacity(), me_c2me_channel_capacity: default_me_c2me_channel_capacity(), me_c2me_send_timeout_ms: default_me_c2me_send_timeout_ms(), diff --git a/src/maestro/me_startup.rs b/src/maestro/me_startup.rs index 21e1ccf..65b37d1 100644 --- a/src/maestro/me_startup.rs +++ b/src/maestro/me_startup.rs @@ -279,6 +279,7 @@ pub(crate) async fn initialize_me_pool( config.general.me_writer_pick_sample_size, config.general.me_socks_kdf_policy, config.general.me_writer_cmd_channel_capacity, + config.general.me_writer_byte_budget_bytes, config.general.me_route_channel_capacity, config.general.me_route_backpressure_enabled, config.general.me_route_fairshare_enabled, diff --git a/src/metrics.rs b/src/metrics.rs index 3859b0c..704f8b8 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -2435,6 +2435,85 @@ async fn render_metrics( } ); + let _ = writeln!( + out, + "# HELP telemt_me_writer_byte_budget_limit_bytes Configured resident-memory budget per ME writer" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_writer_byte_budget_limit_bytes gauge" + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_limit_bytes {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_limit_bytes_gauge() + } else { + 0 + } + ); + let _ = writeln!( + out, + "# HELP telemt_me_writer_byte_budget_reserved_bytes Aggregate ME writer memory reservations by lifecycle state" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_writer_byte_budget_reserved_bytes gauge" + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_reserved_bytes{{state=\"queued\"}} {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_queued_bytes_gauge() + } else { + 0 + } + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_reserved_bytes{{state=\"inflight\"}} {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_inflight_bytes_gauge() + } else { + 0 + } + ); + let _ = writeln!( + out, + "# HELP telemt_me_writer_byte_budget_events_total ME writer byte-budget outcomes" + ); + let _ = writeln!( + out, + "# TYPE telemt_me_writer_byte_budget_events_total counter" + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_events_total{{result=\"wait\"}} {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_wait_total() + } else { + 0 + } + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_events_total{{result=\"timeout\"}} {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_timeout_total() + } else { + 0 + } + ); + let _ = writeln!( + out, + "telemt_me_writer_byte_budget_events_total{{result=\"oversize\"}} {}", + if me_allows_normal { + stats.get_me_writer_byte_budget_oversize_total() + } else { + 0 + } + ); + let _ = writeln!( out, "# HELP telemt_me_writer_pick_total ME writer-pick outcomes by mode and result" diff --git a/src/stats/me_getters.rs b/src/stats/me_getters.rs index 0737089..4ab6dac 100644 --- a/src/stats/me_getters.rs +++ b/src/stats/me_getters.rs @@ -269,6 +269,35 @@ impl Stats { self.me_writer_pick_mode_switch_total .load(Ordering::Relaxed) } + /// Returns the configured resident-memory limit per ME writer. + pub fn get_me_writer_byte_budget_limit_bytes_gauge(&self) -> u64 { + self.me_writer_byte_budget_limit_bytes_gauge + .load(Ordering::Relaxed) + } + /// Returns aggregate queued or enqueueing writer memory reservations. + pub fn get_me_writer_byte_budget_queued_bytes_gauge(&self) -> u64 { + self.me_writer_byte_budget_queued_bytes_gauge + .load(Ordering::Relaxed) + } + /// Returns aggregate writer reservations currently owned by socket writes. + pub fn get_me_writer_byte_budget_inflight_bytes_gauge(&self) -> u64 { + self.me_writer_byte_budget_inflight_bytes_gauge + .load(Ordering::Relaxed) + } + /// Returns the count of blocking writer byte-budget waits. + pub fn get_me_writer_byte_budget_wait_total(&self) -> u64 { + self.me_writer_byte_budget_wait_total.load(Ordering::Relaxed) + } + /// Returns the count of writer byte-budget wait timeouts. + pub fn get_me_writer_byte_budget_timeout_total(&self) -> u64 { + self.me_writer_byte_budget_timeout_total + .load(Ordering::Relaxed) + } + /// Returns the count of payloads that cannot fit the configured writer budget. + pub fn get_me_writer_byte_budget_oversize_total(&self) -> u64 { + self.me_writer_byte_budget_oversize_total + .load(Ordering::Relaxed) + } pub fn get_me_socks_kdf_strict_reject(&self) -> u64 { self.me_socks_kdf_strict_reject.load(Ordering::Relaxed) } diff --git a/src/stats/mod.rs b/src/stats/mod.rs index 7f05bbf..b56ee5b 100644 --- a/src/stats/mod.rs +++ b/src/stats/mod.rs @@ -292,6 +292,12 @@ pub struct Stats { me_writer_pick_p2c_no_candidate_total: AtomicU64, me_writer_pick_blocking_fallback_total: AtomicU64, me_writer_pick_mode_switch_total: AtomicU64, + me_writer_byte_budget_limit_bytes_gauge: AtomicU64, + me_writer_byte_budget_queued_bytes_gauge: AtomicU64, + me_writer_byte_budget_inflight_bytes_gauge: AtomicU64, + me_writer_byte_budget_wait_total: AtomicU64, + me_writer_byte_budget_timeout_total: AtomicU64, + me_writer_byte_budget_oversize_total: AtomicU64, me_socks_kdf_strict_reject: AtomicU64, me_socks_kdf_compat_fallback: AtomicU64, secure_padding_invalid: AtomicU64, diff --git a/src/stats/writer_counters.rs b/src/stats/writer_counters.rs index 2977202..f8c852c 100644 --- a/src/stats/writer_counters.rs +++ b/src/stats/writer_counters.rs @@ -88,6 +88,56 @@ impl Stats { .fetch_add(1, Ordering::Relaxed); } } + /// Publishes the configured resident-memory limit applied to each ME writer. + pub fn set_me_writer_byte_budget_limit_bytes(&self, bytes: usize) { + self.me_writer_byte_budget_limit_bytes_gauge + .store(bytes as u64, Ordering::Relaxed); + } + pub(crate) fn add_me_writer_byte_budget_queued_bytes(&self, bytes: u64) { + self.me_writer_byte_budget_queued_bytes_gauge + .fetch_add(bytes, Ordering::Relaxed); + } + pub(crate) fn move_me_writer_byte_budget_to_inflight(&self, bytes: u64) { + let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update( + Ordering::Relaxed, + Ordering::Relaxed, + |current| Some(current.saturating_sub(bytes)), + ); + self.me_writer_byte_budget_inflight_bytes_gauge + .fetch_add(bytes, Ordering::Relaxed); + } + pub(crate) fn release_me_writer_byte_budget_queued_bytes(&self, bytes: u64) { + let _ = self.me_writer_byte_budget_queued_bytes_gauge.fetch_update( + Ordering::Relaxed, + Ordering::Relaxed, + |current| Some(current.saturating_sub(bytes)), + ); + } + pub(crate) fn release_me_writer_byte_budget_inflight_bytes(&self, bytes: u64) { + let _ = self.me_writer_byte_budget_inflight_bytes_gauge.fetch_update( + Ordering::Relaxed, + Ordering::Relaxed, + |current| Some(current.saturating_sub(bytes)), + ); + } + pub(crate) fn increment_me_writer_byte_budget_wait_total(&self) { + if self.telemetry_me_allows_normal() { + self.me_writer_byte_budget_wait_total + .fetch_add(1, Ordering::Relaxed); + } + } + pub(crate) fn increment_me_writer_byte_budget_timeout_total(&self) { + if self.telemetry_me_allows_normal() { + self.me_writer_byte_budget_timeout_total + .fetch_add(1, Ordering::Relaxed); + } + } + pub(crate) fn increment_me_writer_byte_budget_oversize_total(&self) { + if self.telemetry_me_allows_normal() { + self.me_writer_byte_budget_oversize_total + .fetch_add(1, Ordering::Relaxed); + } + } pub fn increment_me_socks_kdf_strict_reject(&self) { if self.telemetry_me_allows_normal() { self.me_socks_kdf_strict_reject diff --git a/src/transport/middle_proxy/codec.rs b/src/transport/middle_proxy/codec.rs index 235eb3a..0677def 100644 --- a/src/transport/middle_proxy/codec.rs +++ b/src/transport/middle_proxy/codec.rs @@ -1,10 +1,12 @@ use bytes::Bytes; +use std::sync::Arc; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::OwnedSemaphorePermit; use crate::crypto::{AesCbc, crc32, crc32c}; use crate::error::{ProxyError, Result}; use crate::protocol::constants::*; +use crate::stats::Stats; use crate::stream::PooledBuffer; use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len}; @@ -12,11 +14,65 @@ use super::wire::{append_proxy_req_payload_into, proxy_req_payload_len}; const RPC_WRITER_FRAME_BUF_SHRINK_THRESHOLD: usize = 256 * 1024; const RPC_WRITER_FRAME_BUF_RETAIN: usize = 64 * 1024; +enum WriterBytePermitState { + Queued, + Inflight, +} + +/// Holds one writer memory reservation through queueing and socket write completion. +pub(crate) struct WriterBytePermit { + _permit: OwnedSemaphorePermit, + reserved_bytes: u64, + stats: Arc, + state: WriterBytePermitState, +} + +impl WriterBytePermit { + /// Creates a queued reservation and publishes its rounded resident-byte cost. + pub(crate) fn new( + permit: OwnedSemaphorePermit, + reserved_bytes: usize, + stats: Arc, + ) -> Self { + let reserved_bytes = reserved_bytes as u64; + stats.add_me_writer_byte_budget_queued_bytes(reserved_bytes); + Self { + _permit: permit, + reserved_bytes, + stats, + state: WriterBytePermitState::Queued, + } + } + + /// Moves the reservation from queued to in-flight accounting exactly once. + pub(crate) fn mark_inflight(&mut self) { + if matches!(self.state, WriterBytePermitState::Queued) { + self.stats + .move_me_writer_byte_budget_to_inflight(self.reserved_bytes); + self.state = WriterBytePermitState::Inflight; + } + } +} + +impl Drop for WriterBytePermit { + fn drop(&mut self) { + match self.state { + WriterBytePermitState::Queued => self + .stats + .release_me_writer_byte_budget_queued_bytes(self.reserved_bytes), + WriterBytePermitState::Inflight => self + .stats + .release_me_writer_byte_budget_inflight_bytes(self.reserved_bytes), + } + } +} + /// Commands sent to dedicated writer tasks to avoid mutex contention on TCP writes. pub(crate) enum WriterCommand { Data { payload: Bytes, _permit: Option, + writer_permit: WriterBytePermit, }, DataAndFlush(Bytes), ProxyReq(ProxyReqCommand), @@ -33,6 +89,7 @@ pub(crate) struct ProxyReqCommand { pub(crate) proxy_tag: Option<[u8; 16]>, pub(crate) payload: PooledBuffer, pub(crate) _permit: OwnedSemaphorePermit, + pub(crate) writer_permit: WriterBytePermit, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -87,7 +144,7 @@ fn build_rpc_frame_into( ) { let total_len = 4 + 4 + payload.len() + 4; frame.clear(); - frame.reserve(total_len + 15); + frame.reserve_exact(total_len + 15); let total_len = total_len as u32; frame.extend_from_slice(&total_len.to_le_bytes()); frame.extend_from_slice(&seq_no.to_le_bytes()); @@ -279,7 +336,7 @@ impl RpcWriter { ); let total_len = 4 + 4 + payload_len + 4; self.frame_buf.clear(); - self.frame_buf.reserve(total_len + 15); + self.frame_buf.reserve_exact(total_len + 15); self.frame_buf .extend_from_slice(&(total_len as u32).to_le_bytes()); self.frame_buf.extend_from_slice(&self.seq_no.to_le_bytes()); @@ -327,3 +384,35 @@ impl RpcWriter { self.writer.flush().await.map_err(ProxyError::Io) } } + +#[cfg(test)] +mod tests { + use super::*; + use tokio::sync::Semaphore; + + #[test] + fn writer_byte_permit_tracks_queued_and_inflight_lifecycle() { + let stats = Arc::new(Stats::default()); + let semaphore = Arc::new(Semaphore::new(2)); + let permit = semaphore + .clone() + .try_acquire_many_owned(2) + .expect("writer byte permits must be available"); + let mut writer_permit = WriterBytePermit::new(permit, 32 * 1024, stats.clone()); + + assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 32 * 1024); + assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0); + + writer_permit.mark_inflight(); + assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0); + assert_eq!( + stats.get_me_writer_byte_budget_inflight_bytes_gauge(), + 32 * 1024 + ); + + drop(writer_permit); + assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 0); + assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0); + assert_eq!(semaphore.available_permits(), 2); + } +} diff --git a/src/transport/middle_proxy/health.rs b/src/transport/middle_proxy/health.rs index 0573167..b288f5b 100644 --- a/src/transport/middle_proxy/health.rs +++ b/src/transport/middle_proxy/health.rs @@ -1795,6 +1795,7 @@ mod tests { general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -1821,6 +1822,7 @@ mod tests { ) -> u64 { let (conn_id, _rx) = pool.registry.register().await; let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4000 + writer_id as u16), @@ -1830,6 +1832,7 @@ mod tests { contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), created_at: Instant::now() - Duration::from_secs(writer_id), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -1839,7 +1842,9 @@ mod tests { allow_drain_fallback: Arc::new(AtomicBool::new(false)), }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); assert!( pool.registry @@ -1860,6 +1865,7 @@ mod tests { async fn insert_live_writer(pool: &Arc, writer_id: u64, writer_dc: i32) { let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr: SocketAddr::new( @@ -1877,6 +1883,7 @@ mod tests { contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), created_at: Instant::now(), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -1886,7 +1893,9 @@ mod tests { allow_drain_fallback: Arc::new(AtomicBool::new(false)), }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); } diff --git a/src/transport/middle_proxy/pool.rs b/src/transport/middle_proxy/pool.rs index 077a014..69754de 100644 --- a/src/transport/middle_proxy/pool.rs +++ b/src/transport/middle_proxy/pool.rs @@ -10,7 +10,7 @@ use std::sync::atomic::{ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use arc_swap::ArcSwap; -use tokio::sync::{Mutex, RwLock, mpsc, watch}; +use tokio::sync::{Mutex, RwLock, Semaphore, mpsc, watch}; use tokio_util::sync::CancellationToken; use crate::config::{ @@ -48,6 +48,8 @@ pub struct MeWriter { pub contour: Arc, pub created_at: Instant, pub tx: mpsc::Sender, + /// Aggregate resident-memory budget shared by all data commands for this writer. + pub byte_budget: Arc, pub cancel: CancellationToken, pub degraded: Arc, pub rtt_ema_ms_x10: Arc, @@ -277,6 +279,7 @@ pub(super) struct WriterLifecycleCore { pub(super) me_keepalive_payload_random: bool, pub(super) rpc_proxy_req_every_secs: AtomicU64, pub(super) writer_cmd_channel_capacity: usize, + pub(super) writer_byte_budget_permits: usize, } pub(super) struct RouteRuntimeCore { @@ -553,6 +556,7 @@ impl MePool { me_writer_pick_sample_size: u8, me_socks_kdf_policy: MeSocksKdfPolicy, me_writer_cmd_channel_capacity: usize, + me_writer_byte_budget_bytes: usize, me_route_channel_capacity: usize, me_route_backpressure_enabled: bool, me_route_fairshare_enabled: bool, @@ -583,6 +587,7 @@ impl MePool { ); let (writer_epoch, _) = watch::channel(0u64); let now_epoch_secs = Self::now_epoch_secs(); + stats.set_me_writer_byte_budget_limit_bytes(me_writer_byte_budget_bytes); Arc::new(Self { routing: Arc::new(RoutingCore { registry, @@ -615,6 +620,9 @@ impl MePool { me_keepalive_payload_random, rpc_proxy_req_every_secs: AtomicU64::new(rpc_proxy_req_every_secs), writer_cmd_channel_capacity: me_writer_cmd_channel_capacity.max(1), + writer_byte_budget_permits: me_writer_byte_budget_bytes + .div_ceil(crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES) + .max(1), }), route_runtime: Arc::new(RouteRuntimeCore { me_route_no_writer_mode: AtomicU8::new(me_route_no_writer_mode.as_u8()), @@ -833,6 +841,13 @@ impl MePool { }) } + /// Creates the immutable byte semaphore assigned to one ME writer generation. + pub(crate) fn new_writer_byte_budget(&self) -> Arc { + Arc::new(Semaphore::new( + self.writer_lifecycle.writer_byte_budget_permits, + )) + } + pub fn current_generation(&self) -> u64 { self.reinit.active_generation.load(Ordering::Relaxed) } diff --git a/src/transport/middle_proxy/pool_writer.rs b/src/transport/middle_proxy/pool_writer.rs index 34e52c2..c87c296 100644 --- a/src/transport/middle_proxy/pool_writer.rs +++ b/src/transport/middle_proxy/pool_writer.rs @@ -63,13 +63,19 @@ async fn writer_command_loop( _ = cancel.cancelled() => return Ok(()), cmd = rx.recv() => { match cmd { - Some(WriterCommand::Data { payload, _permit }) => { + Some(WriterCommand::Data { + payload, + _permit, + mut writer_permit, + }) => { + writer_permit.mark_inflight(); rpc_writer.send(&payload).await?; } Some(WriterCommand::DataAndFlush(payload)) => { rpc_writer.send_and_flush(&payload).await?; } - Some(WriterCommand::ProxyReq(command)) => { + Some(WriterCommand::ProxyReq(mut command)) => { + command.writer_permit.mark_inflight(); rpc_writer.send_proxy_req(&command).await?; } Some(WriterCommand::ControlAndFlush(payload)) => { @@ -419,6 +425,7 @@ impl MePool { let draining_started_at_epoch_secs = Arc::new(AtomicU64::new(0)); let drain_deadline_epoch_secs = Arc::new(AtomicU64::new(0)); let allow_drain_fallback = Arc::new(AtomicBool::new(false)); + let byte_budget = self.new_writer_byte_budget(); let (tx, rx) = mpsc::channel::(self.writer_lifecycle.writer_cmd_channel_capacity); let rpc_writer = RpcWriter { @@ -438,6 +445,7 @@ impl MePool { contour: contour.clone(), created_at: Instant::now(), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: cancel.clone(), degraded: degraded.clone(), rtt_ema_ms_x10: rtt_ema_ms_x10.clone(), @@ -449,7 +457,9 @@ impl MePool { self.writers .update(|writers| writers.push(writer.clone())) .await; - self.registry.register_writer(writer_id, tx.clone()).await; + self.registry + .register_writer(writer_id, tx.clone(), byte_budget) + .await; self.registry.mark_writer_idle(writer_id).await; self.conn_count.fetch_add(1, Ordering::Relaxed); self.notify_writer_epoch(); diff --git a/src/transport/middle_proxy/registry.rs b/src/transport/middle_proxy/registry.rs index 71f75f8..b3d3f4b 100644 --- a/src/transport/middle_proxy/registry.rs +++ b/src/transport/middle_proxy/registry.rs @@ -48,6 +48,8 @@ pub struct BoundConn { pub struct ConnWriter { pub writer_id: u64, pub tx: mpsc::Sender, + /// Writer-local memory budget used by the hot bound-client route. + pub byte_budget: Arc, } #[derive(Clone, Debug, Default)] @@ -62,7 +64,13 @@ struct RoutingTable { } struct WriterTable { - map: DashMap>, + map: DashMap, +} + +#[derive(Clone)] +struct WriterRoute { + tx: mpsc::Sender, + byte_budget: Arc, } #[derive(Clone)] diff --git a/src/transport/middle_proxy/registry/tests.rs b/src/transport/middle_proxy/registry/tests.rs index 3567a4b..fe981a5 100644 --- a/src/transport/middle_proxy/registry/tests.rs +++ b/src/transport/middle_proxy/registry/tests.rs @@ -1,10 +1,16 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::sync::Arc; use bytes::Bytes; +use tokio::sync::Semaphore; use super::{ConnMeta, ConnRegistry, RouteResult}; use crate::transport::middle_proxy::MeResponse; +fn writer_byte_budget() -> Arc { + Arc::new(Semaphore::new(2049)) +} + #[tokio::test] async fn writer_activity_snapshot_tracks_writer_and_dc_load() { let registry = ConnRegistry::new(); @@ -14,8 +20,12 @@ async fn writer_activity_snapshot_tracks_writer_and_dc_load() { let (conn_c, _rx_c) = registry.register().await; let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); - registry.register_writer(10, writer_tx_a.clone()).await; - registry.register_writer(20, writer_tx_b.clone()).await; + registry + .register_writer(10, writer_tx_a.clone(), writer_byte_budget()) + .await; + registry + .register_writer(20, writer_tx_b.clone(), writer_byte_budget()) + .await; let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); assert!( @@ -124,8 +134,12 @@ async fn bind_writer_rebinds_conn_atomically() { let (conn_id, _rx) = registry.register().await; let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); - registry.register_writer(10, writer_tx_a).await; - registry.register_writer(20, writer_tx_b).await; + registry + .register_writer(10, writer_tx_a, writer_byte_budget()) + .await; + registry + .register_writer(20, writer_tx_b, writer_byte_budget()) + .await; let client_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); let first_our_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 443); @@ -184,8 +198,12 @@ async fn writer_lost_does_not_drop_rebound_conn() { let (conn_id, _rx) = registry.register().await; let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); - registry.register_writer(10, writer_tx_a).await; - registry.register_writer(20, writer_tx_b).await; + registry + .register_writer(10, writer_tx_a, writer_byte_budget()) + .await; + registry + .register_writer(20, writer_tx_b, writer_byte_budget()) + .await; let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); assert!( @@ -262,8 +280,12 @@ async fn non_empty_writer_ids_returns_only_writers_with_bound_clients() { let (conn_id, _rx) = registry.register().await; let (writer_tx_a, _writer_rx_a) = tokio::sync::mpsc::channel(8); let (writer_tx_b, _writer_rx_b) = tokio::sync::mpsc::channel(8); - registry.register_writer(10, writer_tx_a).await; - registry.register_writer(20, writer_tx_b).await; + registry + .register_writer(10, writer_tx_a, writer_byte_budget()) + .await; + registry + .register_writer(20, writer_tx_b, writer_byte_budget()) + .await; let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 443); assert!( diff --git a/src/transport/middle_proxy/registry/writer.rs b/src/transport/middle_proxy/registry/writer.rs index 4d6ea0d..ea0336f 100644 --- a/src/transport/middle_proxy/registry/writer.rs +++ b/src/transport/middle_proxy/registry/writer.rs @@ -1,4 +1,5 @@ use std::collections::{HashMap, HashSet}; +use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::Duration; @@ -56,7 +57,13 @@ impl ConnRegistry { } } - pub async fn register_writer(&self, writer_id: u64, tx: mpsc::Sender) { + /// Registers one writer command route and its matching memory budget atomically. + pub async fn register_writer( + &self, + writer_id: u64, + tx: mpsc::Sender, + byte_budget: Arc, + ) { let mut binding = self.binding.inner.lock().await; binding .conns_for_writer @@ -70,7 +77,13 @@ impl ConnRegistry { .writer_idle_since_epoch_secs .entry(writer_id) .or_insert_with(Self::now_epoch_secs); - self.writers.map.insert(writer_id, tx); + self.writers.map.insert( + writer_id, + super::WriterRoute { + tx, + byte_budget, + }, + ); } /// Unregister connection, returning associated writer_id if any. @@ -417,7 +430,8 @@ impl ConnRegistry { .map(|entry| entry.value().clone())?; Some(ConnWriter { writer_id, - tx: writer, + tx: writer.tx, + byte_budget: writer.byte_budget, }) } @@ -438,7 +452,8 @@ impl ConnRegistry { Some(( ConnWriter { writer_id, - tx: writer, + tx: writer.tx, + byte_budget: writer.byte_budget, }, meta, )) diff --git a/src/transport/middle_proxy/send.rs b/src/transport/middle_proxy/send.rs index 27120fe..2407949 100644 --- a/src/transport/middle_proxy/send.rs +++ b/src/transport/middle_proxy/send.rs @@ -6,16 +6,18 @@ use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; -use tokio::sync::{OwnedSemaphorePermit, mpsc}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; use tokio::sync::mpsc::error::TrySendError; use tracing::{debug, warn}; use super::MePool; -use super::codec::{ProxyReqCommand, WriterCommand}; +use super::codec::{ProxyReqCommand, WriterBytePermit, WriterCommand}; use super::registry::ConnMeta; -use super::wire::build_proxy_req_payload; +use super::wire::{build_proxy_req_payload, proxy_req_payload_len}; +use crate::config::defaults::ME_WRITER_BYTE_PERMIT_UNIT_BYTES; use crate::config::{MeRouteNoWriterMode, MeWriterPickMode}; use crate::error::{ProxyError, Result}; +use crate::stats::Stats; use crate::stream::PooledBuffer; use rand::seq::SliceRandom; @@ -29,6 +31,8 @@ const PICK_PENALTY_WARM: u64 = 200; const PICK_PENALTY_DRAINING: u64 = 600; const PICK_PENALTY_STALE: u64 = 300; const PICK_PENALTY_DEGRADED: u64 = 250; +const RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES: usize = 27; +const LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES: usize = 128; mod close; mod recovery; @@ -39,6 +43,11 @@ enum WriterCommandReserveError { TimedOut, } +enum WriterByteReserveError { + Closed, + TimedOut, +} + fn proxy_tag_array(tag: Option<&[u8]>) -> Option<[u8; 16]> { tag.and_then(|tag| <[u8; 16]>::try_from(tag).ok()) } @@ -61,11 +70,16 @@ fn payload_permit_from_data_command(cmd: WriterCommand) -> Option, - wait: Option, + deadline: Option, ) -> std::result::Result, WriterCommandReserveError> { let reserve = tx.clone().reserve_owned(); - match wait { - Some(wait) => match tokio::time::timeout(wait, reserve).await { + match deadline { + Some(deadline) => match tokio::time::timeout( + deadline.saturating_duration_since(Instant::now()), + reserve, + ) + .await + { Ok(Ok(permit)) => Ok(permit), Ok(Err(_)) => Err(WriterCommandReserveError::Closed), Err(_) => Err(WriterCommandReserveError::TimedOut), @@ -74,6 +88,88 @@ async fn reserve_writer_command_slot( } } +fn writer_send_deadline(wait: Option) -> Option { + wait.map(|wait| Instant::now() + wait) +} + +fn writer_resident_permits( + source_capacity: usize, + encoded_payload_len: usize, +) -> Option<(u32, usize)> { + let resident_bytes = source_capacity + .checked_add(encoded_payload_len)? + .checked_add(RPC_WRITER_FRAME_CAPACITY_OVERHEAD_BYTES)?; + let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES); + let permits = u32::try_from(permits).ok()?; + let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?; + Some((permits.max(1), reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES))) +} + +fn proxy_req_resident_permits( + source_capacity: usize, + data_len: usize, + proxy_tag: Option<&[u8]>, + proto_flags: u32, +) -> Option<(u32, usize)> { + writer_resident_permits( + source_capacity, + proxy_req_payload_len(data_len, proxy_tag, proto_flags), + ) +} + +fn try_reserve_writer_bytes( + byte_budget: &Arc, + permits: u32, + reserved_bytes: usize, + stats: &Arc, +) -> std::result::Result { + byte_budget + .clone() + .try_acquire_many_owned(permits) + .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) +} + +async fn reserve_writer_bytes( + byte_budget: &Arc, + permits: u32, + reserved_bytes: usize, + deadline: Option, + stats: &Arc, +) -> std::result::Result { + match try_reserve_writer_bytes(byte_budget, permits, reserved_bytes, stats) { + Ok(permit) => return Ok(permit), + Err(TryAcquireError::Closed) => return Err(WriterByteReserveError::Closed), + Err(TryAcquireError::NoPermits) => { + stats.increment_me_writer_byte_budget_wait_total(); + } + } + + let acquire = byte_budget.clone().acquire_many_owned(permits); + match deadline { + Some(deadline) => match tokio::time::timeout( + deadline.saturating_duration_since(Instant::now()), + acquire, + ) + .await + { + Ok(Ok(permit)) => Ok(WriterBytePermit::new( + permit, + reserved_bytes, + stats.clone(), + )), + Ok(Err(_)) => Err(WriterByteReserveError::Closed), + Err(_) => { + stats.increment_me_writer_byte_budget_timeout_total(); + Err(WriterByteReserveError::TimedOut) + } + }, + None => acquire + .await + .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) + .map_err(|_| WriterByteReserveError::Closed), + } +} + impl MePool { /// Send RPC_PROXY_REQ. `tag_override`: per-user ad_tag (from access.user_ad_tags); if None, uses pool default. /// `payload_permit` keeps optional client byte accounting alive until the writer consumes the command. @@ -89,6 +185,35 @@ impl MePool { mut payload_permit: Option, ) -> Result<()> { let tag = tag_override.or(self.proxy_tag.as_deref()); + let Some(source_capacity) = data + .len() + .checked_add(LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES) + else { + self.stats + .increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload residency calculation overflow".into(), + )); + }; + let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits( + source_capacity, + data.len(), + tag, + proto_flags, + ) else { + self.stats + .increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload residency calculation overflow".into(), + )); + }; + if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { + self.stats + .increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload exceeds configured byte budget".into(), + )); + } let build_routed_payload = |effective_our_addr: SocketAddr| { ( build_proxy_req_payload( @@ -129,10 +254,38 @@ impl MePool { loop { if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { + let deadline = writer_send_deadline( + self.route_runtime.me_route_blocking_send_timeout, + ); + let writer_permit = match reserve_writer_bytes( + ¤t.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + deadline, + &self.stats, + ) + .await + { + Ok(permit) => permit, + Err(WriterByteReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer byte budget full within blocking send timeout".into(), + )); + } + Err(WriterByteReserveError::Closed) => { + warn!(writer_id = current.writer_id, "ME writer byte budget closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + continue; + } + }; let (current_payload, _) = build_routed_payload(current_meta.our_addr); let command = WriterCommand::Data { payload: current_payload, _permit: payload_permit.take(), + writer_permit, }; match current.tx.try_send(command) { Ok(()) => { @@ -142,7 +295,7 @@ impl MePool { Err(TrySendError::Full(cmd)) => { match reserve_writer_command_slot( ¤t.tx, - self.route_runtime.me_route_blocking_send_timeout, + deadline, ) .await { @@ -482,11 +635,31 @@ impl MePool { if !self.writer_accepts_new_binding(w) { continue; } - // Keep the advertised proxy IP aligned with the selected ME writer source. - let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); - let (payload, meta) = build_routed_payload(effective_our_addr); + let writer_permit = match try_reserve_writer_bytes( + &w.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + &self.stats, + ) { + Ok(permit) => permit, + Err(TryAcquireError::NoPermits) => { + if fallback_blocking_idx.is_none() { + fallback_blocking_idx = Some(idx); + } + continue; + } + Err(TryAcquireError::Closed) => { + self.stats.increment_me_writer_pick_closed_total(pick_mode); + warn!(writer_id = w.id, "ME writer byte budget closed"); + self.remove_writer_and_close_clients(w.id).await; + continue; + } + }; match w.tx.clone().try_reserve_owned() { Ok(permit) => { + // Keep the advertised proxy IP aligned with the selected ME writer source. + let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); + let (payload, meta) = build_routed_payload(effective_our_addr); if !self.registry.bind_writer(conn_id, w.id, meta).await { debug!( conn_id, @@ -500,6 +673,7 @@ impl MePool { permit.send(WriterCommand::Data { payload, _permit: payload_permit.take(), + writer_permit, }); self.stats .increment_me_writer_pick_success_try_total(pick_mode); @@ -542,51 +716,68 @@ impl MePool { } self.stats .increment_me_writer_pick_blocking_fallback_total(); - // Keep the advertised proxy IP aligned with the selected ME writer source. - let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); - let (payload, meta) = build_routed_payload(effective_our_addr); - let reserve_result = - if let Some(timeout) = self.route_runtime.me_route_blocking_send_timeout { - match tokio::time::timeout(timeout, w.tx.clone().reserve_owned()).await { - Ok(result) => result, - Err(_) => { - self.stats.increment_me_writer_pick_full_total(pick_mode); - continue; - } - } - } else { - w.tx.clone().reserve_owned().await - }; - match reserve_result { - Ok(permit) => { - if !self.registry.bind_writer(conn_id, w.id, meta).await { - debug!( - conn_id, - writer_id = w.id, - "ME writer disappeared before fallback bind commit, pruning stale writer" - ); - drop(permit); - self.remove_writer_and_close_clients(w.id).await; - continue; - } - permit.send(WriterCommand::Data { - payload, - _permit: payload_permit.take(), - }); - self.stats - .increment_me_writer_pick_success_fallback_total(pick_mode); - if w.generation < self.current_generation() { - self.stats.increment_pool_stale_pick_total(); - } - self.note_hybrid_route_success(); - return Ok(()); + let deadline = writer_send_deadline( + self.route_runtime.me_route_blocking_send_timeout, + ); + let writer_permit = match reserve_writer_bytes( + &w.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + deadline, + &self.stats, + ) + .await + { + Ok(permit) => permit, + Err(WriterByteReserveError::TimedOut) => { + self.stats.increment_me_writer_pick_full_total(pick_mode); + continue; } - Err(_) => { + Err(WriterByteReserveError::Closed) => { + self.stats.increment_me_writer_pick_closed_total(pick_mode); + warn!(writer_id = w.id, "ME writer byte budget closed (blocking)"); + self.remove_writer_and_close_clients(w.id).await; + continue; + } + }; + let permit = match reserve_writer_command_slot(&w.tx, deadline).await { + Ok(permit) => permit, + Err(WriterCommandReserveError::TimedOut) => { + self.stats.increment_me_writer_pick_full_total(pick_mode); + continue; + } + Err(WriterCommandReserveError::Closed) => { self.stats.increment_me_writer_pick_closed_total(pick_mode); warn!(writer_id = w.id, "ME writer channel closed (blocking)"); self.remove_writer_and_close_clients(w.id).await; + continue; } + }; + // Keep the advertised proxy IP aligned with the selected ME writer source. + let effective_our_addr = SocketAddr::new(w.source_ip, our_addr.port()); + let (payload, meta) = build_routed_payload(effective_our_addr); + if !self.registry.bind_writer(conn_id, w.id, meta).await { + debug!( + conn_id, + writer_id = w.id, + "ME writer disappeared before fallback bind commit, pruning stale writer" + ); + drop(permit); + self.remove_writer_and_close_clients(w.id).await; + continue; } + permit.send(WriterCommand::Data { + payload, + _permit: payload_permit.take(), + writer_permit, + }); + self.stats + .increment_me_writer_pick_success_fallback_total(pick_mode); + if w.generation < self.current_generation() { + self.stats.increment_pool_stale_pick_total(); + } + self.note_hybrid_route_success(); + return Ok(()); } } @@ -604,8 +795,65 @@ impl MePool { tag_override: Option<[u8; 16]>, ) -> Result<()> { let tag = tag_override.or_else(|| proxy_tag_array(self.proxy_tag.as_deref())); + let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits( + payload.capacity(), + payload.len(), + tag.as_ref().map(|tag| tag.as_slice()), + proto_flags, + ) else { + self.stats + .increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload residency calculation overflow".into(), + )); + }; + if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { + self.stats + .increment_me_writer_byte_budget_oversize_total(); + return Err(ProxyError::Proxy( + "ME writer payload exceeds configured byte budget".into(), + )); + } if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { + let deadline = writer_send_deadline( + self.route_runtime.me_route_blocking_send_timeout, + ); + let writer_permit = match reserve_writer_bytes( + ¤t.byte_budget, + writer_byte_permits, + writer_reserved_bytes, + deadline, + &self.stats, + ) + .await + { + Ok(permit) => permit, + Err(WriterByteReserveError::TimedOut) => { + self.stats + .increment_me_writer_pick_full_total(self.writer_pick_mode()); + return Err(ProxyError::Proxy( + "ME writer byte budget full within blocking send timeout".into(), + )); + } + Err(WriterByteReserveError::Closed) => { + warn!(writer_id = current.writer_id, "ME writer byte budget closed"); + self.remove_writer_and_close_clients(current.writer_id) + .await; + return self + .send_proxy_req( + conn_id, + target_dc, + client_addr, + our_addr, + payload.as_ref(), + proto_flags, + tag.as_ref().map(|tag| tag.as_slice()), + Some(_permit), + ) + .await; + } + }; let command = WriterCommand::ProxyReq(ProxyReqCommand { conn_id, client_addr, @@ -614,6 +862,7 @@ impl MePool { proxy_tag: tag, payload, _permit, + writer_permit, }); match current.tx.try_send(command) { Ok(()) => { @@ -623,7 +872,7 @@ impl MePool { Err(TrySendError::Full(cmd)) => { match reserve_writer_command_slot( ¤t.tx, - self.route_runtime.me_route_blocking_send_timeout, + deadline, ) .await { diff --git a/src/transport/middle_proxy/send/close.rs b/src/transport/middle_proxy/send/close.rs index 9a8055f..df5960d 100644 --- a/src/transport/middle_proxy/send/close.rs +++ b/src/transport/middle_proxy/send/close.rs @@ -10,7 +10,7 @@ use crate::protocol::constants::{RPC_CLOSE_CONN_U32, RPC_CLOSE_EXT_U32}; use super::super::MePool; use super::super::codec::{WriterCommand, build_control_payload}; -use super::{WriterCommandReserveError, reserve_writer_command_slot}; +use super::{WriterCommandReserveError, reserve_writer_command_slot, writer_send_deadline}; const ME_CLOSE_SIGNAL_SEND_TIMEOUT: Duration = Duration::from_millis(50); @@ -22,8 +22,11 @@ impl MePool { match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) { Ok(()) => {} Err(TrySendError::Full(cmd)) => { - match reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)) - .await + match reserve_writer_command_slot( + &w.tx, + writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)), + ) + .await { Ok(permit) => { permit.send(cmd); @@ -63,9 +66,12 @@ impl MePool { match w.tx.try_send(WriterCommand::ControlAndFlush(payload)) { Ok(()) => {} Err(TrySendError::Full(cmd)) => { - let _ = reserve_writer_command_slot(&w.tx, Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)) - .await - .map(|permit| permit.send(cmd)); + let _ = reserve_writer_command_slot( + &w.tx, + writer_send_deadline(Some(ME_CLOSE_SIGNAL_SEND_TIMEOUT)), + ) + .await + .map(|permit| permit.send(cmd)); } Err(TrySendError::Closed(_)) => { debug!(conn_id, "ME close_conn skipped: writer channel closed"); diff --git a/src/transport/middle_proxy/tests/health_adversarial_tests.rs b/src/transport/middle_proxy/tests/health_adversarial_tests.rs index fdbd5f9..80cad99 100644 --- a/src/transport/middle_proxy/tests/health_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/health_adversarial_tests.rs @@ -105,6 +105,7 @@ async fn make_pool( general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -134,6 +135,7 @@ async fn insert_draining_writer( drain_deadline_epoch_secs: u64, ) { let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 6000 + writer_id as u16), @@ -143,6 +145,7 @@ async fn insert_draining_writer( contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), created_at: Instant::now() - Duration::from_secs(writer_id), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -153,7 +156,9 @@ async fn insert_draining_writer( }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); for idx in 0..bound_clients { diff --git a/src/transport/middle_proxy/tests/health_integration_tests.rs b/src/transport/middle_proxy/tests/health_integration_tests.rs index 1dc6abf..72ac165 100644 --- a/src/transport/middle_proxy/tests/health_integration_tests.rs +++ b/src/transport/middle_proxy/tests/health_integration_tests.rs @@ -103,6 +103,7 @@ async fn make_pool( general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -131,6 +132,7 @@ async fn insert_draining_writer( drain_deadline_epoch_secs: u64, ) { let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 5500 + writer_id as u16), @@ -140,6 +142,7 @@ async fn insert_draining_writer( contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), created_at: Instant::now() - Duration::from_secs(writer_id), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -149,7 +152,9 @@ async fn insert_draining_writer( allow_drain_fallback: Arc::new(AtomicBool::new(false)), }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); for idx in 0..bound_clients { let (conn_id, _rx) = pool.registry.register().await; diff --git a/src/transport/middle_proxy/tests/health_regression_tests.rs b/src/transport/middle_proxy/tests/health_regression_tests.rs index c2f5441..85f0764 100644 --- a/src/transport/middle_proxy/tests/health_regression_tests.rs +++ b/src/transport/middle_proxy/tests/health_regression_tests.rs @@ -98,6 +98,7 @@ async fn make_pool(me_pool_drain_threshold: u64) -> Arc { general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -126,6 +127,7 @@ async fn insert_draining_writer( ) -> Vec { let mut conn_ids = Vec::with_capacity(bound_clients); let (tx, _writer_rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 4500 + writer_id as u16), @@ -135,6 +137,7 @@ async fn insert_draining_writer( contour: Arc::new(AtomicU8::new(WriterContour::Draining.as_u8())), created_at: Instant::now() - Duration::from_secs(writer_id), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -144,7 +147,9 @@ async fn insert_draining_writer( allow_drain_fallback: Arc::new(AtomicBool::new(false)), }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); for idx in 0..bound_clients { let (conn_id, _rx) = pool.registry.register().await; diff --git a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs index 9125ad0..9d7fb99 100644 --- a/src/transport/middle_proxy/tests/pool_refill_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_refill_security_tests.rs @@ -87,6 +87,7 @@ async fn make_pool() -> Arc { general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, diff --git a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs index f0514a9..f7908dc 100644 --- a/src/transport/middle_proxy/tests/pool_writer_security_tests.rs +++ b/src/transport/middle_proxy/tests/pool_writer_security_tests.rs @@ -92,6 +92,7 @@ async fn make_pool() -> Arc { general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -120,6 +121,7 @@ async fn insert_writer( created_at: Instant, ) { let (tx, _rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let contour = if draining { WriterContour::Draining } else { @@ -134,6 +136,7 @@ async fn insert_writer( contour: Arc::new(AtomicU8::new(contour.as_u8())), created_at, tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -144,7 +147,9 @@ async fn insert_writer( }; pool.writers.write().await.push(writer); - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; pool.conn_count.fetch_add(1, Ordering::Relaxed); } diff --git a/src/transport/middle_proxy/tests/send_adversarial_tests.rs b/src/transport/middle_proxy/tests/send_adversarial_tests.rs index 61f7aa7..3b6e718 100644 --- a/src/transport/middle_proxy/tests/send_adversarial_tests.rs +++ b/src/transport/middle_proxy/tests/send_adversarial_tests.rs @@ -98,6 +98,7 @@ async fn make_pool() -> (Arc, Arc) { general.me_writer_pick_sample_size, MeSocksKdfPolicy::default(), general.me_writer_cmd_channel_capacity, + general.me_writer_byte_budget_bytes, general.me_route_channel_capacity, general.me_route_backpressure_enabled, general.me_route_fairshare_enabled, @@ -127,6 +128,7 @@ async fn insert_writer( register_in_registry: bool, ) -> mpsc::Receiver { let (tx, rx) = mpsc::channel::(8); + let byte_budget = pool.new_writer_byte_budget(); let writer = MeWriter { id: writer_id, addr, @@ -136,6 +138,7 @@ async fn insert_writer( contour: Arc::new(AtomicU8::new(WriterContour::Active.as_u8())), created_at: Instant::now(), tx: tx.clone(), + byte_budget: byte_budget.clone(), cancel: CancellationToken::new(), degraded: Arc::new(AtomicBool::new(false)), rtt_ema_ms_x10: Arc::new(AtomicU32::new(0)), @@ -154,7 +157,9 @@ async fn insert_writer( } pool.rebuild_endpoint_dc_map().await; if register_in_registry { - pool.registry.register_writer(writer_id, tx).await; + pool.registry + .register_writer(writer_id, tx, byte_budget) + .await; } rx }