mirror of
https://github.com/telemt/telemt.git
synced 2026-10-07 18:05:57 +03:00
Maestro: add in-process runtime generation reload
Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>
This commit is contained in:
@@ -80,7 +80,11 @@ impl Stats {
|
||||
return handle;
|
||||
}
|
||||
|
||||
let entry = self.user_stats.entry(user.to_string()).or_default();
|
||||
let quota = self.quota_store.user(user);
|
||||
let entry = self
|
||||
.user_stats
|
||||
.entry(user.to_string())
|
||||
.or_insert_with(|| Arc::new(UserStats::with_quota(quota)));
|
||||
if entry.last_seen_epoch_secs.load(Ordering::Relaxed) == 0 {
|
||||
self.touch_user_stats(entry.value().as_ref());
|
||||
}
|
||||
@@ -166,10 +170,7 @@ impl Stats {
|
||||
#[inline]
|
||||
pub(crate) fn quota_charge_post_write(&self, user_stats: &UserStats, bytes: u64) -> u64 {
|
||||
self.touch_user_stats(user_stats);
|
||||
user_stats
|
||||
.quota_used
|
||||
.fetch_add(bytes, Ordering::Relaxed)
|
||||
.saturating_add(bytes)
|
||||
user_stats.quota.charge(bytes)
|
||||
}
|
||||
|
||||
pub(super) fn maybe_cleanup_user_stats(&self) {
|
||||
|
||||
+26
-23
@@ -8,6 +8,7 @@ mod core_getters;
|
||||
mod helpers;
|
||||
mod me_counters;
|
||||
mod me_getters;
|
||||
mod quota_store;
|
||||
mod replay;
|
||||
pub mod telemetry;
|
||||
pub mod tls_fingerprints;
|
||||
@@ -20,6 +21,7 @@ use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering};
|
||||
use std::time::Instant;
|
||||
|
||||
pub(crate) use self::quota_store::QuotaStore;
|
||||
#[allow(unused_imports)]
|
||||
pub use self::replay::{ReplayChecker, ReplayStats};
|
||||
use self::telemetry::TelemetryPolicy;
|
||||
@@ -344,6 +346,7 @@ pub struct Stats {
|
||||
cached_epoch_secs: AtomicU64,
|
||||
tls_fingerprints: tls_fingerprints::TlsFingerprintCollector,
|
||||
user_stats: DashMap<String, Arc<UserStats>>,
|
||||
quota_store: Arc<QuotaStore>,
|
||||
user_stats_last_cleanup_epoch_secs: AtomicU64,
|
||||
start_time: parking_lot::RwLock<Option<Instant>>,
|
||||
}
|
||||
@@ -356,12 +359,7 @@ pub struct UserStats {
|
||||
pub octets_to_client: AtomicU64,
|
||||
pub msgs_from_client: AtomicU64,
|
||||
pub msgs_to_client: AtomicU64,
|
||||
/// Total bytes charged against per-user quota admission.
|
||||
///
|
||||
/// This counter is the single source of truth for quota enforcement and
|
||||
/// intentionally tracks attempted traffic, not guaranteed delivery.
|
||||
pub quota_used: AtomicU64,
|
||||
pub quota_last_reset_epoch_secs: AtomicU64,
|
||||
quota: Arc<quota_store::UserQuotaCounters>,
|
||||
pub last_seen_epoch_secs: AtomicU64,
|
||||
}
|
||||
|
||||
@@ -378,9 +376,21 @@ pub enum QuotaReserveError {
|
||||
}
|
||||
|
||||
impl UserStats {
|
||||
fn with_quota(quota: Arc<quota_store::UserQuotaCounters>) -> Self {
|
||||
Self {
|
||||
quota,
|
||||
..Self::default()
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn quota_used(&self) -> u64 {
|
||||
self.quota_used.load(Ordering::Relaxed)
|
||||
self.quota.used()
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn refund_quota(&self, bytes: u64) {
|
||||
self.quota.refund(bytes);
|
||||
}
|
||||
|
||||
/// Attempts one CAS reservation step against the quota counter.
|
||||
@@ -390,27 +400,20 @@ impl UserStats {
|
||||
/// with their own contention strategy.
|
||||
#[inline]
|
||||
pub fn quota_try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
|
||||
let current = self.quota_used.load(Ordering::Relaxed);
|
||||
if bytes > limit.saturating_sub(current) {
|
||||
return Err(QuotaReserveError::LimitExceeded);
|
||||
}
|
||||
|
||||
let next = current.saturating_add(bytes);
|
||||
match self.quota_used.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => Ok(next),
|
||||
Err(_) => Err(QuotaReserveError::Contended),
|
||||
}
|
||||
self.quota.try_reserve(bytes, limit)
|
||||
}
|
||||
}
|
||||
|
||||
impl Stats {
|
||||
pub fn new() -> Self {
|
||||
let stats = Self::default();
|
||||
Self::with_quota_store(Arc::new(QuotaStore::default()))
|
||||
}
|
||||
|
||||
pub(crate) fn with_quota_store(quota_store: Arc<QuotaStore>) -> Self {
|
||||
let stats = Self {
|
||||
quota_store,
|
||||
..Self::default()
|
||||
};
|
||||
stats.apply_telemetry_policy(TelemetryPolicy::default());
|
||||
stats.refresh_cached_epoch_secs();
|
||||
*stats.start_time.write() = Some(Instant::now());
|
||||
|
||||
@@ -0,0 +1,147 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
|
||||
use dashmap::DashMap;
|
||||
|
||||
use super::{QuotaReserveError, UserQuotaSnapshot};
|
||||
|
||||
/// Process-scoped per-user quota accounting shared by runtime generations.
|
||||
#[derive(Default)]
|
||||
pub struct QuotaStore {
|
||||
users: DashMap<String, Arc<UserQuotaCounters>>,
|
||||
}
|
||||
|
||||
/// Atomic quota state for one configured user.
|
||||
#[derive(Default)]
|
||||
pub(crate) struct UserQuotaCounters {
|
||||
used_bytes: AtomicU64,
|
||||
last_reset_epoch_secs: AtomicU64,
|
||||
}
|
||||
|
||||
impl QuotaStore {
|
||||
pub(crate) fn user(&self, user: &str) -> Arc<UserQuotaCounters> {
|
||||
if let Some(existing) = self.users.get(user) {
|
||||
return Arc::clone(existing.value());
|
||||
}
|
||||
Arc::clone(
|
||||
self.users
|
||||
.entry(user.to_string())
|
||||
.or_insert_with(|| Arc::new(UserQuotaCounters::default()))
|
||||
.value(),
|
||||
)
|
||||
}
|
||||
|
||||
pub(crate) fn used(&self, user: &str) -> u64 {
|
||||
self.users.get(user).map(|state| state.used()).unwrap_or(0)
|
||||
}
|
||||
|
||||
pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
|
||||
let state = self.user(user);
|
||||
state.used_bytes.store(used_bytes, Ordering::Relaxed);
|
||||
state
|
||||
.last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
}
|
||||
|
||||
pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot {
|
||||
let state = self.user(user);
|
||||
state.used_bytes.store(0, Ordering::Relaxed);
|
||||
state
|
||||
.last_reset_epoch_secs
|
||||
.store(now_epoch_secs, Ordering::Relaxed);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs: now_epoch_secs,
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
||||
let mut out = HashMap::new();
|
||||
for entry in self.users.iter() {
|
||||
let state = entry.value();
|
||||
let used_bytes = state.used();
|
||||
let last_reset_epoch_secs = state.last_reset_epoch_secs.load(Ordering::Relaxed);
|
||||
if used_bytes == 0 && last_reset_epoch_secs == 0 {
|
||||
continue;
|
||||
}
|
||||
out.insert(
|
||||
entry.key().clone(),
|
||||
UserQuotaSnapshot {
|
||||
used_bytes,
|
||||
last_reset_epoch_secs,
|
||||
},
|
||||
);
|
||||
}
|
||||
out
|
||||
}
|
||||
}
|
||||
|
||||
impl UserQuotaCounters {
|
||||
#[inline]
|
||||
pub(crate) fn used(&self) -> u64 {
|
||||
self.used_bytes.load(Ordering::Relaxed)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn charge(&self, bytes: u64) -> u64 {
|
||||
self.used_bytes
|
||||
.fetch_add(bytes, Ordering::Relaxed)
|
||||
.saturating_add(bytes)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn refund(&self, bytes: u64) {
|
||||
let mut current = self.used_bytes.load(Ordering::Relaxed);
|
||||
loop {
|
||||
let next = current.saturating_sub(bytes);
|
||||
match self.used_bytes.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => return,
|
||||
Err(observed) => current = observed,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub(crate) fn try_reserve(&self, bytes: u64, limit: u64) -> Result<u64, QuotaReserveError> {
|
||||
let current = self.used_bytes.load(Ordering::Relaxed);
|
||||
if bytes > limit.saturating_sub(current) {
|
||||
return Err(QuotaReserveError::LimitExceeded);
|
||||
}
|
||||
|
||||
let next = current.saturating_add(bytes);
|
||||
match self.used_bytes.compare_exchange_weak(
|
||||
current,
|
||||
next,
|
||||
Ordering::Relaxed,
|
||||
Ordering::Relaxed,
|
||||
) {
|
||||
Ok(_) => Ok(next),
|
||||
Err(_) => Err(QuotaReserveError::Contended),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::stats::Stats;
|
||||
|
||||
#[test]
|
||||
fn quota_counters_are_shared_across_stats_generations() {
|
||||
let store = Arc::new(QuotaStore::default());
|
||||
let first = Stats::with_quota_store(store.clone());
|
||||
store.user("alice").charge(512);
|
||||
assert_eq!(first.get_user_quota_used("alice"), 512);
|
||||
|
||||
let second = Stats::with_quota_store(store);
|
||||
assert_eq!(second.get_user_quota_used("alice"), 512);
|
||||
second.reset_user_quota("alice");
|
||||
assert_eq!(first.get_user_quota_used("alice"), 0);
|
||||
}
|
||||
}
|
||||
+5
-35
@@ -119,51 +119,21 @@ impl Stats {
|
||||
}
|
||||
|
||||
pub fn get_user_quota_used(&self, user: &str) -> u64 {
|
||||
self.user_stats
|
||||
.get(user)
|
||||
.map(|s| s.quota_used.load(Ordering::Relaxed))
|
||||
.unwrap_or(0)
|
||||
self.quota_store.used(user)
|
||||
}
|
||||
|
||||
pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
|
||||
let stats = self.get_or_create_user_stats_handle(user);
|
||||
stats.quota_used.store(used_bytes, Ordering::Relaxed);
|
||||
stats
|
||||
.quota_last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
self.quota_store
|
||||
.load(user, used_bytes, last_reset_epoch_secs);
|
||||
}
|
||||
|
||||
pub fn reset_user_quota(&self, user: &str) -> UserQuotaSnapshot {
|
||||
let stats = self.get_or_create_user_stats_handle(user);
|
||||
let last_reset_epoch_secs = Self::now_epoch_secs();
|
||||
stats.quota_used.store(0, Ordering::Relaxed);
|
||||
stats
|
||||
.quota_last_reset_epoch_secs
|
||||
.store(last_reset_epoch_secs, Ordering::Relaxed);
|
||||
UserQuotaSnapshot {
|
||||
used_bytes: 0,
|
||||
last_reset_epoch_secs,
|
||||
}
|
||||
self.quota_store.reset(user, last_reset_epoch_secs)
|
||||
}
|
||||
|
||||
pub fn user_quota_snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
|
||||
let mut out = HashMap::new();
|
||||
for entry in self.user_stats.iter() {
|
||||
let stats = entry.value();
|
||||
let used_bytes = stats.quota_used.load(Ordering::Relaxed);
|
||||
let last_reset_epoch_secs = stats.quota_last_reset_epoch_secs.load(Ordering::Relaxed);
|
||||
if used_bytes == 0 && last_reset_epoch_secs == 0 {
|
||||
continue;
|
||||
}
|
||||
out.insert(
|
||||
entry.key().clone(),
|
||||
UserQuotaSnapshot {
|
||||
used_bytes,
|
||||
last_reset_epoch_secs,
|
||||
},
|
||||
);
|
||||
}
|
||||
out
|
||||
self.quota_store.snapshot()
|
||||
}
|
||||
|
||||
pub fn get_handshake_timeouts(&self) -> u64 {
|
||||
|
||||
Reference in New Issue
Block a user