astelm/telemt
1
0
Fork 0
mirror of https://github.com/telemt/telemt.git synced 2026-11-02 05:42:00 +03:00
Code Issues Packages Projects Releases Wiki Activity

Hardswap Invariants in tests + Quota fixes

Browse Source
This commit is contained in:
Alexey
2026-09-20 00:28:52 +03:00
parent 89dacbd17e
commit d706b3f3ba
66 changed files with 1821 additions and 505 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/proxy/relay/adaptive_copy.rs
+4 -1
View File
@@ -19,7 +19,7 @@ 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 crate::stats::{Stats, UserQuotaHandle};
use super::WATCHDOG_INTERVAL;
use super::io::{SharedCounters, StatsIo, is_quota_io_error};
@@ -141,6 +141,7 @@ pub(crate) async fn relay_direct_adaptive<CR, CW, SR, SW>(
max_connections: u32,
user: &str,
stats: Arc<Stats>,
quota_handle: UserQuotaHandle,
quota_limit: Option<u64>,
traffic_lease: Option<Arc<TrafficLease>>,
activity_timeout: Duration,
@@ -200,6 +201,7 @@ where
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
quota_handle.clone(),
traffic_lease.clone(),
quota_limit,
Arc::clone(&quota_exceeded),
@@ -210,6 +212,7 @@ where
Arc::clone(&counters),
Arc::clone(&stats),
user_owned.clone(),
quota_handle,
traffic_lease,
quota_limit,
Arc::clone(&quota_exceeded),
src/proxy/relay/io.rs
+14 -9
View File
@@ -1,5 +1,5 @@
use crate::proxy::traffic_limiter::{RateDirection, TrafficLease, next_refill_delay};
use crate::stats::{Stats, UserStats};
use crate::stats::{Stats, UserQuotaHandle, UserStats};
use std::io;
use std::pin::Pin;
use std::sync::Arc;
@@ -40,6 +40,7 @@ pub(super) struct StatsIo<S> {
stats: Arc<Stats>,
user: String,
user_stats: Arc<UserStats>,
quota_handle: UserQuotaHandle,
traffic_lease: Option<Arc<TrafficLease>>,
c2s_rate_debt_bytes: u64,
c2s_wait: RateWaitState,
@@ -71,11 +72,13 @@ impl<S> StatsIo<S> {
quota_exceeded: Arc<AtomicBool>,
epoch: Instant,
) -> Self {
let quota_handle = stats.current_user_quota_handle(&user);
Self::new_with_traffic_lease(
inner,
counters,
stats,
user,
quota_handle,
None,
quota_limit,
quota_exceeded,
@@ -88,6 +91,7 @@ impl<S> StatsIo<S> {
counters: Arc<SharedCounters>,
stats: Arc<Stats>,
user: String,
quota_handle: UserQuotaHandle,
traffic_lease: Option<Arc<TrafficLease>>,
quota_limit: Option<u64>,
quota_exceeded: Arc<AtomicBool>,
@@ -102,6 +106,7 @@ impl<S> StatsIo<S> {
stats,
user,
user_stats,
quota_handle,
traffic_lease,
c2s_rate_debt_bytes: 0,
c2s_wait: RateWaitState::default(),
@@ -213,7 +218,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
let mut quota_reservation = None;
let mut read_limit = buf.remaining();
if let Some(limit) = this.quota_limit {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -230,7 +235,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
let mut reserve_rounds = 0usize;
while quota_reservation.is_none() {
for _ in 0..QUOTA_RESERVE_SPIN_RETRIES {
match this.user_stats.quota_reserve(desired, limit) {
match this.quota_handle.try_reserve(desired, limit) {
Ok(reservation) => {
quota_reservation = Some(reservation);
break;
@@ -305,7 +310,7 @@ impl<S: AsyncRead + Unpin> AsyncRead for StatsIo<S> {
}
}
if let Some(limit) = this.quota_limit
&& this.user_stats.quota_used() >= limit
&& this.quota_handle.used() >= limit
{
this.quota_exceeded.store(true, Ordering::Release);
}
@@ -401,7 +406,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
if !write_buf.is_empty() {
let mut reserve_rounds = 0usize;
while quota_reservation.is_none() {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -412,7 +417,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
let desired = remaining.min(write_buf.len() as u64);
let mut saw_contention = false;
for _ in 0..QUOTA_RESERVE_SPIN_RETRIES {
match this.user_stats.quota_reserve(desired, limit) {
match this.quota_handle.try_reserve(desired, limit) {
Ok(reservation) => {
quota_reservation = Some(reservation);
write_buf = &write_buf[..desired as usize];
@@ -442,7 +447,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
}
}
} else {
let used_before = this.user_stats.quota_used();
let used_before = this.quota_handle.used();
let remaining = limit.saturating_sub(used_before);
if remaining == 0 {
this.quota_exceeded.store(true, Ordering::Release);
@@ -481,7 +486,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
if let (Some(limit), Some(remaining)) = (this.quota_limit, remaining_before) {
if should_immediate_quota_check(remaining, n_to_charge) {
this.quota_bytes_since_check = 0;
if this.user_stats.quota_used() >= limit {
if this.quota_handle.used() >= limit {
this.quota_exceeded.store(true, Ordering::Release);
}
} else {
@@ -490,7 +495,7 @@ impl<S: AsyncWrite + Unpin> AsyncWrite for StatsIo<S> {
let interval = quota_adaptive_interval_bytes(remaining);
if this.quota_bytes_since_check >= interval {
this.quota_bytes_since_check = 0;
if this.user_stats.quota_used() >= limit {
if this.quota_handle.used() >= limit {
this.quota_exceeded.store(true, Ordering::Release);
}
}