diff --git a/src/config/defaults.rs b/src/config/defaults.rs index 3b8ae55..0f78af1 100644 --- a/src/config/defaults.rs +++ b/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 } diff --git a/src/config/load.rs b/src/config/load.rs index cf13ba9..aa97b58 100644 --- a/src/config/load.rs +++ b/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(), diff --git a/src/config/load/strict_keys.rs b/src/config/load/strict_keys.rs index 47bc3e0..21fecb6 100644 --- a/src/config/load/strict_keys.rs +++ b/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", diff --git a/src/config/tests/load_memory_envelope_tests.rs b/src/config/tests/load_memory_envelope_tests.rs index 0702c5c..3592d27 100644 --- a/src/config/tests/load_memory_envelope_tests.rs +++ b/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); diff --git a/src/config/types.rs b/src/config/types.rs index 0d4d8ab..0994a37 100644 --- a/src/config/types.rs +++ b/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(), diff --git a/src/maestro/mod.rs b/src/maestro/mod.rs index e7f349c..a848fda 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -34,6 +34,10 @@ use crate::crypto::SecureRandom; use crate::ip_tracker::UserIpTracker; use crate::network::probe::{decide_network_capabilities, log_probe_result, run_probe}; use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; +use crate::proxy::direct_buffer_budget::{ + DirectBufferBudget, resolve_direct_buffer_hard_limit, + spawn_direct_buffer_budget_controller, +}; use crate::proxy::shared_state::ProxySharedState; use crate::startup::{ COMPONENT_API_BOOTSTRAP, COMPONENT_CONFIG_LOAD, COMPONENT_DC_CONNECTIVITY_PING, @@ -473,7 +477,18 @@ async fn run_telemt_core( config.network.dns_overrides.len() ); } - let shared_state = ProxySharedState::new(); + let direct_buffer_hard_limit = resolve_direct_buffer_hard_limit( + config.general.direct_relay_buffer_budget_max_bytes, + ) + .await; + let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit); + info!( + hard_limit_bytes = direct_buffer_hard_limit, + configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes, + "Direct relay buffer budget initialized" + ); + let shared_state = + ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone()); shared_state.apply_user_enabled_config(&config.access.user_enabled); shared_state.traffic_limiter.apply_policy( config.access.user_rate_limits.clone(), @@ -882,6 +897,12 @@ async fn run_telemt_core( stats.clone(), shared_state.clone(), ); + spawn_direct_buffer_budget_controller( + direct_buffer_budget, + stats.clone(), + shared_state.clone(), + config.server.max_connections, + ); let bound = listeners::bind_listeners( &config, diff --git a/src/metrics.rs b/src/metrics.rs index 704f8b8..a7e008a 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -596,6 +596,74 @@ async fn render_metrics( stats.get_buffer_pool_in_use_gauge() ); + let direct_budget = shared_state.direct_buffer_budget.snapshot(); + let _ = writeln!( + out, + "# HELP telemt_direct_relay_buffer_budget_bytes Direct relay copy-buffer budget and memory inputs" + ); + let _ = writeln!( + out, + "# TYPE telemt_direct_relay_buffer_budget_bytes gauge" + ); + for (kind, value) in [ + ("hard_limit", direct_budget.hard_limit_bytes), + ("target", direct_budget.target_bytes), + ("reserved", direct_budget.reserved_bytes), + ("memory_total", direct_budget.memory_total_bytes), + ("memory_available", direct_budget.memory_available_bytes), + ("process_rss", direct_budget.process_rss_bytes), + ] { + let _ = writeln!( + out, + "telemt_direct_relay_buffer_budget_bytes{{kind=\"{}\"}} {}", + kind, value + ); + } + let _ = writeln!( + out, + "# HELP telemt_direct_relay_buffer_budget_events_total Direct relay buffer-budget lifecycle events" + ); + let _ = writeln!( + out, + "# TYPE telemt_direct_relay_buffer_budget_events_total counter" + ); + for (result, value) in [ + ("promotion", direct_budget.promotion_total), + ("promotion_denied", direct_budget.promotion_denied_total), + ("minimum_fallback", direct_budget.minimum_fallback_total), + ("admission_rejected", direct_budget.admission_rejected_total), + ("quiet_demotion", direct_budget.quiet_demotion_total), + ( + "write_pressure_demotion", + direct_budget.write_pressure_demotion_total, + ), + ( + "global_pressure_demotion", + direct_budget.global_pressure_demotion_total, + ), + ] { + let _ = writeln!( + out, + "telemt_direct_relay_buffer_budget_events_total{{result=\"{}\"}} {}", + result, value + ); + } + let _ = writeln!( + out, + "# HELP telemt_direct_relay_buffer_sessions Current Direct relay sessions by adaptive tier" + ); + let _ = writeln!(out, "# TYPE telemt_direct_relay_buffer_sessions gauge"); + for (tier, value) in ["base", "tier1", "tier2", "tier3"] + .into_iter() + .zip(direct_budget.tier_sessions) + { + let _ = writeln!( + out, + "telemt_direct_relay_buffer_sessions{{tier=\"{}\"}} {}", + tier, value + ); + } + let _ = writeln!( out, "# HELP telemt_tls_fetch_profile_cache_entries Current adaptive TLS fetch profile-cache entries" diff --git a/src/proxy/adaptive_buffers.rs b/src/proxy/adaptive_buffers.rs index a04f4e8..9870522 100644 --- a/src/proxy/adaptive_buffers.rs +++ b/src/proxy/adaptive_buffers.rs @@ -1,7 +1,4 @@ -#![allow(dead_code)] - -// Adaptive buffer policy is staged and retained for deterministic rollout. -// Keep definitions compiled for compatibility and security test scaffolding. +// Adaptive buffer policy shared by active Direct relay sessions. use dashmap::DashMap; use std::cmp::max; @@ -13,29 +10,42 @@ const PROFILE_TTL: Duration = Duration::from_secs(300); const THROUGHPUT_UP_BPS: f64 = 8_000_000.0; const THROUGHPUT_DOWN_BPS: f64 = 2_000_000.0; const RATIO_CONFIRM_THRESHOLD: f64 = 1.12; -const TIER1_HOLD_TICKS: u32 = 8; -const TIER2_HOLD_TICKS: u32 = 4; -const QUIET_DEMOTE_TICKS: u32 = 480; -const HARD_COOLDOWN_TICKS: u32 = 20; +const TIER1_HOLD: Duration = Duration::from_secs(2); +const TIER2_HOLD: Duration = Duration::from_secs(1); +const QUIET_DEMOTE: Duration = Duration::from_secs(120); +const HARD_COOLDOWN: Duration = Duration::from_secs(5); +const SUSTAINED_PRESSURE_DEMOTE: Duration = Duration::from_secs(30); +const PRESSURE_DEMOTE_COOLDOWN: Duration = Duration::from_secs(60); const HARD_PENDING_THRESHOLD: u32 = 3; const HARD_PARTIAL_RATIO_THRESHOLD: f64 = 0.25; +#[cfg(test)] const DIRECT_C2S_CAP_BYTES: usize = 128 * 1024; +#[cfg(test)] const DIRECT_S2C_CAP_BYTES: usize = 512 * 1024; +#[cfg(test)] const ME_FRAMES_CAP: usize = 96; +#[cfg(test)] const ME_BYTES_CAP: usize = 384 * 1024; +#[cfg(test)] const ME_DELAY_MIN_US: u64 = 150; const MAX_USER_PROFILES_ENTRIES: usize = 50_000; const MAX_USER_KEY_BYTES: usize = 512; #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +/// Per-session Direct copy-buffer capacity tier. pub enum AdaptiveTier { + /// Conservative baseline capacity. Base = 0, + /// First throughput promotion. Tier1 = 1, + /// Sustained bidirectional pressure promotion. Tier2 = 2, + /// Configured per-direction ceilings. Tier3 = 3, } impl AdaptiveTier { + /// Returns the next larger tier, saturating at `Tier3`. pub fn promote(self) -> Self { match self { Self::Base => Self::Tier1, @@ -45,6 +55,7 @@ impl AdaptiveTier { } } + /// Returns the next smaller tier, saturating at `Base`. pub fn demote(self) -> Self { match self { Self::Base => Self::Base, @@ -54,6 +65,7 @@ impl AdaptiveTier { } } + #[cfg(test)] fn ratio(self) -> (usize, usize) { match self { Self::Base => (1, 1), @@ -63,49 +75,70 @@ impl AdaptiveTier { } } + /// Returns the stable numeric tier used by bounded metrics. pub fn as_u8(self) -> u8 { self as u8 } } #[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Signal that caused an accepted adaptive tier transition. pub enum TierTransitionReason { + /// Sustained throughput and directional ratio confirmation. SoftConfirmed, + /// Short pending or partial-write pressure burst. HardPressure, + /// Sustained low-throughput period. QuietDemotion, + /// Sustained pending or partial-write pressure. + SustainedWritePressure, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] +/// Proposed transition emitted by the per-session controller. pub struct TierTransition { + /// Tier active before the observation. pub from: AdaptiveTier, + /// Tier requested after the observation. pub to: AdaptiveTier, + /// Pressure or throughput condition that requested the transition. pub reason: TierTransitionReason, } #[derive(Debug, Clone, Copy, Default)] +/// Directional byte and write-pressure deltas for one observation period. pub struct RelaySignalSample { + /// Client-to-DC bytes copied during the period. pub c2s_bytes: u64, + /// Bytes offered to DC-to-client writes during the period. pub s2c_requested_bytes: u64, + /// Bytes accepted by DC-to-client writes during the period. pub s2c_written_bytes: u64, + /// Successful DC-to-client write operations during the period. pub s2c_write_ops: u64, + /// Partial DC-to-client write operations during the period. pub s2c_partial_writes: u64, + /// Consecutive pending DC-to-client writes at sample time. pub s2c_consecutive_pending_writes: u32, } #[derive(Debug, Clone, Copy)] +/// Stateful hysteresis controller for one active Direct session. pub struct SessionAdaptiveController { tier: AdaptiveTier, max_tier_seen: AdaptiveTier, throughput_ema_bps: f64, incoming_ema_bps: f64, outgoing_ema_bps: f64, - tier1_hold_ticks: u32, - tier2_hold_ticks: u32, - quiet_ticks: u32, - hard_cooldown_ticks: u32, + tier1_hold: Duration, + tier2_hold: Duration, + quiet: Duration, + hard_cooldown: Duration, + sustained_pressure: Duration, } impl SessionAdaptiveController { + /// Creates a controller at the tier whose memory reservation was accepted. pub fn new(initial_tier: AdaptiveTier) -> Self { Self { tier: initial_tier, @@ -113,25 +146,33 @@ impl SessionAdaptiveController { throughput_ema_bps: 0.0, incoming_ema_bps: 0.0, outgoing_ema_bps: 0.0, - tier1_hold_ticks: 0, - tier2_hold_ticks: 0, - quiet_ticks: 0, - hard_cooldown_ticks: 0, + tier1_hold: Duration::ZERO, + tier2_hold: Duration::ZERO, + quiet: Duration::ZERO, + hard_cooldown: Duration::ZERO, + sustained_pressure: Duration::ZERO, } } + /// Returns the highest tier proposed by controller observations. + #[allow(dead_code)] pub fn max_tier_seen(&self) -> AdaptiveTier { self.max_tier_seen } + /// Returns the controller's current logical tier. + pub fn tier(&self) -> AdaptiveTier { + self.tier + } + + /// Observes one period and returns at most one hysteresis-controlled transition. pub fn observe(&mut self, sample: RelaySignalSample, tick_secs: f64) -> Option { if tick_secs <= f64::EPSILON { return None; } - if self.hard_cooldown_ticks > 0 { - self.hard_cooldown_ticks -= 1; - } + let tick = Duration::from_secs_f64(tick_secs); + self.hard_cooldown = self.hard_cooldown.saturating_sub(tick); let c2s_bps = (sample.c2s_bytes as f64 * 8.0) / tick_secs; let incoming_bps = (sample.s2c_requested_bytes as f64 * 8.0) / tick_secs; @@ -144,9 +185,9 @@ impl SessionAdaptiveController { let tier1_now = self.throughput_ema_bps >= THROUGHPUT_UP_BPS; if tier1_now { - self.tier1_hold_ticks = self.tier1_hold_ticks.saturating_add(1); + self.tier1_hold = self.tier1_hold.saturating_add(tick); } else { - self.tier1_hold_ticks = 0; + self.tier1_hold = Duration::ZERO; } let ratio = if self.outgoing_ema_bps <= f64::EPSILON { @@ -156,9 +197,9 @@ impl SessionAdaptiveController { }; let tier2_now = ratio >= RATIO_CONFIRM_THRESHOLD; if tier2_now { - self.tier2_hold_ticks = self.tier2_hold_ticks.saturating_add(1); + self.tier2_hold = self.tier2_hold.saturating_add(tick); } else { - self.tier2_hold_ticks = 0; + self.tier2_hold = Duration::ZERO; } let partial_ratio = if sample.s2c_write_ops == 0 { @@ -169,24 +210,37 @@ impl SessionAdaptiveController { let hard_now = sample.s2c_consecutive_pending_writes >= HARD_PENDING_THRESHOLD || partial_ratio >= HARD_PARTIAL_RATIO_THRESHOLD; - if hard_now && self.hard_cooldown_ticks == 0 { - return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN_TICKS); + if hard_now { + self.sustained_pressure = self.sustained_pressure.saturating_add(tick); + if self.sustained_pressure >= SUSTAINED_PRESSURE_DEMOTE { + self.sustained_pressure = Duration::ZERO; + return self.demote( + TierTransitionReason::SustainedWritePressure, + PRESSURE_DEMOTE_COOLDOWN, + ); + } + } else { + self.sustained_pressure = Duration::ZERO; } - if self.tier1_hold_ticks >= TIER1_HOLD_TICKS && self.tier2_hold_ticks >= TIER2_HOLD_TICKS { - return self.promote(TierTransitionReason::SoftConfirmed, 0); + if hard_now && self.hard_cooldown.is_zero() { + return self.promote(TierTransitionReason::HardPressure, HARD_COOLDOWN); + } + + if self.tier1_hold >= TIER1_HOLD && self.tier2_hold >= TIER2_HOLD { + return self.promote(TierTransitionReason::SoftConfirmed, Duration::ZERO); } let demote_candidate = self.throughput_ema_bps < THROUGHPUT_DOWN_BPS && !tier2_now && !hard_now; if demote_candidate { - self.quiet_ticks = self.quiet_ticks.saturating_add(1); - if self.quiet_ticks >= QUIET_DEMOTE_TICKS { - self.quiet_ticks = 0; - return self.demote(TierTransitionReason::QuietDemotion); + self.quiet = self.quiet.saturating_add(tick); + if self.quiet >= QUIET_DEMOTE { + self.quiet = Duration::ZERO; + return self.demote(TierTransitionReason::QuietDemotion, Duration::ZERO); } } else { - self.quiet_ticks = 0; + self.quiet = Duration::ZERO; } None @@ -195,7 +249,7 @@ impl SessionAdaptiveController { fn promote( &mut self, reason: TierTransitionReason, - hard_cooldown_ticks: u32, + hard_cooldown: Duration, ) -> Option { let from = self.tier; let to = from.promote(); @@ -204,22 +258,27 @@ impl SessionAdaptiveController { } self.tier = to; self.max_tier_seen = max(self.max_tier_seen, to); - self.hard_cooldown_ticks = hard_cooldown_ticks; - self.tier1_hold_ticks = 0; - self.tier2_hold_ticks = 0; - self.quiet_ticks = 0; + self.hard_cooldown = hard_cooldown; + self.tier1_hold = Duration::ZERO; + self.tier2_hold = Duration::ZERO; + self.quiet = Duration::ZERO; Some(TierTransition { from, to, reason }) } - fn demote(&mut self, reason: TierTransitionReason) -> Option { + fn demote( + &mut self, + reason: TierTransitionReason, + hard_cooldown: Duration, + ) -> Option { let from = self.tier; let to = from.demote(); if from == to { return None; } self.tier = to; - self.tier1_hold_ticks = 0; - self.tier2_hold_ticks = 0; + self.hard_cooldown = hard_cooldown; + self.tier1_hold = Duration::ZERO; + self.tier2_hold = Duration::ZERO; Some(TierTransition { from, to, reason }) } } @@ -235,6 +294,8 @@ fn profiles() -> &'static DashMap { USER_PROFILES.get_or_init(DashMap::new) } +/// Returns a fresh user's recent successful Direct tier, or `Base` when stale. +#[allow(dead_code)] pub fn seed_tier_for_user(user: &str) -> AdaptiveTier { if user.len() > MAX_USER_KEY_BYTES { return AdaptiveTier::Base; @@ -253,6 +314,8 @@ pub fn seed_tier_for_user(user: &str) -> AdaptiveTier { AdaptiveTier::Base } +/// Records the highest successfully allocated tier for bounded session seeding. +#[allow(dead_code)] pub fn record_user_tier(user: &str, tier: AdaptiveTier) { if user.len() > MAX_USER_KEY_BYTES { return; @@ -282,6 +345,8 @@ pub fn record_user_tier(user: &str, tier: AdaptiveTier) { } } +#[cfg(test)] +/// Returns the legacy staged scaling policy retained by security fixtures. pub fn direct_copy_buffers_for_tier( tier: AdaptiveTier, base_c2s: usize, @@ -294,6 +359,32 @@ pub fn direct_copy_buffers_for_tier( ) } +/// Maps an adaptive tier to independent capacities within configured ceilings. +pub(crate) fn direct_copy_buffers_for_tier_with_ceilings( + tier: AdaptiveTier, + base_c2s: usize, + base_s2c: usize, + ceiling_c2s: usize, + ceiling_s2c: usize, +) -> (usize, usize) { + ( + direct_direction_size(tier, base_c2s, ceiling_c2s), + direct_direction_size(tier, base_s2c, ceiling_s2c), + ) +} + +fn direct_direction_size(tier: AdaptiveTier, base: usize, ceiling: usize) -> usize { + let target = match tier { + AdaptiveTier::Base => base, + AdaptiveTier::Tier1 => ceiling / 4, + AdaptiveTier::Tier2 => ceiling / 2, + AdaptiveTier::Tier3 => ceiling, + }; + target.max(base).min(ceiling.max(base)).max(1) +} + +#[cfg(test)] +/// Returns the staged Middle-End flush policy retained by security fixtures. pub fn me_flush_policy_for_tier( tier: AdaptiveTier, base_frames: usize, @@ -323,6 +414,7 @@ fn ema(prev: f64, value: f64) -> f64 { } } +#[cfg(test)] fn scale(base: usize, numerator: usize, denominator: usize, cap: usize) -> usize { let scaled = base .saturating_mul(numerator) @@ -338,6 +430,10 @@ mod adaptive_buffers_security_tests; #[path = "tests/adaptive_buffers_record_race_security_tests.rs"] mod adaptive_buffers_record_race_security_tests; +#[cfg(test)] +#[path = "tests/adaptive_direct_budget_policy_tests.rs"] +mod adaptive_direct_budget_policy_tests; + #[cfg(test)] mod tests { use super::*; @@ -396,7 +492,7 @@ mod tests { fn test_quiet_demotion_is_slow_and_stepwise() { let mut ctrl = SessionAdaptiveController::new(AdaptiveTier::Tier2); let mut demotion = None; - for _ in 0..QUIET_DEMOTE_TICKS { + for _ in 0..480 { demotion = ctrl.observe(sample(1, 1, 1, 1, 0, 0), 0.25); } diff --git a/src/proxy/direct_buffer_budget.rs b/src/proxy/direct_buffer_budget.rs new file mode 100644 index 0000000..bb25a04 --- /dev/null +++ b/src/proxy/direct_buffer_budget.rs @@ -0,0 +1,549 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; + +use tokio::sync::watch; + +use crate::stats::Stats; + +use super::shared_state::ProxySharedState; + +/// Accounting granularity for process-wide Direct copy-buffer reservations. +pub(crate) const DIRECT_BUFFER_UNIT_BYTES: usize = 4 * 1024; +/// Minimum client-to-DC copy-buffer capacity for one Direct session. +pub(crate) const DIRECT_BASE_C2S_BYTES: usize = 4 * 1024; +/// Minimum DC-to-client copy-buffer capacity for one Direct session. +pub(crate) const DIRECT_BASE_S2C_BYTES: usize = 8 * 1024; + +const AUTO_HARD_MIN_BYTES: usize = 64 * 1024 * 1024; +const AUTO_HARD_MAX_BYTES: usize = 2 * 1024 * 1024 * 1024; +const AUTO_HARD_FALLBACK_BYTES: usize = 512 * 1024 * 1024; +const TARGET_FLOOR_MIN_BYTES: usize = 16 * 1024 * 1024; +const CONTROL_INTERVAL: Duration = Duration::from_secs(1); +const HEALTHY_RECOVERY_SAMPLES: u8 = 30; + +#[derive(Debug, Clone, Copy, Default)] +/// Lock-free observability snapshot of the Direct copy-buffer envelope. +pub(crate) struct DirectBufferBudgetSnapshot { + /// Absolute process-wide copy-buffer ceiling. + pub(crate) hard_limit_bytes: u64, + /// Current pressure-adjusted promotion target. + pub(crate) target_bytes: u64, + /// Bytes currently covered by active session leases. + pub(crate) reserved_bytes: u64, + /// Effective host or cgroup memory limit. + pub(crate) memory_total_bytes: u64, + /// Effective host or cgroup memory headroom. + pub(crate) memory_available_bytes: u64, + /// Current process resident set size. + pub(crate) process_rss_bytes: u64, + /// Successful tier growth reservations. + pub(crate) promotion_total: u64, + /// Tier growth attempts rejected by the adaptive target. + pub(crate) promotion_denied_total: u64, + /// Sessions admitted at minimum size above the adaptive target. + pub(crate) minimum_fallback_total: u64, + /// Sessions rejected by the absolute ceiling. + pub(crate) admission_rejected_total: u64, + /// Quiet-period tier reductions. + pub(crate) quiet_demotion_total: u64, + /// Sustained write-pressure tier reductions. + pub(crate) write_pressure_demotion_total: u64, + /// Process-wide pressure tier reductions. + pub(crate) global_pressure_demotion_total: u64, + /// Current sessions for Base through Tier3. + pub(crate) tier_sessions: [u64; 4], +} + +#[derive(Debug, Clone, Copy, Default)] +struct SystemMemorySample { + total_bytes: u64, + available_bytes: u64, + process_rss_bytes: u64, +} + +/// Process-wide hard envelope and adaptive target for Direct copy buffers. +pub(crate) struct DirectBufferBudget { + hard_limit_bytes: u64, + target_bytes: AtomicU64, + reserved_bytes: AtomicU64, + pressure_generation: AtomicU64, + pressure_tx: watch::Sender, + memory_total_bytes: AtomicU64, + memory_available_bytes: AtomicU64, + process_rss_bytes: AtomicU64, + promotion_total: AtomicU64, + promotion_denied_total: AtomicU64, + minimum_fallback_total: AtomicU64, + admission_rejected_total: AtomicU64, + quiet_demotion_total: AtomicU64, + write_pressure_demotion_total: AtomicU64, + global_pressure_demotion_total: AtomicU64, + tier_sessions: [AtomicU64; 4], +} + +impl DirectBufferBudget { + /// Creates an envelope with a fixed absolute ceiling. + pub(crate) fn new(hard_limit_bytes: usize) -> Arc { + let hard_limit_bytes = align_down(hard_limit_bytes.max(DIRECT_BUFFER_UNIT_BYTES)) as u64; + let (pressure_tx, _) = watch::channel(0); + Arc::new(Self { + hard_limit_bytes, + target_bytes: AtomicU64::new(hard_limit_bytes), + reserved_bytes: AtomicU64::new(0), + pressure_generation: AtomicU64::new(0), + pressure_tx, + memory_total_bytes: AtomicU64::new(0), + memory_available_bytes: AtomicU64::new(0), + process_rss_bytes: AtomicU64::new(0), + promotion_total: AtomicU64::new(0), + promotion_denied_total: AtomicU64::new(0), + minimum_fallback_total: AtomicU64::new(0), + admission_rejected_total: AtomicU64::new(0), + quiet_demotion_total: AtomicU64::new(0), + write_pressure_demotion_total: AtomicU64::new(0), + global_pressure_demotion_total: AtomicU64::new(0), + tier_sessions: std::array::from_fn(|_| AtomicU64::new(0)), + }) + } + + /// Returns the current pressure-adjusted reservation target. + pub(crate) fn target_bytes(&self) -> usize { + self.target_bytes.load(Ordering::Relaxed) as usize + } + + /// Subscribes to target reductions that require prompt session demotion. + pub(crate) fn subscribe_pressure(&self) -> watch::Receiver { + self.pressure_tx.subscribe() + } + + /// Reserves bytes against either the adaptive target or the absolute ceiling. + pub(crate) fn try_reserve( + self: &Arc, + bytes: usize, + allow_above_target: bool, + ) -> Option { + let bytes = align_up(bytes) as u64; + let limit = if allow_above_target { + self.hard_limit_bytes + } else { + self.target_bytes + .load(Ordering::Relaxed) + .min(self.hard_limit_bytes) + }; + if !self.try_add_reserved(bytes, limit) { + return None; + } + self.tier_sessions[0].fetch_add(1, Ordering::Relaxed); + Some(DirectBufferLease { + budget: Arc::clone(self), + reserved_bytes: bytes, + tier: 0, + }) + } + + fn try_add_reserved(&self, bytes: u64, limit: u64) -> bool { + let mut current = self.reserved_bytes.load(Ordering::Acquire); + loop { + if bytes > limit.saturating_sub(current) { + return false; + } + match self.reserved_bytes.compare_exchange_weak( + current, + current + bytes, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(observed) => current = observed, + } + } + } + + fn target_floor_bytes(&self) -> u64 { + (self.hard_limit_bytes / 8) + .max(TARGET_FLOOR_MIN_BYTES as u64) + .min(self.hard_limit_bytes) + } + + fn set_target_bytes(&self, target: u64) { + let target = align_down( + target + .clamp(self.target_floor_bytes(), self.hard_limit_bytes) + as usize, + ) as u64; + let previous = self.target_bytes.swap(target, Ordering::AcqRel); + if target < previous { + let generation = self + .pressure_generation + .fetch_add(1, Ordering::AcqRel) + .wrapping_add(1); + self.pressure_tx.send_replace(generation); + } + } + + fn update_system_sample(&self, sample: SystemMemorySample) { + self.memory_total_bytes + .store(sample.total_bytes, Ordering::Relaxed); + self.memory_available_bytes + .store(sample.available_bytes, Ordering::Relaxed); + self.process_rss_bytes + .store(sample.process_rss_bytes, Ordering::Relaxed); + } + + /// Records a session that had to bypass the adaptive target at minimum size. + pub(crate) fn increment_minimum_fallback(&self) { + self.minimum_fallback_total + .fetch_add(1, Ordering::Relaxed); + } + + /// Records a session rejected because the absolute ceiling was exhausted. + pub(crate) fn increment_admission_rejected(&self) { + self.admission_rejected_total + .fetch_add(1, Ordering::Relaxed); + } + + /// Records a tier reduction after sustained low throughput. + pub(crate) fn increment_quiet_demotion(&self) { + self.quiet_demotion_total.fetch_add(1, Ordering::Relaxed); + } + + /// Records a tier reduction after sustained partial or pending writes. + pub(crate) fn increment_write_pressure_demotion(&self) { + self.write_pressure_demotion_total + .fetch_add(1, Ordering::Relaxed); + } + + /// Records a tier reduction requested by the process-wide controller. + pub(crate) fn increment_global_pressure_demotion(&self) { + self.global_pressure_demotion_total + .fetch_add(1, Ordering::Relaxed); + } + + /// Captures all bounded metrics without allocating or locking. + pub(crate) fn snapshot(&self) -> DirectBufferBudgetSnapshot { + DirectBufferBudgetSnapshot { + hard_limit_bytes: self.hard_limit_bytes, + target_bytes: self.target_bytes.load(Ordering::Relaxed), + reserved_bytes: self.reserved_bytes.load(Ordering::Relaxed), + memory_total_bytes: self.memory_total_bytes.load(Ordering::Relaxed), + memory_available_bytes: self.memory_available_bytes.load(Ordering::Relaxed), + process_rss_bytes: self.process_rss_bytes.load(Ordering::Relaxed), + promotion_total: self.promotion_total.load(Ordering::Relaxed), + promotion_denied_total: self.promotion_denied_total.load(Ordering::Relaxed), + minimum_fallback_total: self.minimum_fallback_total.load(Ordering::Relaxed), + admission_rejected_total: self.admission_rejected_total.load(Ordering::Relaxed), + quiet_demotion_total: self.quiet_demotion_total.load(Ordering::Relaxed), + write_pressure_demotion_total: self + .write_pressure_demotion_total + .load(Ordering::Relaxed), + global_pressure_demotion_total: self + .global_pressure_demotion_total + .load(Ordering::Relaxed), + tier_sessions: std::array::from_fn(|index| { + self.tier_sessions[index].load(Ordering::Relaxed) + }), + } + } +} + +/// Returns the conservative ceiling used when memory discovery is unavailable. +pub(crate) fn fallback_direct_buffer_hard_limit() -> usize { + AUTO_HARD_FALLBACK_BYTES +} + +/// RAII ownership of all copy-buffer bytes retained by one Direct session. +pub(crate) struct DirectBufferLease { + budget: Arc, + reserved_bytes: u64, + tier: usize, +} + +impl DirectBufferLease { + /// Returns the currently covered allocation rounded to accounting units. + pub(crate) fn reserved_bytes(&self) -> usize { + self.reserved_bytes as usize + } + + /// Attempts to cover a larger tier before its buffers are resized. + pub(crate) fn try_grow_to(&mut self, bytes: usize) -> bool { + let bytes = align_up(bytes) as u64; + if bytes <= self.reserved_bytes { + return true; + } + let delta = bytes - self.reserved_bytes; + let limit = self + .budget + .target_bytes + .load(Ordering::Relaxed) + .min(self.budget.hard_limit_bytes); + if !self.budget.try_add_reserved(delta, limit) { + self.budget + .promotion_denied_total + .fetch_add(1, Ordering::Relaxed); + return false; + } + self.reserved_bytes = bytes; + self.budget.promotion_total.fetch_add(1, Ordering::Relaxed); + true + } + + /// Releases bytes only after both directional buffers report smaller coverage. + pub(crate) fn shrink_to(&mut self, bytes: usize) { + let bytes = align_up(bytes) as u64; + if bytes >= self.reserved_bytes { + return; + } + let released = self.reserved_bytes - bytes; + self.reserved_bytes = bytes; + self.budget + .reserved_bytes + .fetch_sub(released, Ordering::AcqRel); + } + + /// Updates bounded per-tier session gauges for an accepted transition. + pub(crate) fn set_tier(&mut self, tier: usize) { + let tier = tier.min(self.budget.tier_sessions.len() - 1); + if tier == self.tier { + return; + } + decrement_saturating(&self.budget.tier_sessions[self.tier]); + self.budget.tier_sessions[tier].fetch_add(1, Ordering::Relaxed); + self.tier = tier; + } +} + +impl Drop for DirectBufferLease { + fn drop(&mut self) { + self.budget + .reserved_bytes + .fetch_sub(self.reserved_bytes, Ordering::AcqRel); + decrement_saturating(&self.budget.tier_sessions[self.tier]); + } +} + +/// Resolves the startup hard ceiling from config, cgroup, and host memory. +pub(crate) async fn resolve_direct_buffer_hard_limit(configured: usize) -> usize { + if configured != 0 { + return align_down(configured); + } + let sample = read_system_memory_sample().await; + if sample.total_bytes == 0 { + return AUTO_HARD_FALLBACK_BYTES; + } + let derived = (sample.total_bytes / 4) + .clamp(AUTO_HARD_MIN_BYTES as u64, AUTO_HARD_MAX_BYTES as u64) + .min(sample.total_bytes); + align_down(derived as usize).max(DIRECT_BUFFER_UNIT_BYTES) +} + +/// Starts the single control-plane task that adjusts the runtime target. +pub(crate) fn spawn_direct_buffer_budget_controller( + budget: Arc, + stats: Arc, + shared: Arc, + max_connections: u32, +) { + tokio::spawn(async move { + let mut interval = tokio::time::interval(CONTROL_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + let mut healthy_streak = 0u8; + let mut previous_denied = 0u64; + let mut previous_fallback = 0u64; + let mut previous_rejected = 0u64; + + loop { + interval.tick().await; + let sample = read_system_memory_sample().await; + budget.update_system_sample(sample); + + let snapshot = budget.snapshot(); + let denied_delta = snapshot.promotion_denied_total.saturating_sub(previous_denied); + previous_denied = snapshot.promotion_denied_total; + let fallback_delta = snapshot + .minimum_fallback_total + .saturating_sub(previous_fallback); + previous_fallback = snapshot.minimum_fallback_total; + let rejected_delta = snapshot + .admission_rejected_total + .saturating_sub(previous_rejected); + previous_rejected = snapshot.admission_rejected_total; + + let connection_pct = connection_fill_pct(stats.as_ref(), max_connections); + let memory_available_pct = percentage(sample.available_bytes, sample.total_bytes); + let target_utilization_pct = percentage(snapshot.reserved_bytes, snapshot.target_bytes); + let pressure = shared.conntrack_pressure_active() + || connection_pct.is_some_and(|value| value >= 85) + || memory_available_pct.is_some_and(|value| value <= 15) + || target_utilization_pct.is_some_and(|value| value >= 90) + || denied_delta > 0 + || fallback_delta > 0 + || rejected_delta > 0; + + let headroom_target = if sample.total_bytes == 0 { + snapshot.hard_limit_bytes + } else { + snapshot + .reserved_bytes + .saturating_add(sample.available_bytes / 4) + .min(snapshot.hard_limit_bytes) + }; + + if pressure { + healthy_streak = 0; + let reduced = snapshot.target_bytes.saturating_mul(3) / 4; + budget.set_target_bytes(reduced.min(headroom_target)); + continue; + } + + let healthy = memory_available_pct.is_none_or(|value| value >= 30) + && connection_pct.is_none_or(|value| value <= 70); + if !healthy { + healthy_streak = 0; + if headroom_target < snapshot.target_bytes { + budget.set_target_bytes(headroom_target); + } + continue; + } + + healthy_streak = healthy_streak.saturating_add(1); + if healthy_streak >= HEALTHY_RECOVERY_SAMPLES { + healthy_streak = 0; + let increment = (snapshot.target_bytes / 16).max(4 * 1024 * 1024); + budget.set_target_bytes( + snapshot + .target_bytes + .saturating_add(increment) + .min(headroom_target), + ); + } + } + }); +} + +fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option { + if max_connections == 0 { + return None; + } + Some( + ((stats + .get_current_connections_total() + .saturating_mul(100)) + / u64::from(max_connections)) + .min(100) as u8, + ) +} + +fn percentage(value: u64, total: u64) -> Option { + if total == 0 { + return None; + } + Some(((value.saturating_mul(100)) / total).min(100) as u8) +} + +async fn read_system_memory_sample() -> SystemMemorySample { + #[cfg(target_os = "linux")] + { + let meminfo = tokio::fs::read_to_string("/proc/meminfo") + .await + .unwrap_or_default(); + let status = tokio::fs::read_to_string("/proc/self/status") + .await + .unwrap_or_default(); + let host_total = parse_kib_field(&meminfo, "MemTotal:"); + let host_available = parse_kib_field(&meminfo, "MemAvailable:"); + let process_rss = parse_kib_field(&status, "VmRSS:"); + + let cgroup_v2_max = read_cgroup_limit("/sys/fs/cgroup/memory.max").await; + let cgroup_v2_current = read_u64_file("/sys/fs/cgroup/memory.current").await; + let cgroup_v1_max = read_cgroup_limit("/sys/fs/cgroup/memory/memory.limit_in_bytes").await; + let cgroup_v1_current = + read_u64_file("/sys/fs/cgroup/memory/memory.usage_in_bytes").await; + let cgroup_max = cgroup_v2_max.or(cgroup_v1_max); + let cgroup_current = cgroup_v2_current.or(cgroup_v1_current); + + let total = match (host_total, cgroup_max) { + (0, Some(limit)) => limit, + (host, Some(limit)) => host.min(limit), + (host, None) => host, + }; + let cgroup_available = cgroup_max.zip(cgroup_current).map(|(limit, current)| { + limit.saturating_sub(current) + }); + let available = match (host_available, cgroup_available) { + (0, Some(value)) => value, + (host, Some(value)) => host.min(value), + (host, None) => host, + }; + return SystemMemorySample { + total_bytes: total, + available_bytes: available, + process_rss_bytes: process_rss, + }; + } + #[cfg(not(target_os = "linux"))] + { + SystemMemorySample::default() + } +} + +#[cfg(target_os = "linux")] +async fn read_cgroup_limit(path: &str) -> Option { + let raw = tokio::fs::read_to_string(path).await.ok()?; + let raw = raw.trim(); + if raw == "max" { + return None; + } + let value = raw.parse::().ok()?; + (value < (1u64 << 60)).then_some(value) +} + +#[cfg(target_os = "linux")] +async fn read_u64_file(path: &str) -> Option { + tokio::fs::read_to_string(path) + .await + .ok()? + .trim() + .parse() + .ok() +} + +#[cfg(target_os = "linux")] +fn parse_kib_field(raw: &str, key: &str) -> u64 { + raw.lines() + .find_map(|line| { + let value = line.strip_prefix(key)?.split_whitespace().next()?; + value.parse::().ok() + }) + .unwrap_or(0) + .saturating_mul(1024) +} + +fn align_up(bytes: usize) -> usize { + bytes + .div_ceil(DIRECT_BUFFER_UNIT_BYTES) + .saturating_mul(DIRECT_BUFFER_UNIT_BYTES) +} + +fn align_down(bytes: usize) -> usize { + bytes / DIRECT_BUFFER_UNIT_BYTES * DIRECT_BUFFER_UNIT_BYTES +} + +fn decrement_saturating(value: &AtomicU64) { + let mut current = value.load(Ordering::Relaxed); + while current != 0 { + match value.compare_exchange_weak( + current, + current - 1, + Ordering::Relaxed, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(observed) => current = observed, + } + } +} + +#[cfg(test)] +#[path = "tests/direct_buffer_budget_tests.rs"] +mod tests; diff --git a/src/proxy/direct_relay.rs b/src/proxy/direct_relay.rs index 4f7cb20..5f46272 100644 --- a/src/proxy/direct_relay.rs +++ b/src/proxy/direct_relay.rs @@ -345,21 +345,21 @@ where } else { Duration::from_secs(1800) }; - let relay_result = - crate::proxy::relay::relay_bidirectional_with_activity_timeout_lease_and_cancel( + let relay_result = crate::proxy::relay::relay_direct_adaptive( client_reader, client_writer, tg_reader, tg_writer, config.general.direct_relay_copy_buf_c2s_bytes, config.general.direct_relay_copy_buf_s2c_bytes, + config.server.max_connections, user, Arc::clone(&stats), config.access.user_data_quota.get(user).copied(), - buffer_pool, traffic_lease, relay_activity_timeout, session_cancel.clone(), + Arc::clone(&shared.direct_buffer_budget), ); tokio::pin!(relay_result); let relay_result = loop { diff --git a/src/proxy/mod.rs b/src/proxy/mod.rs index 4e1827e..904fff8 100644 --- a/src/proxy/mod.rs +++ b/src/proxy/mod.rs @@ -60,6 +60,8 @@ pub mod adaptive_buffers; pub mod client; +// Process-wide Direct relay copy-buffer ownership and pressure policy. +pub(crate) mod direct_buffer_budget; pub mod direct_relay; pub mod handshake; pub mod masking; diff --git a/src/proxy/relay.rs b/src/proxy/relay.rs index 36af33a..7b6d6c9 100644 --- a/src/proxy/relay.rs +++ b/src/proxy/relay.rs @@ -85,6 +85,9 @@ fn watchdog_delta(current: u64, previous: u64) -> u64 { } mod io; +mod adaptive_copy; + +pub(crate) use self::adaptive_copy::relay_direct_adaptive; use self::io::{CombinedStream, SharedCounters, StatsIo, is_quota_io_error}; #[cfg(test)] @@ -217,6 +220,7 @@ where .await } +#[allow(dead_code)] pub async fn relay_bidirectional_with_activity_timeout_lease_and_cancel( client_reader: CR, client_writer: CW, diff --git a/src/proxy/relay/adaptive_copy.rs b/src/proxy/relay/adaptive_copy.rs new file mode 100644 index 0000000..c267343 --- /dev/null +++ b/src/proxy/relay/adaptive_copy.rs @@ -0,0 +1,532 @@ +use std::io; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::task::{Context, Poll}; +use std::time::Duration; + +use tokio::io::{AsyncBufRead, AsyncRead, AsyncWrite, AsyncWriteExt, ReadBuf, copy_buf}; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +use crate::error::{ProxyError, Result}; +use crate::proxy::adaptive_buffers::{ + AdaptiveTier, RelaySignalSample, SessionAdaptiveController, TierTransitionReason, + direct_copy_buffers_for_tier_with_ceilings, +}; +use crate::proxy::direct_buffer_budget::{ + DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES, DirectBufferBudget, DirectBufferLease, +}; +use crate::proxy::traffic_limiter::TrafficLease; +use crate::stats::Stats; + +use super::WATCHDOG_INTERVAL; +use super::io::{SharedCounters, StatsIo, is_quota_io_error}; +use super::watchdog_delta; + +mod write_pressure; + +use self::write_pressure::WritePressureIo; + +struct AdaptiveBufferState { + desired_bytes: AtomicUsize, + actual_bytes: AtomicUsize, +} + +impl AdaptiveBufferState { + fn new(bytes: usize) -> Arc { + Arc::new(Self { + desired_bytes: AtomicUsize::new(bytes.max(1)), + actual_bytes: AtomicUsize::new(bytes.max(1)), + }) + } +} + +struct AdaptiveBufReader { + inner: R, + buffer: Box<[u8]>, + pos: usize, + cap: usize, + state: Arc, +} + +impl AdaptiveBufReader { + fn new(inner: R, state: Arc) -> Self { + let bytes = state.actual_bytes.load(Ordering::Relaxed).max(1); + Self { + inner, + buffer: vec![0; bytes].into_boxed_slice(), + pos: 0, + cap: 0, + state, + } + } + + fn resize_if_drained(&mut self) { + if self.pos != self.cap { + return; + } + let desired = self.state.desired_bytes.load(Ordering::Acquire).max(1); + if desired == self.buffer.len() { + return; + } + self.buffer = vec![0; desired].into_boxed_slice(); + self.pos = 0; + self.cap = 0; + self.state.actual_bytes.store(desired, Ordering::Release); + } +} + +impl AsyncRead for AdaptiveBufReader { + fn poll_read( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + output: &mut ReadBuf<'_>, + ) -> Poll> { + let this = self.get_mut(); + if this.pos < this.cap { + let available = &this.buffer[this.pos..this.cap]; + let copied = available.len().min(output.remaining()); + output.put_slice(&available[..copied]); + this.pos += copied; + return Poll::Ready(Ok(())); + } + this.resize_if_drained(); + Pin::new(&mut this.inner).poll_read(cx, output) + } +} + +impl AsyncBufRead for AdaptiveBufReader { + fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.pos < this.cap { + return Poll::Ready(Ok(&this.buffer[this.pos..this.cap])); + } + + this.resize_if_drained(); + let mut read_buf = ReadBuf::new(&mut this.buffer); + match Pin::new(&mut this.inner).poll_read(cx, &mut read_buf) { + Poll::Ready(Ok(())) => { + this.pos = 0; + this.cap = read_buf.filled().len(); + Poll::Ready(Ok(&this.buffer[..this.cap])) + } + Poll::Ready(Err(error)) => Poll::Ready(Err(error)), + Poll::Pending => Poll::Pending, + } + } + + fn consume(self: Pin<&mut Self>, amount: usize) { + let this = self.get_mut(); + this.pos = this.pos.saturating_add(amount).min(this.cap); + } +} + +enum AdaptiveRelayOutcome { + Copy(io::Result<(u64, u64)>), + ActivityTimeout, + UserDisabled, +} + +#[allow(clippy::too_many_arguments)] +/// Relays one Direct session with independently resizable directional buffers. +pub(crate) async fn relay_direct_adaptive( + client_reader: CR, + client_writer: CW, + server_reader: SR, + server_writer: SW, + ceiling_c2s_bytes: usize, + ceiling_s2c_bytes: usize, + max_connections: u32, + user: &str, + stats: Arc, + quota_limit: Option, + traffic_lease: Option>, + activity_timeout: Duration, + session_cancel: CancellationToken, + budget: Arc, +) -> Result<()> +where + CR: AsyncRead + Unpin + Send + 'static, + CW: AsyncWrite + Unpin + Send + 'static, + SR: AsyncRead + Unpin + Send + 'static, + SW: AsyncWrite + Unpin + Send + 'static, +{ + let activity_timeout = activity_timeout.max(Duration::from_secs(1)); + let epoch = Instant::now(); + let counters = Arc::new(SharedCounters::new()); + let quota_exceeded = Arc::new(AtomicBool::new(false)); + let user_owned = user.to_string(); + + let (base_c2s, base_s2c) = initial_base_sizes( + ceiling_c2s_bytes, + ceiling_s2c_bytes, + max_connections, + budget.target_bytes(), + ); + let base_total = base_c2s.saturating_add(base_s2c); + let mut lease = match budget.try_reserve(base_total, false) { + Some(lease) => lease, + None => { + let minimum_total = DIRECT_BASE_C2S_BYTES + DIRECT_BASE_S2C_BYTES; + match budget.try_reserve(minimum_total, true) { + Some(lease) => { + budget.increment_minimum_fallback(); + lease + } + None => { + budget.increment_admission_rejected(); + return Err(ProxyError::Proxy( + "Direct relay buffer pressure: budget exhausted".to_string(), + )); + } + } + } + }; + + let effective_base = if lease.reserved_bytes() < base_total { + (DIRECT_BASE_C2S_BYTES, DIRECT_BASE_S2C_BYTES) + } else { + (base_c2s, base_s2c) + }; + let c2s_state = AdaptiveBufferState::new(effective_base.0); + let s2c_state = AdaptiveBufferState::new(effective_base.1); + + let mut controller = SessionAdaptiveController::new(AdaptiveTier::Base); + + let c2s_client = StatsIo::new_with_traffic_lease( + client_reader, + Arc::clone(&counters), + Arc::clone(&stats), + user_owned.clone(), + traffic_lease.clone(), + quota_limit, + Arc::clone("a_exceeded), + epoch, + ); + let client_writer = StatsIo::new_with_traffic_lease( + client_writer, + Arc::clone(&counters), + Arc::clone(&stats), + user_owned.clone(), + traffic_lease, + quota_limit, + Arc::clone("a_exceeded), + epoch, + ); + let mut client_writer = WritePressureIo::new(client_writer, Arc::clone(&counters)); + let mut c2s_reader = AdaptiveBufReader::new(c2s_client, Arc::clone(&c2s_state)); + let mut s2c_reader = AdaptiveBufReader::new(server_reader, Arc::clone(&s2c_state)); + let mut server_writer = server_writer; + let mut pressure_rx = budget.subscribe_pressure(); + + let relay_outcome = { + let copy = async { + let c2s = async { + let copied = copy_buf(&mut c2s_reader, &mut server_writer).await?; + server_writer.shutdown().await?; + Ok::(copied) + }; + let s2c = async { + let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?; + client_writer.shutdown().await?; + Ok::(copied) + }; + tokio::try_join!(c2s, s2c) + }; + tokio::pin!(copy); + + let mut interval = tokio::time::interval(WATCHDOG_INTERVAL); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval.tick().await; + let mut previous = RelaySignalSample::default(); + let mut previous_log_c2s = 0u64; + let mut previous_log_s2c = 0u64; + let mut previous_sample_at = epoch; + + loop { + tokio::select! { + result = &mut copy => break AdaptiveRelayOutcome::Copy(result), + _ = session_cancel.cancelled() => break AdaptiveRelayOutcome::UserDisabled, + changed = pressure_rx.changed() => { + if changed.is_ok() { + apply_global_pressure_demotion( + &mut controller, + &mut lease, + &c2s_state, + &s2c_state, + effective_base, + (ceiling_c2s_bytes, ceiling_s2c_bytes), + budget.as_ref(), + ); + reconcile_reservation(&mut lease, &c2s_state, &s2c_state); + } + } + _ = interval.tick() => { + let now = Instant::now(); + let idle = counters.idle_duration(now, epoch); + if quota_exceeded.load(Ordering::Acquire) { + warn!(user = %user_owned, "User data quota reached, closing relay"); + break AdaptiveRelayOutcome::ActivityTimeout; + } + if idle >= activity_timeout { + warn!( + user = %user_owned, + c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed), + s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed), + idle_secs = idle.as_secs(), + "Activity timeout" + ); + break AdaptiveRelayOutcome::ActivityTimeout; + } + + let sample = current_sample(counters.as_ref()); + let c2s_delta = watchdog_delta(sample.c2s_bytes, previous_log_c2s); + let s2c_delta = watchdog_delta(sample.s2c_written_bytes, previous_log_s2c); + if c2s_delta > 0 || s2c_delta > 0 { + let secs = now.saturating_duration_since(previous_sample_at).as_secs_f64(); + debug!( + user = %user_owned, + c2s_kbps = (c2s_delta as f64 / secs / 1024.0) as u64, + s2c_kbps = (s2c_delta as f64 / secs / 1024.0) as u64, + c2s_total = sample.c2s_bytes, + s2c_total = sample.s2c_written_bytes, + "Relay active" + ); + } + + let delta = sample_delta(sample, previous); + let tick_secs = now.saturating_duration_since(previous_sample_at).as_secs_f64(); + if let Some(transition) = controller.observe(delta, tick_secs) { + apply_controller_transition( + transition, + &mut controller, + &mut lease, + &c2s_state, + &s2c_state, + effective_base, + (ceiling_c2s_bytes, ceiling_s2c_bytes), + budget.as_ref(), + ); + } + reconcile_reservation(&mut lease, &c2s_state, &s2c_state); + previous = sample; + previous_log_c2s = sample.c2s_bytes; + previous_log_s2c = sample.s2c_written_bytes; + previous_sample_at = now; + } + } + } + }; + + let _ = client_writer.shutdown().await; + let _ = server_writer.shutdown().await; + let c2s_ops = counters.c2s_ops.load(Ordering::Relaxed); + let s2c_ops = counters.s2c_ops.load(Ordering::Relaxed); + let duration = epoch.elapsed(); + match relay_outcome { + AdaptiveRelayOutcome::Copy(Ok((c2s, s2c))) => { + debug!( + user = %user_owned, + c2s_bytes = c2s, + s2c_bytes = s2c, + c2s_msgs = c2s_ops, + s2c_msgs = s2c_ops, + duration_secs = duration.as_secs(), + "Relay finished" + ); + Ok(()) + } + AdaptiveRelayOutcome::Copy(Err(error)) if is_quota_io_error(&error) => { + warn!( + user = %user_owned, + c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed), + s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed), + c2s_msgs = c2s_ops, + s2c_msgs = s2c_ops, + duration_secs = duration.as_secs(), + "Data quota reached, closing relay" + ); + Err(ProxyError::DataQuotaExceeded { user: user_owned }) + } + AdaptiveRelayOutcome::Copy(Err(error)) => { + debug!( + user = %user_owned, + c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed), + s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed), + c2s_msgs = c2s_ops, + s2c_msgs = s2c_ops, + duration_secs = duration.as_secs(), + error = %error, + "Relay error" + ); + Err(error.into()) + } + AdaptiveRelayOutcome::ActivityTimeout => { + debug!( + user = %user_owned, + c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed), + s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed), + c2s_msgs = c2s_ops, + s2c_msgs = s2c_ops, + duration_secs = duration.as_secs(), + "Relay finished (activity timeout)" + ); + Ok(()) + } + AdaptiveRelayOutcome::UserDisabled => { + debug!( + user = %user_owned, + c2s_bytes = counters.c2s_bytes.load(Ordering::Relaxed), + s2c_bytes = counters.s2c_bytes.load(Ordering::Relaxed), + c2s_msgs = c2s_ops, + s2c_msgs = s2c_ops, + duration_secs = duration.as_secs(), + "Relay finished (user disabled)" + ); + Err(ProxyError::UserDisabled { user: user_owned }) + } + } +} + +fn initial_base_sizes( + ceiling_c2s: usize, + ceiling_s2c: usize, + max_connections: u32, + target_bytes: usize, +) -> (usize, usize) { + let configured_total = ceiling_c2s.saturating_add(ceiling_s2c); + let configured_worst_case = configured_total.saturating_mul(max_connections as usize); + if max_connections != 0 && configured_worst_case <= target_bytes { + return (ceiling_c2s, ceiling_s2c); + } + ( + DIRECT_BASE_C2S_BYTES.min(ceiling_c2s), + DIRECT_BASE_S2C_BYTES.min(ceiling_s2c), + ) +} + +fn current_sample(counters: &SharedCounters) -> RelaySignalSample { + RelaySignalSample { + c2s_bytes: counters.c2s_bytes.load(Ordering::Relaxed), + s2c_requested_bytes: counters.s2c_requested_bytes.load(Ordering::Relaxed), + s2c_written_bytes: counters.s2c_bytes.load(Ordering::Relaxed), + s2c_write_ops: counters.s2c_ops.load(Ordering::Relaxed), + s2c_partial_writes: counters.s2c_partial_writes.load(Ordering::Relaxed), + s2c_consecutive_pending_writes: counters + .s2c_consecutive_pending_writes + .load(Ordering::Relaxed), + } +} + +fn sample_delta(current: RelaySignalSample, previous: RelaySignalSample) -> RelaySignalSample { + RelaySignalSample { + c2s_bytes: current.c2s_bytes.saturating_sub(previous.c2s_bytes), + s2c_requested_bytes: current + .s2c_requested_bytes + .saturating_sub(previous.s2c_requested_bytes), + s2c_written_bytes: current + .s2c_written_bytes + .saturating_sub(previous.s2c_written_bytes), + s2c_write_ops: current.s2c_write_ops.saturating_sub(previous.s2c_write_ops), + s2c_partial_writes: current + .s2c_partial_writes + .saturating_sub(previous.s2c_partial_writes), + s2c_consecutive_pending_writes: current.s2c_consecutive_pending_writes, + } +} + +#[allow(clippy::too_many_arguments)] +fn apply_controller_transition( + transition: crate::proxy::adaptive_buffers::TierTransition, + controller: &mut SessionAdaptiveController, + lease: &mut DirectBufferLease, + c2s_state: &AdaptiveBufferState, + s2c_state: &AdaptiveBufferState, + base: (usize, usize), + ceilings: (usize, usize), + budget: &DirectBufferBudget, +) { + let sizes = direct_copy_buffers_for_tier_with_ceilings( + transition.to, + base.0, + base.1, + ceilings.0, + ceilings.1, + ); + if transition.to > transition.from { + if !lease.try_grow_to(sizes.0.saturating_add(sizes.1)) { + *controller = SessionAdaptiveController::new(transition.from); + return; + } + } else { + match transition.reason { + TierTransitionReason::QuietDemotion => budget.increment_quiet_demotion(), + TierTransitionReason::SustainedWritePressure => { + budget.increment_write_pressure_demotion(); + } + TierTransitionReason::SoftConfirmed | TierTransitionReason::HardPressure => {} + } + } + set_desired_sizes(c2s_state, s2c_state, sizes); + lease.set_tier(transition.to.as_u8() as usize); +} + +fn apply_global_pressure_demotion( + controller: &mut SessionAdaptiveController, + lease: &mut DirectBufferLease, + c2s_state: &AdaptiveBufferState, + s2c_state: &AdaptiveBufferState, + base: (usize, usize), + ceilings: (usize, usize), + budget: &DirectBufferBudget, +) { + let current = controller.tier(); + let target = current.demote(); + if target == current { + return; + } + *controller = SessionAdaptiveController::new(target); + let sizes = direct_copy_buffers_for_tier_with_ceilings( + target, + base.0, + base.1, + ceilings.0, + ceilings.1, + ); + set_desired_sizes(c2s_state, s2c_state, sizes); + lease.set_tier(target.as_u8() as usize); + budget.increment_global_pressure_demotion(); +} + +fn set_desired_sizes( + c2s_state: &AdaptiveBufferState, + s2c_state: &AdaptiveBufferState, + sizes: (usize, usize), +) { + c2s_state + .desired_bytes + .store(sizes.0.max(1), Ordering::Release); + s2c_state + .desired_bytes + .store(sizes.1.max(1), Ordering::Release); +} + +fn reconcile_reservation( + lease: &mut DirectBufferLease, + c2s_state: &AdaptiveBufferState, + s2c_state: &AdaptiveBufferState, +) { + // Promotion reserves the desired allocation before either reader grows. + // Demotion keeps the actual allocation covered until its buffered bytes drain. + let covered_c2s = c2s_state + .actual_bytes + .load(Ordering::Acquire) + .max(c2s_state.desired_bytes.load(Ordering::Acquire)); + let covered_s2c = s2c_state + .actual_bytes + .load(Ordering::Acquire) + .max(s2c_state.desired_bytes.load(Ordering::Acquire)); + lease.shrink_to(covered_c2s.saturating_add(covered_s2c)); +} diff --git a/src/proxy/relay/adaptive_copy/write_pressure.rs b/src/proxy/relay/adaptive_copy/write_pressure.rs new file mode 100644 index 0000000..52ddd50 --- /dev/null +++ b/src/proxy/relay/adaptive_copy/write_pressure.rs @@ -0,0 +1,72 @@ +use std::io; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::task::{Context, Poll}; + +use tokio::io::AsyncWrite; + +use super::super::io::SharedCounters; + +/// Direct-only writer wrapper that exposes bounded backpressure signals. +pub(super) struct WritePressureIo { + inner: W, + counters: Arc, +} + +impl WritePressureIo { + /// Wraps the client writer without changing its I/O or error contract. + pub(super) fn new(inner: W, counters: Arc) -> Self { + Self { inner, counters } + } +} + +impl AsyncWrite for WritePressureIo { + fn poll_write( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + buffer: &[u8], + ) -> Poll> { + let this = self.get_mut(); + if !buffer.is_empty() { + this.counters + .s2c_requested_bytes + .fetch_add(buffer.len() as u64, Ordering::Relaxed); + } + match Pin::new(&mut this.inner).poll_write(cx, buffer) { + Poll::Ready(Ok(written)) => { + this.counters + .s2c_consecutive_pending_writes + .store(0, Ordering::Relaxed); + if written < buffer.len() { + this.counters + .s2c_partial_writes + .fetch_add(1, Ordering::Relaxed); + } + Poll::Ready(Ok(written)) + } + Poll::Ready(Err(error)) => { + this.counters + .s2c_consecutive_pending_writes + .store(0, Ordering::Relaxed); + Poll::Ready(Err(error)) + } + Poll::Pending => { + let _ = this.counters.s2c_consecutive_pending_writes.fetch_update( + Ordering::Relaxed, + Ordering::Relaxed, + |current| Some(current.saturating_add(1)), + ); + Poll::Pending + } + } + } + + fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_flush(cx) + } + + fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + Pin::new(&mut self.get_mut().inner).poll_shutdown(cx) + } +} diff --git a/src/proxy/relay/io/counters.rs b/src/proxy/relay/io/counters.rs index 963b88d..3bad32d 100644 --- a/src/proxy/relay/io/counters.rs +++ b/src/proxy/relay/io/counters.rs @@ -1,4 +1,4 @@ -use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; use std::time::Duration; use tokio::time::Instant; @@ -20,6 +20,12 @@ pub(in crate::proxy::relay) struct SharedCounters { pub(in crate::proxy::relay) c2s_ops: AtomicU64, /// Number of poll_write completions (≈ S→C chunks) pub(in crate::proxy::relay) s2c_ops: AtomicU64, + /// Bytes presented to client writes, including retried pending writes. + pub(in crate::proxy::relay) s2c_requested_bytes: AtomicU64, + /// Successful client writes that consumed only part of the offered slice. + pub(in crate::proxy::relay) s2c_partial_writes: AtomicU64, + /// Consecutive pending client writes observed by the active copy loop. + pub(in crate::proxy::relay) s2c_consecutive_pending_writes: AtomicU32, /// Milliseconds since relay epoch of last I/O activity last_activity_ms: AtomicU64, } @@ -31,6 +37,9 @@ impl SharedCounters { s2c_bytes: AtomicU64::new(0), c2s_ops: AtomicU64::new(0), s2c_ops: AtomicU64::new(0), + s2c_requested_bytes: AtomicU64::new(0), + s2c_partial_writes: AtomicU64::new(0), + s2c_consecutive_pending_writes: AtomicU32::new(0), last_activity_ms: AtomicU64::new(0), } } diff --git a/src/proxy/shared_state.rs b/src/proxy/shared_state.rs index 6a47761..e036a41 100644 --- a/src/proxy/shared_state.rs +++ b/src/proxy/shared_state.rs @@ -10,6 +10,9 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio_util::sync::CancellationToken; use crate::proxy::handshake::{AuthProbeSaturationState, AuthProbeState}; +use crate::proxy::direct_buffer_budget::{ + DirectBufferBudget, fallback_direct_buffer_hard_limit, +}; use crate::proxy::middle_relay::{DesyncDedupRotationState, RelayIdleCandidateRegistry}; use crate::proxy::traffic_limiter::TrafficLimiter; @@ -69,6 +72,7 @@ pub(crate) struct ProxySharedState { pub(crate) handshake: HandshakeSharedState, pub(crate) middle_relay: MiddleRelaySharedState, pub(crate) traffic_limiter: Arc, + pub(crate) direct_buffer_budget: Arc, disabled_users: DashMap, active_user_sessions: DashMap<(String, u64), CancellationToken>, pub(crate) conntrack_pressure_active: AtomicBool, @@ -101,6 +105,15 @@ impl Drop for UserSessionGuard { impl ProxySharedState { pub(crate) fn new() -> Arc { + Self::new_with_direct_buffer_budget(DirectBufferBudget::new( + fallback_direct_buffer_hard_limit(), + )) + } + + /// Creates process state with the startup-resolved Direct buffer envelope. + pub(crate) fn new_with_direct_buffer_budget( + direct_buffer_budget: Arc, + ) -> Arc { Arc::new(Self { handshake: HandshakeSharedState { auth_probe: DashMap::new(), @@ -129,6 +142,7 @@ impl ProxySharedState { relay_idle_mark_seq: AtomicU64::new(0), }, traffic_limiter: TrafficLimiter::new(), + direct_buffer_budget, disabled_users: DashMap::new(), active_user_sessions: DashMap::new(), conntrack_pressure_active: AtomicBool::new(false), diff --git a/src/proxy/tests/adaptive_direct_budget_policy_tests.rs b/src/proxy/tests/adaptive_direct_budget_policy_tests.rs new file mode 100644 index 0000000..f7cee57 --- /dev/null +++ b/src/proxy/tests/adaptive_direct_budget_policy_tests.rs @@ -0,0 +1,58 @@ +use super::*; + +#[test] +fn configured_direct_sizes_are_strict_tier_ceilings() { + let base = (4 * 1024, 8 * 1024); + let ceilings = (64 * 1024, 256 * 1024); + assert_eq!( + direct_copy_buffers_for_tier_with_ceilings( + AdaptiveTier::Base, + base.0, + base.1, + ceilings.0, + ceilings.1, + ), + base + ); + assert_eq!( + direct_copy_buffers_for_tier_with_ceilings( + AdaptiveTier::Tier3, + base.0, + base.1, + ceilings.0, + ceilings.1, + ), + ceilings + ); +} + +#[test] +fn sustained_pending_pressure_demotes_after_transient_promotion() { + let mut controller = SessionAdaptiveController::new(AdaptiveTier::Tier1); + let pressure = RelaySignalSample { + c2s_bytes: 0, + s2c_requested_bytes: 1024, + s2c_written_bytes: 0, + s2c_write_ops: 0, + s2c_partial_writes: 0, + s2c_consecutive_pending_writes: 3, + }; + + let first = controller + .observe(pressure, 10.0) + .expect("transient pressure must retain the staged promotion"); + assert_eq!(first.reason, TierTransitionReason::HardPressure); + + let second = controller + .observe(pressure, 10.0) + .expect("bounded transient pressure may promote one additional tier"); + assert_eq!(second.reason, TierTransitionReason::HardPressure); + let sustained = controller + .observe(pressure, 10.0) + .expect("sustained pressure must release one tier"); + assert_eq!( + sustained.reason, + TierTransitionReason::SustainedWritePressure + ); + assert_eq!(sustained.to, AdaptiveTier::Tier2); +} diff --git a/src/proxy/tests/direct_buffer_budget_tests.rs b/src/proxy/tests/direct_buffer_budget_tests.rs new file mode 100644 index 0000000..796e7fc --- /dev/null +++ b/src/proxy/tests/direct_buffer_budget_tests.rs @@ -0,0 +1,38 @@ +use super::*; + +#[test] +fn lease_drop_releases_the_complete_reservation() { + let budget = DirectBufferBudget::new(16 * 1024); + { + let lease = budget + .try_reserve(12 * 1024, false) + .expect("minimum reservation must fit"); + assert_eq!(lease.reserved_bytes(), 12 * 1024); + assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024); + } + assert_eq!(budget.snapshot().reserved_bytes, 0); +} + +#[test] +fn absolute_ceiling_rejects_excess_minimum_reservations() { + let budget = DirectBufferBudget::new(16 * 1024); + let _lease = budget + .try_reserve(12 * 1024, true) + .expect("first minimum reservation must fit"); + assert!(budget.try_reserve(8 * 1024, true).is_none()); + assert_eq!(budget.snapshot().reserved_bytes, 12 * 1024); +} + +#[test] +fn growth_and_shrink_keep_accounting_balanced() { + let budget = DirectBufferBudget::new(32 * 1024); + let mut lease = budget + .try_reserve(12 * 1024, false) + .expect("base reservation must fit"); + assert!(lease.try_grow_to(24 * 1024)); + assert_eq!(budget.snapshot().reserved_bytes, 24 * 1024); + lease.shrink_to(16 * 1024); + assert_eq!(budget.snapshot().reserved_bytes, 16 * 1024); + drop(lease); + assert_eq!(budget.snapshot().reserved_bytes, 0); +}