Bound Direct relay buffers with an adaptive global memory envelope

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
Alexey
2026-07-11 20:56:54 +03:00
parent d4c4980e5a
commit 96425f15c8
18 changed files with 1586 additions and 48 deletions
+6
View File
@@ -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
}
+20
View File
@@ -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(),
+1
View File
@@ -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",
@@ -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);
+9 -2
View File
@@ -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(),
+22 -1
View File
@@ -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,
+68
View File
@@ -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"
+137 -41
View File
@@ -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<TierTransition> {
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<TierTransition> {
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<TierTransition> {
fn demote(
&mut self,
reason: TierTransitionReason,
hard_cooldown: Duration,
) -> Option<TierTransition> {
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<String, UserAdaptiveProfile> {
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);
}
+549
View File
@@ -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<u64>,
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<Self> {
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<u64> {
self.pressure_tx.subscribe()
}
/// Reserves bytes against either the adaptive target or the absolute ceiling.
pub(crate) fn try_reserve(
self: &Arc<Self>,
bytes: usize,
allow_above_target: bool,
) -> Option<DirectBufferLease> {
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<DirectBufferBudget>,
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<DirectBufferBudget>,
stats: Arc<Stats>,
shared: Arc<ProxySharedState>,
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<u8> {
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<u8> {
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<u64> {
let raw = tokio::fs::read_to_string(path).await.ok()?;
let raw = raw.trim();
if raw == "max" {
return None;
}
let value = raw.parse::<u64>().ok()?;
(value < (1u64 << 60)).then_some(value)
}
#[cfg(target_os = "linux")]
async fn read_u64_file(path: &str) -> Option<u64> {
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::<u64>().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;
+3 -3
View File
@@ -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 {
+2
View File
@@ -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;
+4
View File
@@ -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<CR, CW, SR, SW>(
client_reader: CR,
client_writer: CW,
+532
View File
@@ -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<Self> {
Arc::new(Self {
desired_bytes: AtomicUsize::new(bytes.max(1)),
actual_bytes: AtomicUsize::new(bytes.max(1)),
})
}
}
struct AdaptiveBufReader<R> {
inner: R,
buffer: Box<[u8]>,
pos: usize,
cap: usize,
state: Arc<AdaptiveBufferState>,
}
impl<R> AdaptiveBufReader<R> {
fn new(inner: R, state: Arc<AdaptiveBufferState>) -> 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<R: AsyncRead + Unpin> AsyncRead for AdaptiveBufReader<R> {
fn poll_read(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
output: &mut ReadBuf<'_>,
) -> Poll<io::Result<()>> {
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<R: AsyncRead + Unpin> AsyncBufRead for AdaptiveBufReader<R> {
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<&[u8]>> {
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<CR, CW, SR, SW>(
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<Stats>,
quota_limit: Option<u64>,
traffic_lease: Option<Arc<TrafficLease>>,
activity_timeout: Duration,
session_cancel: CancellationToken,
budget: Arc<DirectBufferBudget>,
) -> 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(&quota_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(&quota_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::<u64, io::Error>(copied)
};
let s2c = async {
let copied = copy_buf(&mut s2c_reader, &mut client_writer).await?;
client_writer.shutdown().await?;
Ok::<u64, io::Error>(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));
}
@@ -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<W> {
inner: W,
counters: Arc<SharedCounters>,
}
impl<W> WritePressureIo<W> {
/// Wraps the client writer without changing its I/O or error contract.
pub(super) fn new(inner: W, counters: Arc<SharedCounters>) -> Self {
Self { inner, counters }
}
}
impl<W: AsyncWrite + Unpin> AsyncWrite for WritePressureIo<W> {
fn poll_write(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
buffer: &[u8],
) -> Poll<io::Result<usize>> {
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<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_flush(cx)
}
fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
}
}
+10 -1
View File
@@ -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),
}
}
+14
View File
@@ -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<TrafficLimiter>,
pub(crate) direct_buffer_budget: Arc<DirectBufferBudget>,
disabled_users: DashMap<String, ()>,
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> {
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<DirectBufferBudget>,
) -> Arc<Self> {
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),
@@ -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);
}
@@ -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);
}