diff --git a/src/config/tests/load_memory_envelope_tests.rs b/src/config/tests/load_memory_envelope_tests.rs index 3592d27..85e274c 100644 --- a/src/config/tests/load_memory_envelope_tests.rs +++ b/src/config/tests/load_memory_envelope_tests.rs @@ -145,8 +145,7 @@ me_writer_byte_budget_bytes = 268451840 "#, ); - let err = - ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail"); + let err = ProxyConfig::load(&path).expect_err("writer byte budget above hard cap must fail"); let msg = err.to_string(); assert!( msg.contains("general.me_writer_byte_budget_bytes must be within [33570816, 268435456]"), @@ -165,8 +164,7 @@ direct_relay_buffer_budget_max_bytes = 16777217 "#, ); - let err = ProxyConfig::load(&path) - .expect_err("unaligned direct relay buffer budget must fail"); + 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" @@ -184,8 +182,8 @@ direct_relay_buffer_budget_max_bytes = 2147487744 "#, ); - let err = ProxyConfig::load(&path) - .expect_err("direct relay buffer budget above hard cap must fail"); + 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]" )); diff --git a/src/config/types.rs b/src/config/types.rs index 0994a37..4607732 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -1126,8 +1126,7 @@ 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(), + 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 5143f9c..e01670b 100644 --- a/src/maestro/mod.rs +++ b/src/maestro/mod.rs @@ -33,11 +33,10 @@ use crate::conntrack_control; 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, + DirectBufferBudget, resolve_direct_buffer_hard_limit, spawn_direct_buffer_budget_controller, }; +use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController}; use crate::proxy::shared_state::ProxySharedState; use crate::startup::{ COMPONENT_API_BOOTSTRAP, COMPONENT_CONFIG_LOAD, COMPONENT_DC_CONNECTIVITY_PING, @@ -477,10 +476,8 @@ async fn run_telemt_core( config.network.dns_overrides.len() ); } - let direct_buffer_hard_limit = resolve_direct_buffer_hard_limit( - config.general.direct_relay_buffer_budget_max_bytes, - ) - .await; + 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, diff --git a/src/metrics.rs b/src/metrics.rs index acbe502..f2816e2 100644 --- a/src/metrics.rs +++ b/src/metrics.rs @@ -611,10 +611,7 @@ async fn render_metrics( 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" - ); + 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), @@ -2517,10 +2514,7 @@ async fn render_metrics( out, "# HELP telemt_me_writer_byte_budget_limit_bytes Configured resident-memory budget per ME writer" ); - let _ = writeln!( - out, - "# TYPE telemt_me_writer_byte_budget_limit_bytes gauge" - ); + let _ = writeln!(out, "# TYPE telemt_me_writer_byte_budget_limit_bytes gauge"); let _ = writeln!( out, "telemt_me_writer_byte_budget_limit_bytes {}", diff --git a/src/proxy/direct_buffer_budget.rs b/src/proxy/direct_buffer_budget.rs index 44ed537..eea3c0d 100644 --- a/src/proxy/direct_buffer_budget.rs +++ b/src/proxy/direct_buffer_budget.rs @@ -170,11 +170,9 @@ impl DirectBufferBudget { } 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 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 @@ -196,8 +194,7 @@ impl DirectBufferBudget { /// 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); + self.minimum_fallback_total.fetch_add(1, Ordering::Relaxed); } /// Records a session rejected because the absolute ceiling was exhausted. @@ -369,7 +366,9 @@ pub(crate) fn spawn_direct_buffer_budget_controller( budget.update_system_sample(sample); let snapshot = budget.snapshot(); - let denied_delta = snapshot.promotion_denied_total.saturating_sub(previous_denied); + let denied_delta = snapshot + .promotion_denied_total + .saturating_sub(previous_denied); previous_denied = snapshot.promotion_denied_total; let fallback_delta = snapshot .minimum_fallback_total @@ -404,9 +403,7 @@ pub(crate) fn spawn_direct_buffer_budget_controller( pool_snapshot.allocated, pool_snapshot.allocated.saturating_sub(pool_snapshot.pooled), ); - stats.set_buffer_pool_replaced_nonstandard_total( - pool_snapshot.replaced_nonstandard, - ); + stats.set_buffer_pool_replaced_nonstandard_total(pool_snapshot.replaced_nonstandard); let headroom_target = if sample.total_bytes == 0 { snapshot.hard_limit_bytes @@ -454,11 +451,8 @@ fn connection_fill_pct(stats: &Stats, max_connections: u32) -> Option { return None; } Some( - ((stats - .get_current_connections_total() - .saturating_mul(100)) - / u64::from(max_connections)) - .min(100) as u8, + ((stats.get_current_connections_total().saturating_mul(100)) / u64::from(max_connections)) + .min(100) as u8, ) } @@ -485,8 +479,7 @@ async fn read_system_memory_sample() -> SystemMemorySample { 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_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); @@ -495,9 +488,9 @@ async fn read_system_memory_sample() -> SystemMemorySample { (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 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), diff --git a/src/proxy/direct_relay.rs b/src/proxy/direct_relay.rs index d3a9d73..54a55b6 100644 --- a/src/proxy/direct_relay.rs +++ b/src/proxy/direct_relay.rs @@ -346,21 +346,21 @@ where Duration::from_secs(1800) }; 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(), - traffic_lease, - relay_activity_timeout, - session_cancel.clone(), - Arc::clone(&shared.direct_buffer_budget), - ); + 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(), + traffic_lease, + relay_activity_timeout, + session_cancel.clone(), + Arc::clone(&shared.direct_buffer_budget), + ); tokio::pin!(relay_result); let relay_result = loop { if let Some(cutover) = diff --git a/src/proxy/relay.rs b/src/proxy/relay.rs index 7b6d6c9..6695026 100644 --- a/src/proxy/relay.rs +++ b/src/proxy/relay.rs @@ -84,8 +84,8 @@ fn watchdog_delta(current: u64, previous: u64) -> u64 { current.saturating_sub(previous) } -mod io; mod adaptive_copy; +mod io; pub(crate) use self::adaptive_copy::relay_direct_adaptive; diff --git a/src/proxy/relay/adaptive_copy.rs b/src/proxy/relay/adaptive_copy.rs index c267343..cae723c 100644 --- a/src/proxy/relay/adaptive_copy.rs +++ b/src/proxy/relay/adaptive_copy.rs @@ -488,13 +488,8 @@ fn apply_global_pressure_demotion( return; } *controller = SessionAdaptiveController::new(target); - let sizes = direct_copy_buffers_for_tier_with_ceilings( - target, - base.0, - base.1, - ceilings.0, - ceilings.1, - ); + 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(); diff --git a/src/proxy/shared_state.rs b/src/proxy/shared_state.rs index e036a41..de897dd 100644 --- a/src/proxy/shared_state.rs +++ b/src/proxy/shared_state.rs @@ -9,10 +9,8 @@ use dashmap::DashMap; use tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc}; use tokio_util::sync::CancellationToken; +use crate::proxy::direct_buffer_budget::{DirectBufferBudget, fallback_direct_buffer_hard_limit}; 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; diff --git a/src/stats/me_getters.rs b/src/stats/me_getters.rs index 34a4e08..ce04423 100644 --- a/src/stats/me_getters.rs +++ b/src/stats/me_getters.rs @@ -292,7 +292,8 @@ impl Stats { } /// Returns the count of blocking writer byte-budget waits. pub fn get_me_writer_byte_budget_wait_total(&self) -> u64 { - self.me_writer_byte_budget_wait_total.load(Ordering::Relaxed) + self.me_writer_byte_budget_wait_total + .load(Ordering::Relaxed) } /// Returns the count of writer byte-budget wait timeouts. pub fn get_me_writer_byte_budget_timeout_total(&self) -> u64 { diff --git a/src/stats/writer_counters.rs b/src/stats/writer_counters.rs index 02c74ec..4e80479 100644 --- a/src/stats/writer_counters.rs +++ b/src/stats/writer_counters.rs @@ -114,11 +114,11 @@ impl Stats { ); } pub(crate) fn release_me_writer_byte_budget_inflight_bytes(&self, bytes: u64) { - let _ = self.me_writer_byte_budget_inflight_bytes_gauge.fetch_update( - Ordering::Relaxed, - Ordering::Relaxed, - |current| Some(current.saturating_sub(bytes)), - ); + let _ = self + .me_writer_byte_budget_inflight_bytes_gauge + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + Some(current.saturating_sub(bytes)) + }); } pub(crate) fn increment_me_writer_byte_budget_wait_total(&self) { if self.telemetry_me_allows_normal() { diff --git a/src/transport/middle_proxy/codec.rs b/src/transport/middle_proxy/codec.rs index 0677def..edb5d93 100644 --- a/src/transport/middle_proxy/codec.rs +++ b/src/transport/middle_proxy/codec.rs @@ -400,7 +400,10 @@ mod tests { .expect("writer byte permits must be available"); let mut writer_permit = WriterBytePermit::new(permit, 32 * 1024, stats.clone()); - assert_eq!(stats.get_me_writer_byte_budget_queued_bytes_gauge(), 32 * 1024); + assert_eq!( + stats.get_me_writer_byte_budget_queued_bytes_gauge(), + 32 * 1024 + ); assert_eq!(stats.get_me_writer_byte_budget_inflight_bytes_gauge(), 0); writer_permit.mark_inflight(); diff --git a/src/transport/middle_proxy/registry/writer.rs b/src/transport/middle_proxy/registry/writer.rs index ea0336f..a3d375f 100644 --- a/src/transport/middle_proxy/registry/writer.rs +++ b/src/transport/middle_proxy/registry/writer.rs @@ -77,13 +77,9 @@ impl ConnRegistry { .writer_idle_since_epoch_secs .entry(writer_id) .or_insert_with(Self::now_epoch_secs); - self.writers.map.insert( - writer_id, - super::WriterRoute { - tx, - byte_budget, - }, - ); + self.writers + .map + .insert(writer_id, super::WriterRoute { tx, byte_budget }); } /// Unregister connection, returning associated writer_id if any. diff --git a/src/transport/middle_proxy/send.rs b/src/transport/middle_proxy/send.rs index 2407949..c74aa74 100644 --- a/src/transport/middle_proxy/send.rs +++ b/src/transport/middle_proxy/send.rs @@ -6,8 +6,8 @@ use std::sync::Arc; use std::sync::atomic::Ordering; use std::time::{Duration, Instant}; -use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; use tokio::sync::mpsc::error::TrySendError; +use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, mpsc}; use tracing::{debug, warn}; use super::MePool; @@ -74,16 +74,15 @@ async fn reserve_writer_command_slot( ) -> std::result::Result, WriterCommandReserveError> { let reserve = tx.clone().reserve_owned(); match deadline { - Some(deadline) => match tokio::time::timeout( - deadline.saturating_duration_since(Instant::now()), - reserve, - ) - .await - { - Ok(Ok(permit)) => Ok(permit), - Ok(Err(_)) => Err(WriterCommandReserveError::Closed), - Err(_) => Err(WriterCommandReserveError::TimedOut), - }, + Some(deadline) => { + match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), reserve) + .await + { + Ok(Ok(permit)) => Ok(permit), + Ok(Err(_)) => Err(WriterCommandReserveError::Closed), + Err(_) => Err(WriterCommandReserveError::TimedOut), + } + } None => reserve.await.map_err(|_| WriterCommandReserveError::Closed), } } @@ -102,7 +101,10 @@ fn writer_resident_permits( let permits = resident_bytes.div_ceil(ME_WRITER_BYTE_PERMIT_UNIT_BYTES); let permits = u32::try_from(permits).ok()?; let reserved_bytes = (permits as usize).checked_mul(ME_WRITER_BYTE_PERMIT_UNIT_BYTES)?; - Some((permits.max(1), reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES))) + Some(( + permits.max(1), + reserved_bytes.max(ME_WRITER_BYTE_PERMIT_UNIT_BYTES), + )) } fn proxy_req_resident_permits( @@ -146,23 +148,18 @@ async fn reserve_writer_bytes( let acquire = byte_budget.clone().acquire_many_owned(permits); match deadline { - Some(deadline) => match tokio::time::timeout( - deadline.saturating_duration_since(Instant::now()), - acquire, - ) - .await - { - Ok(Ok(permit)) => Ok(WriterBytePermit::new( - permit, - reserved_bytes, - stats.clone(), - )), - Ok(Err(_)) => Err(WriterByteReserveError::Closed), - Err(_) => { - stats.increment_me_writer_byte_budget_timeout_total(); - Err(WriterByteReserveError::TimedOut) + Some(deadline) => { + match tokio::time::timeout(deadline.saturating_duration_since(Instant::now()), acquire) + .await + { + Ok(Ok(permit)) => Ok(WriterBytePermit::new(permit, reserved_bytes, stats.clone())), + Ok(Err(_)) => Err(WriterByteReserveError::Closed), + Err(_) => { + stats.increment_me_writer_byte_budget_timeout_total(); + Err(WriterByteReserveError::TimedOut) + } } - }, + } None => acquire .await .map(|permit| WriterBytePermit::new(permit, reserved_bytes, stats.clone())) @@ -189,27 +186,21 @@ impl MePool { .len() .checked_add(LEGACY_PROXY_REQ_SOURCE_CAPACITY_OVERHEAD_BYTES) else { - self.stats - .increment_me_writer_byte_budget_oversize_total(); + self.stats.increment_me_writer_byte_budget_oversize_total(); return Err(ProxyError::Proxy( "ME writer payload residency calculation overflow".into(), )); }; - let Some((writer_byte_permits, writer_reserved_bytes)) = proxy_req_resident_permits( - source_capacity, - data.len(), - tag, - proto_flags, - ) else { - self.stats - .increment_me_writer_byte_budget_oversize_total(); + let Some((writer_byte_permits, writer_reserved_bytes)) = + proxy_req_resident_permits(source_capacity, data.len(), tag, proto_flags) + else { + self.stats.increment_me_writer_byte_budget_oversize_total(); return Err(ProxyError::Proxy( "ME writer payload residency calculation overflow".into(), )); }; if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { - self.stats - .increment_me_writer_byte_budget_oversize_total(); + self.stats.increment_me_writer_byte_budget_oversize_total(); return Err(ProxyError::Proxy( "ME writer payload exceeds configured byte budget".into(), )); @@ -254,9 +245,8 @@ impl MePool { loop { if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { - let deadline = writer_send_deadline( - self.route_runtime.me_route_blocking_send_timeout, - ); + let deadline = + writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); let writer_permit = match reserve_writer_bytes( ¤t.byte_budget, writer_byte_permits, @@ -275,7 +265,10 @@ impl MePool { )); } Err(WriterByteReserveError::Closed) => { - warn!(writer_id = current.writer_id, "ME writer byte budget closed"); + warn!( + writer_id = current.writer_id, + "ME writer byte budget closed" + ); self.remove_writer_and_close_clients(current.writer_id) .await; continue; @@ -293,12 +286,7 @@ impl MePool { return Ok(()); } Err(TrySendError::Full(cmd)) => { - match reserve_writer_command_slot( - ¤t.tx, - deadline, - ) - .await - { + match reserve_writer_command_slot(¤t.tx, deadline).await { Ok(permit) => { permit.send(cmd); self.note_hybrid_route_success(); @@ -716,9 +704,7 @@ impl MePool { } self.stats .increment_me_writer_pick_blocking_fallback_total(); - let deadline = writer_send_deadline( - self.route_runtime.me_route_blocking_send_timeout, - ); + let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); let writer_permit = match reserve_writer_bytes( &w.byte_budget, writer_byte_permits, @@ -801,24 +787,20 @@ impl MePool { tag.as_ref().map(|tag| tag.as_slice()), proto_flags, ) else { - self.stats - .increment_me_writer_byte_budget_oversize_total(); + self.stats.increment_me_writer_byte_budget_oversize_total(); return Err(ProxyError::Proxy( "ME writer payload residency calculation overflow".into(), )); }; if writer_byte_permits as usize > self.writer_lifecycle.writer_byte_budget_permits { - self.stats - .increment_me_writer_byte_budget_oversize_total(); + self.stats.increment_me_writer_byte_budget_oversize_total(); return Err(ProxyError::Proxy( "ME writer payload exceeds configured byte budget".into(), )); } if let Some((current, current_meta)) = self.registry.get_writer_with_meta(conn_id).await { - let deadline = writer_send_deadline( - self.route_runtime.me_route_blocking_send_timeout, - ); + let deadline = writer_send_deadline(self.route_runtime.me_route_blocking_send_timeout); let writer_permit = match reserve_writer_bytes( ¤t.byte_budget, writer_byte_permits, @@ -837,7 +819,10 @@ impl MePool { )); } Err(WriterByteReserveError::Closed) => { - warn!(writer_id = current.writer_id, "ME writer byte budget closed"); + warn!( + writer_id = current.writer_id, + "ME writer byte budget closed" + ); self.remove_writer_and_close_clients(current.writer_id) .await; return self @@ -870,12 +855,7 @@ impl MePool { return Ok(()); } Err(TrySendError::Full(cmd)) => { - match reserve_writer_command_slot( - ¤t.tx, - deadline, - ) - .await - { + match reserve_writer_command_slot(¤t.tx, deadline).await { Ok(permit) => { permit.send(cmd); self.note_hybrid_route_success();