This commit is contained in:
Alexey
2026-10-19 00:15:27 +03:00
parent a99b548aa1
commit 6fdc2a5ba4
43 changed files with 2764 additions and 159 deletions
+3
View File
@@ -467,3 +467,6 @@ mod decoy_dns_tests;
// Exact host groups retain the source shape and process-deferred limit checks.
#[cfg(test)]
mod host_list_tests;
#[cfg(test)]
mod quota_period_tests;
+19
View File
@@ -0,0 +1,19 @@
use super::*;
#[tokio::test]
async fn quota_periods_do_not_enable_access_mutation_through_generic_config_patch() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
let body = "[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n";
std::fs::write(&path, body).unwrap();
let error = apply_patch_to_path(
&path,
&serde_json::json!({"access":{"user_quota_periods":{"alice":"daily"}}}),
None,
)
.await
.err()
.unwrap();
assert_eq!(error.code, "access_not_editable");
assert_eq!(std::fs::read_to_string(path).unwrap(), body);
}
+3
View File
@@ -29,6 +29,8 @@ pub(super) enum AccessSection {
UserMaxTcpConns,
UserExpirations,
UserDataQuota,
/// Per-user recurring UTC calendar settings.
UserQuotaPeriods,
UserRateLimits,
UserMaxUniqueIps,
}
@@ -42,6 +44,7 @@ impl AccessSection {
Self::UserMaxTcpConns => "access.user_max_tcp_conns",
Self::UserExpirations => "access.user_expirations",
Self::UserDataQuota => "access.user_data_quota",
Self::UserQuotaPeriods => "access.user_quota_periods",
Self::UserRateLimits => "access.user_rate_limits",
Self::UserMaxUniqueIps => "access.user_max_unique_ips",
}
+5
View File
@@ -231,6 +231,10 @@ pub(super) fn render_access_section(
.collect();
serialize_table_body(&rows)?
}
AccessSection::UserQuotaPeriods => {
let rows: BTreeMap<_, _> = cfg.access.user_quota_periods.iter().collect();
serialize_table_body(&rows)?
}
AccessSection::UserRateLimits => {
let rows: BTreeMap<String, RateLimitBps> = cfg
.access
@@ -269,6 +273,7 @@ fn access_section_is_empty(cfg: &ProxyConfig, section: AccessSection) -> bool {
AccessSection::UserMaxTcpConns => cfg.access.user_max_tcp_conns.is_empty(),
AccessSection::UserExpirations => cfg.access.user_expirations.is_empty(),
AccessSection::UserDataQuota => cfg.access.user_data_quota.is_empty(),
AccessSection::UserQuotaPeriods => cfg.access.user_quota_periods.is_empty(),
AccessSection::UserRateLimits => cfg.access.user_rate_limits.is_empty(),
AccessSection::UserMaxUniqueIps => cfg.access.user_max_unique_ips.is_empty(),
}
+7 -1
View File
@@ -200,7 +200,13 @@ pub(super) async fn handle(
("GET", "/v1/stats/users/quota") => {
let revision = current_revision(&shared.config_path).await?;
let disk_cfg = load_config_from_disk(&shared.config_path).await?;
let data = build_user_quota_list(&disk_cfg, shared.stats.as_ref());
let _quota_guard = shared.quota_state.lock().await;
let mut data = build_user_quota_list(&disk_cfg, shared.stats.as_ref());
for entry in &mut data.users {
entry.schedule = shared
.quota_state
.schedule_status(&entry.username, entry.quota_period);
}
Ok(success_response(StatusCode::OK, data, revision))
}
+12
View File
@@ -23,6 +23,8 @@ pub(in crate::api) struct UserInfo {
pub(in crate::api) max_tcp_conns: Option<usize>,
pub(in crate::api) expiration_rfc3339: Option<String>,
pub(in crate::api) data_quota_bytes: Option<u64>,
/// Desired UTC period; runtime activation is reported by the quota endpoint.
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
pub(in crate::api) rate_limit_up_bps: Option<u64>,
pub(in crate::api) rate_limit_down_bps: Option<u64>,
pub(in crate::api) max_unique_ips: Option<usize>,
@@ -69,6 +71,11 @@ pub(in crate::api) struct UserQuotaListData {
pub(in crate::api) struct UserQuotaEntry {
pub(in crate::api) username: String,
pub(in crate::api) data_quota_bytes: u64,
/// Configured period may precede asynchronous runtime activation.
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
/// Effective schedule comes from the process authority, never desired TOML.
#[serde(flatten)]
pub(in crate::api) schedule: crate::stats::QuotaScheduleStatus,
pub(in crate::api) used_bytes: u64,
pub(in crate::api) last_reset_epoch_secs: u64,
}
@@ -81,6 +88,8 @@ pub(in crate::api) struct CreateUserRequest {
pub(in crate::api) max_tcp_conns: Option<usize>,
pub(in crate::api) expiration_rfc3339: Option<String>,
pub(in crate::api) data_quota_bytes: Option<u64>,
/// Desired UTC period; runtime activation is reported by the quota endpoint.
pub(in crate::api) quota_period: Option<crate::config::UserQuotaPeriod>,
pub(in crate::api) rate_limit_up_bps: Option<u64>,
pub(in crate::api) rate_limit_down_bps: Option<u64>,
pub(in crate::api) max_unique_ips: Option<usize>,
@@ -98,6 +107,9 @@ pub(in crate::api) struct PatchUserRequest {
pub(in crate::api) expiration_rfc3339: Patch<String>,
#[serde(default, deserialize_with = "patch_field")]
pub(in crate::api) data_quota_bytes: Patch<u64>,
/// Omitted preserves the period; null removes it.
#[serde(default, deserialize_with = "patch_field")]
pub(in crate::api) quota_period: Patch<crate::config::UserQuotaPeriod>,
#[serde(default, deserialize_with = "patch_field")]
pub(in crate::api) rate_limit_up_bps: Patch<u64>,
#[serde(default, deserialize_with = "patch_field")]
+3
View File
@@ -35,3 +35,6 @@ pub(super) use view::{build_user_quota_list, users_from_config};
#[cfg(test)]
mod tests;
#[cfg(test)]
mod period_tests;
+18
View File
@@ -58,6 +58,7 @@ async fn create_user_to_completion(
let credential_id = credential_id_from_hex(&secret)
.ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?;
let _guard = shared.mutation_lock.lock().await;
let quota_guard = shared.quota_state.lock().await;
let (mut cfg, base_revision) =
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
@@ -112,10 +113,18 @@ async fn create_user_to_completion(
cfg.access.user_enabled.insert(body.username.clone(), false);
}
if let Some(period) = body.quota_period {
cfg.access
.user_quota_periods
.insert(body.username.clone(), period);
}
cfg.validate()
.map_err(|e| ApiFailure::bad_request(format!("config validation failed: {}", e)))?;
let mut touched_sections = vec![AccessSection::Users];
if body.quota_period.is_some() {
touched_sections.push(AccessSection::UserQuotaPeriods);
}
if touches_user_ad_tags {
touched_sections.push(AccessSection::UserAdTags);
}
@@ -151,6 +160,14 @@ async fn create_user_to_completion(
cfg.access.is_user_enabled(&body.username),
);
shared.quota_state.fence_fresh_schedule(
&body.username,
&shared.active_runtime.load_full().config(),
credential_id,
cfg.access.user_quota_periods.get(&body.username).copied(),
cfg.access.user_data_quota.get(&body.username).copied(),
);
drop(quota_guard);
if let Some(limit) = updated_limit {
shared
.ip_tracker
@@ -187,6 +204,7 @@ async fn create_user_to_completion(
.then_some(cfg.access.user_max_tcp_conns_global_each)),
expiration_rfc3339: None,
data_quota_bytes: None,
quota_period: body.quota_period,
rate_limit_up_bps: body.rate_limit_up_bps.filter(|limit| *limit > 0),
rate_limit_down_bps: body.rate_limit_down_bps.filter(|limit| *limit > 0),
max_unique_ips: updated_limit,
+8 -1
View File
@@ -33,6 +33,7 @@ async fn rotate_secret_to_completion(
.ok_or_else(|| ApiFailure::internal("validated user secret could not be decoded"))?;
let _guard = shared.mutation_lock.lock().await;
let quota_guard = shared.quota_state.lock().await;
let (mut cfg, base_revision) =
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
@@ -59,6 +60,8 @@ async fn rotate_secret_to_completion(
credential_id,
cfg.access.is_user_enabled(user),
);
shared.quota_state.refresh_fresh_schedule(user, &cfg);
drop(quota_guard);
drop(_guard);
let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips();
@@ -106,6 +109,7 @@ async fn delete_user_to_completion(
shared: &ApiShared,
) -> Result<(String, String), ApiFailure> {
let _guard = shared.mutation_lock.lock().await;
let quota_guard = shared.quota_state.lock().await;
let (mut cfg, base_revision) =
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
@@ -138,6 +142,9 @@ async fn delete_user_to_completion(
if cfg.access.user_expirations.remove(user).is_some() {
touched_sections.push(AccessSection::UserExpirations);
}
if cfg.access.user_quota_periods.remove(user).is_some() {
touched_sections.push(AccessSection::UserQuotaPeriods);
}
if cfg.access.user_data_quota.remove(user).is_some() {
touched_sections.push(AccessSection::UserDataQuota);
}
@@ -161,7 +168,7 @@ async fn delete_user_to_completion(
let configured_users = cfg.access.users.keys().cloned().collect();
if let Err(error) = shared
.quota_state
.remove_user(&configured_users, user)
.remove_user_locked(&configured_users, user, quota_guard)
.await
{
warn!(
+445
View File
@@ -0,0 +1,445 @@
use super::*;
use crate::api::{ApiRuntimeState, events::ApiEventStore};
use crate::config::UserQuotaPeriod;
use crate::maestro::{
control_plane::ProcessControlPlane, generation::test_runtime_generation, reload::ReloadControl,
};
use crate::proxy::{shared_state::ProxySharedState, user_admission::UserAdmissionAuthority};
use crate::quota_state::QuotaStateOwner;
use crate::stats::QuotaStore;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64};
use tokio::sync::{Mutex, watch};
struct Fixture {
_dir: tempfile::TempDir,
shared: ApiShared,
store: Arc<QuotaStore>,
config_tx: watch::Sender<Arc<ProxyConfig>>,
}
impl Fixture {
fn new() -> Self {
let dir = tempfile::tempdir().unwrap();
let config_path = dir.path().join("config.toml");
std::fs::write(&config_path, "# preserve this operator comment\n[access.users]\nadmin = '00000000000000000000000000000000'\n").unwrap();
let cfg = ProxyConfig::load(&config_path).unwrap();
let store = Arc::new(QuotaStore::default());
let stats = Arc::new(Stats::with_quota_store(store.clone()));
let mut runtime = test_runtime_generation(1, cfg.clone());
let authority = UserAdmissionAuthority::new_with_quota_store(store.clone());
let proxy_shared = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
runtime.proxy_shared.direct_buffer_budget.clone(),
authority,
);
proxy_shared.activate_user_config_source(
1,
None,
&cfg.access.users,
&cfg.access.user_enabled,
);
let (config_tx, config_rx) = watch::channel(Arc::new(cfg.clone()));
let mutable = Arc::get_mut(&mut runtime).unwrap();
mutable.stats = stats.clone();
mutable.proxy_shared = proxy_shared.clone();
mutable.config_rx = config_rx;
let shared = ApiShared {
stats,
ip_tracker: runtime.ip_tracker.clone(),
me_pool: runtime.me_pool_runtime.clone(),
upstream_manager: runtime.upstream_manager.clone(),
config_path,
quota_state: QuotaStateOwner::new(dir.path().join("quota.json"), store.clone()),
detected_ips_rx: watch::channel((None, None)).1,
mutation_lock: Arc::new(Mutex::new(())),
minimal_cache: Arc::new(Mutex::new(None)),
runtime_edge_connections_cache: Arc::new(Mutex::new(None)),
runtime_edge_recompute_lock: Arc::new(Mutex::new(())),
cache_generation: Arc::new(AtomicU64::new(1)),
runtime_events: Arc::new(ApiEventStore::new(64)),
request_id: Arc::new(AtomicU64::new(1)),
runtime_state: Arc::new(ApiRuntimeState {
process_started_at_epoch_secs: 0,
config_reload_count: AtomicU64::new(0),
last_config_reload_epoch_secs: AtomicU64::new(0),
admission_open: AtomicBool::new(true),
}),
startup_tracker: Arc::new(crate::startup::StartupTracker::new(0)),
route_runtime: runtime.route_runtime.clone(),
proxy_shared,
reload_control: ReloadControl::channel(1).0,
active_runtime: Arc::new(arc_swap::ArcSwap::from(runtime)),
web_trace: crate::web::trace::WebTraceStore::new(
cfg.web.debug.clone(),
&cfg.web.limits,
),
web_runtime_rx: crate::web::control::WebRuntimeControl::new().subscribe(),
control_plane: ProcessControlPlane::new(),
};
Self {
_dir: dir,
shared,
store,
config_tx,
}
}
async fn publish(&self) {
let cfg = Arc::new(ProxyConfig::load(&self.shared.config_path).unwrap());
let guard = self.shared.quota_state.lock().await;
self.config_tx.send_replace(cfg.clone());
assert!(
self.shared
.proxy_shared
.apply_user_config_from_source(1, &cfg.access.users, &cfg.access.user_enabled)
.is_some()
);
self.shared.proxy_shared.publish_quota_config(1, &cfg, true);
drop(guard);
self.shared
.quota_state
.process_schedules(crate::quota_state::now_epoch_secs(), false)
.await
.unwrap();
}
async fn finish(&self) {
assert!(
self.shared
.control_plane
.shutdown(std::time::Duration::from_secs(1))
.await
);
}
}
fn patch(value: serde_json::Value) -> PatchUserRequest {
serde_json::from_value(value).unwrap()
}
#[tokio::test]
async fn quota_period_api_updates_included_owner_and_preserves_root_bytes() {
let fixture = Fixture::new();
let root = "# root remains unchanged\ninclude = \"users.toml\"\n";
let included = fixture.shared.config_path.with_file_name("users.toml");
std::fs::write(&included, "[access.users]\nadmin = '00000000000000000000000000000000'\n[access.user_data_quota]\nadmin = 1024\n[access.user_quota_periods]\nadmin = 'daily'\n").unwrap();
std::fs::write(&fixture.shared.config_path, root).unwrap();
patch_user(
"admin",
patch(serde_json::json!({"quota_period":"monthly"})),
None,
&fixture.shared,
)
.await
.unwrap();
assert_eq!(
std::fs::read_to_string(&fixture.shared.config_path).unwrap(),
root
);
let cfg = ProxyConfig::load(&fixture.shared.config_path).unwrap();
assert_eq!(
cfg.access.user_quota_periods["admin"],
UserQuotaPeriod::Monthly
);
assert!(
std::fs::read_to_string(included)
.unwrap()
.contains("monthly")
);
fixture.finish().await;
}
#[tokio::test]
async fn quota_period_prepared_before_recreation_cannot_attach_retired_schedule() {
let fixture = Fixture::new();
let create = |period| {
serde_json::from_value(serde_json::json!({
"username":"alice", "secret":"0123456789abcdef0123456789abcdef",
"data_quota_bytes":1024, "quota_period":period
}))
.unwrap()
};
create_user(create("daily"), None, &fixture.shared)
.await
.unwrap();
fixture.publish().await;
// Capture a real source before recreation; its asynchronous preparation finishes later.
let mut stale = ProxyConfig::parse_source(&fixture.shared.config_path).unwrap();
stale.config.general.ad_tag = Some("11111111111111111111111111111111".into());
delete_user("alice", None, &fixture.shared).await.unwrap();
create_user(create("monthly"), None, &fixture.shared)
.await
.unwrap();
let stale = Arc::new(stale.prepare().await.unwrap().config);
let guard = fixture.shared.quota_state.lock().await;
fixture.config_tx.send_replace(stale.clone());
fixture.shared.proxy_shared.apply_user_config_from_source(
1,
&stale.access.users,
&stale.access.user_enabled,
);
fixture
.shared
.proxy_shared
.publish_quota_config(1, &stale, true);
drop(guard);
fixture
.shared
.quota_state
.process_schedules(crate::quota_state::now_epoch_secs(), false)
.await
.unwrap();
assert_eq!(
fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.quota_period,
None
);
assert_eq!(
fixture
.shared
.quota_state
.schedule_status("alice", Some(UserQuotaPeriod::Monthly))
.quota_schedule_state,
"pending"
);
fixture.publish().await;
assert_eq!(
fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.quota_period,
Some(UserQuotaPeriod::Monthly)
);
fixture.finish().await;
}
#[tokio::test]
async fn quota_period_fresh_patch_and_rotation_accept_matching_stale_epoch_candidate() {
let fixture = Fixture::new();
create_user(
serde_json::from_value(serde_json::json!({
"username":"alice", "data_quota_bytes":1024, "quota_period":"daily"
}))
.unwrap(),
None,
&fixture.shared,
)
.await
.unwrap();
patch_user(
"alice",
patch(serde_json::json!({"quota_period":"monthly"})),
None,
&fixture.shared,
)
.await
.unwrap();
rotate_secret(
"alice",
RotateSecretRequest {
secret: Some("11111111111111111111111111111111".into()),
},
None,
&fixture.shared,
)
.await
.unwrap();
let cfg = Arc::new(ProxyConfig::load(&fixture.shared.config_path).unwrap());
let guard = fixture.shared.quota_state.lock().await;
fixture
.shared
.proxy_shared
.publish_quota_config(1, &cfg, false);
drop(guard);
fixture
.shared
.quota_state
.process_schedules(crate::quota_state::now_epoch_secs(), false)
.await
.unwrap();
assert_eq!(
fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.quota_period,
Some(UserQuotaPeriod::Monthly)
);
fixture.finish().await;
}
#[tokio::test]
async fn quota_period_api_create_patch_remove_delete_and_effective_status() {
let fixture = Fixture::new();
let created = create_user(serde_json::from_value(serde_json::json!({
"username": "alice", "secret": "0123456789abcdef0123456789abcdef", "data_quota_bytes": 1024, "quota_period": "daily"
})).unwrap(), None, &fixture.shared).await.unwrap().0;
assert_eq!(created.user.quota_period, Some(UserQuotaPeriod::Daily));
assert_eq!(
fixture
.shared
.quota_state
.schedule_status("alice", Some(UserQuotaPeriod::Daily))
.quota_schedule_state,
"pending"
);
fixture.store.current_handle("alice").unwrap().charge(300);
fixture.publish().await;
assert_eq!(fixture.store.used("alice"), 300);
let before = fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint();
patch_user(
"alice",
patch(serde_json::json!({"quota_period":"monthly"})),
None,
&fixture.shared,
)
.await
.unwrap();
let status = fixture
.shared
.quota_state
.schedule_status("alice", Some(UserQuotaPeriod::Monthly));
assert_eq!(status.active_quota_period, Some(UserQuotaPeriod::Daily));
assert_eq!(status.quota_schedule_state, "pending");
fixture.publish().await;
assert_eq!(fixture.store.used("alice"), 300);
assert!(
fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.period_cursor_epoch_secs
>= before.period_cursor_epoch_secs
);
let bytes = std::fs::read(&fixture.shared.config_path).unwrap();
assert!(
patch_user(
"alice",
patch(serde_json::json!({"data_quota_bytes":null})),
None,
&fixture.shared
)
.await
.is_err()
);
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), bytes);
patch_user(
"alice",
patch(serde_json::json!({"data_quota_bytes":null,"quota_period":null})),
None,
&fixture.shared,
)
.await
.unwrap();
fixture.publish().await;
assert_eq!(fixture.store.used("alice"), 300);
assert_eq!(
fixture
.shared
.quota_state
.schedule_status("alice", None)
.quota_schedule_state,
"disabled"
);
patch_user(
"alice",
patch(serde_json::json!({"data_quota_bytes":1024,"quota_period":"weekly"})),
None,
&fixture.shared,
)
.await
.unwrap();
fixture.publish().await;
delete_user("alice", None, &fixture.shared).await.unwrap();
let cfg = ProxyConfig::load(&fixture.shared.config_path).unwrap();
assert!(!cfg.access.users.contains_key("alice"));
assert!(!cfg.access.user_quota_periods.contains_key("alice"));
assert!(fixture.store.current_handle("alice").is_none());
assert!(
std::fs::read_to_string(&fixture.shared.config_path)
.unwrap()
.contains("# preserve this operator comment")
);
fixture.finish().await;
}
#[tokio::test]
async fn quota_period_api_early_rotation_and_revision_conflicts_are_safe() {
let fixture = Fixture::new();
let (_, old_revision) = create_user(
serde_json::from_value(serde_json::json!({
"username":"alice", "data_quota_bytes":1024, "quota_period":"monthly"
}))
.unwrap(),
None,
&fixture.shared,
)
.await
.unwrap();
fixture.store.current_handle("alice").unwrap().charge(200);
rotate_secret(
"alice",
RotateSecretRequest {
secret: Some("11111111111111111111111111111111".into()),
},
None,
&fixture.shared,
)
.await
.unwrap();
fixture.publish().await;
assert_eq!(fixture.store.used("alice"), 200);
assert_eq!(
fixture
.shared
.quota_state
.schedule_status("alice", Some(UserQuotaPeriod::Monthly))
.quota_schedule_state,
"active"
);
let bytes = std::fs::read(&fixture.shared.config_path).unwrap();
let error = patch_user(
"alice",
patch(serde_json::json!({"quota_period":"daily"})),
Some(old_revision),
&fixture.shared,
)
.await
.err()
.unwrap();
assert_eq!(error.code, "revision_conflict");
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), bytes);
fixture.finish().await;
}
#[test]
fn quota_period_patch_preserves_omitted_null_and_set_semantics() {
assert!(matches!(
patch(serde_json::json!({})).quota_period,
Patch::Unchanged
));
assert!(matches!(
patch(serde_json::json!({"quota_period":null})).quota_period,
Patch::Remove
));
assert!(matches!(
patch(serde_json::json!({"quota_period":"weekly"})).quota_period,
Patch::Set(UserQuotaPeriod::Weekly)
));
assert!(
serde_json::from_value::<PatchUserRequest>(serde_json::json!({"quota_period":"Weekly"}))
.is_err()
);
}
+21
View File
@@ -48,6 +48,7 @@ async fn patch_user_to_completion(
}
let expiration = parse_patch_expiration(&body.expiration_rfc3339)?;
let _guard = shared.mutation_lock.lock().await;
let quota_guard = shared.quota_state.lock().await;
let (mut cfg, base_revision) =
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
@@ -93,6 +94,18 @@ async fn patch_user_to_completion(
.insert(user.to_string(), expiration);
}
}
let touches_user_quota_periods = !matches!(body.quota_period, Patch::Unchanged);
match body.quota_period {
Patch::Unchanged => {}
Patch::Remove => {
cfg.access.user_quota_periods.remove(user);
}
Patch::Set(period) => {
cfg.access
.user_quota_periods
.insert(user.to_string(), period);
}
}
match body.data_quota_bytes {
Patch::Unchanged => {}
Patch::Remove => {
@@ -169,6 +182,9 @@ async fn patch_user_to_completion(
};
let mut touched_sections = Vec::new();
if touches_user_quota_periods {
touched_sections.push(AccessSection::UserQuotaPeriods);
}
if touches_users {
touched_sections.push(AccessSection::Users);
}
@@ -212,6 +228,8 @@ async fn patch_user_to_completion(
cfg.access.is_user_enabled(user),
);
}
shared.quota_state.refresh_fresh_schedule(user, &cfg);
drop(quota_guard);
match max_unique_ips_change {
Some(Some(limit)) => shared.ip_tracker.set_user_limit(user, limit).await,
Some(None) => shared.ip_tracker.remove_user_limit(user).await,
@@ -259,6 +277,7 @@ async fn set_user_enabled_to_completion(
shared: &ApiShared,
) -> Result<(UserInfo, String), ApiFailure> {
let _guard = shared.mutation_lock.lock().await;
let quota_guard = shared.quota_state.lock().await;
let (mut cfg, base_revision) =
load_config_for_mutation(&shared.config_path, expected_revision.as_deref()).await?;
@@ -294,6 +313,8 @@ async fn set_user_enabled_to_completion(
shared
.proxy_shared
.stage_user_credential(user, credential_id, enabled);
shared.quota_state.refresh_fresh_schedule(user, &cfg);
drop(quota_guard);
drop(_guard);
let (detected_ip_v4, detected_ip_v6) = shared.detected_link_ips();
+7 -1
View File
@@ -51,6 +51,7 @@ pub(in crate::api) async fn users_from_config(
.get(&username)
.map(chrono::DateTime::<chrono::Utc>::to_rfc3339),
data_quota_bytes: cfg.access.user_data_quota.get(&username).copied(),
quota_period: cfg.access.user_quota_periods.get(&username).copied(),
rate_limit_up_bps: cfg
.access
.user_rate_limits
@@ -102,7 +103,12 @@ pub(in crate::api) fn build_user_quota_list(cfg: &ProxyConfig, stats: &Stats) ->
.map(|entry| (entry.used_bytes, entry.last_reset_epoch_secs))
.unwrap_or((0, 0));
users.push(UserQuotaEntry {
username,
username: username.clone(),
quota_period: cfg.access.user_quota_periods.get(&username).copied(),
schedule: crate::stats::QuotaScheduleStatus {
quota_schedule_state: "disabled",
..Default::default()
},
data_quota_bytes,
used_bytes,
last_reset_epoch_secs,
+3
View File
@@ -80,3 +80,6 @@ mod conveyor_tests;
// Host groups are validated and published as one complete immutable snapshot.
#[cfg(test)]
mod host_list_tests;
#[cfg(test)]
mod quota_period_tests;
+5
View File
@@ -84,6 +84,8 @@ pub struct HotFields {
pub user_max_tcp_conns_global_each: usize,
pub user_expirations: std::collections::HashMap<String, chrono::DateTime<chrono::Utc>>,
pub user_data_quota: std::collections::HashMap<String, u64>,
/// Recurring UTC calendars accepted by the active runtime source.
pub user_quota_periods: std::collections::HashMap<String, crate::config::UserQuotaPeriod>,
pub user_rate_limits: std::collections::HashMap<String, crate::config::RateLimitBps>,
pub cidr_rate_limits: std::collections::HashMap<CidrRateLimitKey, crate::config::RateLimitBps>,
pub user_max_unique_ips: std::collections::HashMap<String, usize>,
@@ -215,6 +217,7 @@ impl HotFields {
user_max_tcp_conns_global_each: cfg.access.user_max_tcp_conns_global_each,
user_expirations: cfg.access.user_expirations.clone(),
user_data_quota: cfg.access.user_data_quota.clone(),
user_quota_periods: cfg.access.user_quota_periods.clone(),
user_rate_limits: cfg.access.user_rate_limits.clone(),
cidr_rate_limits: cfg.access.cidr_rate_limits.clone(),
user_max_unique_ips: cfg.access.user_max_unique_ips.clone(),
@@ -228,6 +231,7 @@ impl HotFields {
pub(super) fn overlay_hot_fields(old: &ProxyConfig, new: &ProxyConfig) -> ProxyConfig {
let mut cfg = old.clone();
cfg.quota_source_read_at = new.quota_source_read_at;
cfg.general.log_level = new.general.log_level.clone();
cfg.general.ad_tag = new.general.ad_tag.clone();
@@ -336,6 +340,7 @@ pub(super) fn overlay_hot_fields(old: &ProxyConfig, new: &ProxyConfig) -> ProxyC
cfg.access.user_max_tcp_conns_global_each = new.access.user_max_tcp_conns_global_each;
cfg.access.user_expirations = new.access.user_expirations.clone();
cfg.access.user_data_quota = new.access.user_data_quota.clone();
cfg.access.user_quota_periods = new.access.user_quota_periods.clone();
cfg.access.user_rate_limits = new.access.user_rate_limits.clone();
cfg.access.cidr_rate_limits = new.access.cidr_rate_limits.clone();
cfg.access.user_max_unique_ips = new.access.user_max_unique_ips.clone();
@@ -0,0 +1,33 @@
use super::*;
#[test]
fn quota_period_reload_is_hot_and_invalid_candidates_preserve_the_snapshot() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
let source = "[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n[access.user_data_quota]\nalice = 1024\n";
std::fs::write(&path, source).unwrap();
let initial = Arc::new(ProxyConfig::load(&path).unwrap());
let (config_tx, _config_rx) = watch::channel(initial.clone());
let (log_tx, _log_rx) = watch::channel(initial.general.log_level.clone());
let mut state = ReloadState::new(None);
for (period, valid) in [
(Some("daily"), true),
(Some("monthly"), true),
(Some("yearly"), false),
(None, true),
] {
let suffix = period
.map(|p| format!("[access.user_quota_periods]\nalice = '{p}'\n"))
.unwrap_or_default();
std::fs::write(&path, format!("{source}{suffix}")).unwrap();
let before = config_tx.borrow().clone();
reload_config(&path, &config_tx, &log_tx, None, None, &mut state);
let after = config_tx.borrow().clone();
if valid {
let json = serde_json::to_value(after.access.user_quota_periods.get("alice")).unwrap();
assert_eq!(json, serde_json::to_value(period).unwrap());
} else {
assert!(Arc::ptr_eq(&before, &after));
}
}
}
+6
View File
@@ -463,6 +463,12 @@ pub(super) fn log_changes(
new_hot.user_expirations.len()
);
}
if old_hot.user_quota_periods != new_hot.user_quota_periods {
info!(
"config reload: user_quota_periods updated ({} entries)",
new_hot.user_quota_periods.len()
);
}
if old_hot.user_data_quota != new_hot.user_data_quota {
info!(
"config reload: user_data_quota updated ({} entries)",
+2
View File
@@ -372,6 +372,8 @@ async fn candidate_watcher_waits_for_activation_and_reconciles_disk() {
cancellation.clone(),
None,
Some(activation_rx),
None,
None,
);
let watcher = tokio::spawn(watcher);
+20 -1
View File
@@ -156,6 +156,7 @@ fn reload_config_with_resolver(
detected_ip_v6: Option<IpAddr>,
reload_state: &mut ReloadState,
dns_resolver: Option<&crate::network::dns_overrides::GenerationDnsResolver>,
refresh_quota_source: bool,
) -> Option<WatchManifest> {
let (loaded, applied_cfg) = prepared;
let LoadedConfig {
@@ -178,6 +179,9 @@ fn reload_config_with_resolver(
}
if !hot_changed {
if refresh_quota_source {
config_tx.send(Arc::new(applied_cfg)).ok();
}
reload_state.mark_applied(rendered_hash);
return Some(next_manifest);
}
@@ -228,6 +232,7 @@ pub(super) fn reload_config(
detected_ip_v6,
reload_state,
None,
false,
)
}
@@ -248,6 +253,8 @@ pub fn spawn_config_watcher(
cancellation: tokio_util::sync::CancellationToken,
dns_resolver: Option<Arc<crate::network::dns_overrides::GenerationDnsResolver>>,
mut activation: Option<watch::Receiver<bool>>,
quota_mutation: Option<Arc<tokio::sync::Mutex<()>>>,
quota_source_refresh: Option<Arc<dyn Fn() -> bool + Send + Sync>>,
) -> (
watch::Receiver<Arc<ProxyConfig>>,
watch::Receiver<LogLevel>,
@@ -365,7 +372,12 @@ pub fn spawn_config_watcher(
if poll_watcher.is_some() {
info!("config watcher: poll watcher active (Docker/NFS safe)");
}
if initial_loaded.is_some() && !initial_matches_disk {
if initial_loaded.is_some()
&& (!initial_matches_disk
|| quota_source_refresh
.as_ref()
.is_some_and(|refresh| refresh()))
{
let _ = notify_tx.try_send(());
}
@@ -409,6 +421,8 @@ pub fn spawn_config_watcher(
&mut reload_state,
dns_resolver.clone(),
force,
quota_mutation.as_ref(),
quota_source_refresh.as_deref(),
)
.await;
if next_manifest.is_none() {
@@ -423,6 +437,8 @@ pub fn spawn_config_watcher(
&mut reload_state,
dns_resolver.clone(),
force,
quota_mutation.as_ref(),
quota_source_refresh.as_deref(),
)
.await;
}
@@ -461,3 +477,6 @@ mod path_tests {
#[cfg(test)]
mod decoy_dns_tests;
#[cfg(test)]
mod quota_period_tests;
@@ -48,6 +48,7 @@ async fn decoy_dns_watcher_publishes_dns_only_changes_once() {
None,
&mut state,
None,
false,
)
.unwrap();
assert!(config_rx.has_changed().unwrap());
@@ -62,13 +63,14 @@ async fn decoy_dns_watcher_publishes_dns_only_changes_once() {
None,
&mut state,
None,
false,
)
.unwrap();
assert!(!config_rx.has_changed().unwrap());
assert!(Arc::ptr_eq(&published, &config_rx.borrow()));
preparation::reload(
&path, &config_tx, &log_tx, None, None, &mut state, None, false,
&path, &config_tx, &log_tx, None, None, &mut state, None, false, None, None,
)
.await
.unwrap();
@@ -77,7 +79,7 @@ async fn decoy_dns_watcher_publishes_dns_only_changes_once() {
// Explicit reload bypasses the unchanged source hash and replaces the injected endpoint.
preparation::reload(
&path, &config_tx, &log_tx, None, None, &mut state, None, true,
&path, &config_tx, &log_tx, None, None, &mut state, None, true, None, None,
)
.await
.unwrap();
@@ -87,7 +89,7 @@ async fn decoy_dns_watcher_publishes_dns_only_changes_once() {
let current = config_rx.borrow_and_update().clone();
std::fs::write(&path, SOURCE.replace("localhost:18081", "localhost:18080")).unwrap();
preparation::reload(
&path, &config_tx, &log_tx, None, None, &mut state, None, false,
&path, &config_tx, &log_tx, None, None, &mut state, None, false, None, None,
)
.await
.unwrap();
@@ -111,7 +113,7 @@ async fn decoy_dns_unchanged_watcher_source_skips_resolution() {
let mut state = ReloadState::new(Some(source.rendered_hash));
assert!(preparation::source_matches_active(&initial, &source));
preparation::reload(
&path, &config_tx, &log_tx, None, None, &mut state, None, false,
&path, &config_tx, &log_tx, None, None, &mut state, None, false, None, None,
)
.await
.unwrap();
@@ -34,9 +34,30 @@ pub(super) async fn reload(
reload_state: &mut ReloadState,
dns_resolver: Option<Arc<crate::network::dns_overrides::GenerationDnsResolver>>,
force: bool,
quota_mutation: Option<&Arc<tokio::sync::Mutex<()>>>,
quota_source_refresh: Option<&(dyn Fn() -> bool + Send + Sync)>,
) -> Option<WatchManifest> {
// Serialize source acquisition with API lifecycle commits, then release before DNS.
// The source timestamp fences any later mutation during asynchronous preparation.
let guard = match quota_mutation {
Some(gate) => Some(gate.lock().await),
None => None,
};
let source = read_source(path).await?;
let manifest = WatchManifest::from_source_files(&source.source_files);
// A new file edit can restore the exact pre-creation bytes. Refresh only source
// evidence for fenced accounts, retaining prepared DNS and immutable WEB state.
if !force
&& quota_source_refresh.is_some_and(|refresh| refresh())
&& source_matches_active(&config_tx.borrow(), &source)
{
let mut refreshed = (**config_tx.borrow()).clone();
refreshed.quota_source_read_at = source.config.quota_source_read_at;
config_tx.send(Arc::new(refreshed)).ok();
reload_state.mark_applied(source.rendered_hash);
return Some(manifest);
}
drop(guard);
if !force && reload_state.is_applied(source.rendered_hash) {
return Some(manifest);
}
@@ -66,6 +87,7 @@ pub(super) async fn reload(
detected_ip_v6,
reload_state,
dns_resolver.as_deref(),
quota_source_refresh.is_some_and(|refresh| refresh()),
),
Ok(Err(error)) => {
error!(
@@ -0,0 +1,105 @@
use super::*;
#[tokio::test]
async fn quota_period_unchanged_source_refreshes_only_fenced_accounts_without_dns() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
std::fs::write(
&path,
"[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n",
)
.unwrap();
let source = ProxyConfig::parse_source(&path).unwrap();
let initial = Arc::new(source.config);
let (tx, mut rx) = watch::channel(initial.clone());
let (log, _) = watch::channel(initial.general.log_level.clone());
let mut state = ReloadState::new(Some(source.rendered_hash));
let gate = Arc::new(tokio::sync::Mutex::new(()));
preparation::reload(
&path,
&tx,
&log,
None,
None,
&mut state,
None,
false,
Some(&gate),
Some(&|| true),
)
.await
.unwrap();
assert!(rx.has_changed().unwrap());
let refreshed = rx.borrow_and_update().clone();
assert!(config_equal(&initial, &refreshed));
assert!(refreshed.quota_source_read_at > initial.quota_source_read_at);
preparation::reload(
&path,
&tx,
&log,
None,
None,
&mut state,
None,
false,
Some(&gate),
Some(&|| false),
)
.await
.unwrap();
assert!(!rx.has_changed().unwrap());
assert!(Arc::ptr_eq(&refreshed, &rx.borrow()));
preparation::reload(
&path,
&tx,
&log,
None,
None,
&mut state,
None,
true,
Some(&gate),
Some(&|| true),
)
.await
.unwrap();
assert!(rx.has_changed().unwrap());
assert!(rx.borrow().quota_source_read_at > refreshed.quota_source_read_at);
assert!(config_equal(&refreshed, &rx.borrow()));
}
#[tokio::test]
async fn quota_period_watcher_initialization_refreshes_pending_identical_source() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.toml");
std::fs::write(
&path,
"[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n",
)
.unwrap();
let initial = Arc::new(ProxyConfig::load(&path).unwrap());
let cancel = tokio_util::sync::CancellationToken::new();
let (mut rx, _, task) = spawn_config_watcher(
path,
initial.clone(),
None,
None,
cancel.clone(),
None,
None,
Some(Arc::new(tokio::sync::Mutex::new(()))),
Some(Arc::new(|| true)),
);
let task = tokio::spawn(task);
tokio::time::timeout(Duration::from_secs(5), rx.changed())
.await
.unwrap()
.unwrap();
assert!(config_equal(&initial, &rx.borrow()));
assert!(rx.borrow().quota_source_read_at > initial.quota_source_read_at);
cancel.cancel();
tokio::time::timeout(Duration::from_secs(5), task)
.await
.unwrap()
.unwrap();
}
+14 -2
View File
@@ -156,6 +156,10 @@ pub struct ProxyConfig {
#[serde(default)]
pub default_dc: Option<u8>,
/// Monotonic start of source acquisition, used only for cold quota policy fencing.
#[serde(skip)]
pub(crate) quota_source_read_at: Option<std::time::Instant>,
/// Precomputed authentication snapshot for handshake hot paths.
#[serde(skip)]
pub(crate) runtime_user_auth: Option<Arc<UserAuthSnapshot>>,
@@ -177,8 +181,11 @@ impl ProxyConfig {
path: P,
source_overrides: &BTreeMap<PathBuf, String>,
) -> Result<LoadedConfig> {
let read_at = std::time::Instant::now();
let graph = Self::read_source_graph_with_overrides(path, source_overrides)?;
Self::load_source_graph(graph)
let mut loaded = Self::load_source_graph(graph)?;
loaded.config.quota_source_read_at = Some(read_at);
Ok(loaded)
}
/// Captures the raw include graph without requiring typed config validity.
@@ -259,7 +266,11 @@ impl ProxyConfig {
path: P,
overrides: &BTreeMap<PathBuf, String>,
) -> Result<ParsedConfigSource> {
pipeline::parse_source_graph(Self::read_source_graph_with_overrides(path, overrides)?)
let read_at = std::time::Instant::now();
let mut parsed =
pipeline::parse_source_graph(Self::read_source_graph_with_overrides(path, overrides)?)?;
parsed.config.quota_source_read_at = Some(read_at);
Ok(parsed)
}
/// Loads and prepares a complete runtime candidate without blocking async workers.
@@ -318,6 +329,7 @@ impl ProxyConfig {
/// Validates cross-field configuration invariants after deserialization.
pub fn validate(&self) -> Result<()> {
self.access.validate_quota_periods()?;
if self.access.users.is_empty() {
return Err(ProxyError::Config("No users configured".to_string()));
}
+1
View File
@@ -455,6 +455,7 @@ const ACCESS_CONFIG_KEYS: &[&str] = &[
"user_max_tcp_conns_global_each",
"user_expirations",
"user_data_quota",
"user_quota_periods",
"user_rate_limits",
"cidr_rate_limits",
"user_max_unique_ips",
+1
View File
@@ -1,6 +1,7 @@
use super::*;
pub(super) fn validate(config: &mut ProxyConfig) -> Result<()> {
config.access.validate_quota_periods()?;
if config.general.me_writer_cmd_channel_capacity == 0 {
return Err(ProxyError::Config(
"general.me_writer_cmd_channel_capacity must be > 0".to_string(),
+2
View File
@@ -46,6 +46,8 @@ mod legacy_policy_tests;
mod me_route_tests;
#[path = "load_basic_tests/me_startup_tests.rs"]
mod me_startup_tests;
#[path = "load_basic_tests/quota_period_tests.rs"]
mod quota_period_tests;
#[path = "load_basic_tests/source_security_tests.rs"]
mod source_security_tests;
#[path = "load_basic_tests/synlimit_mss_tests.rs"]
@@ -0,0 +1,69 @@
use super::*;
#[test]
fn quota_period_configuration_survives_round_trip() {
let cfg: ProxyConfig = toml::from_str(
r#"
[access.users]
alice = "0123456789abcdef0123456789abcdef"
[access.user_data_quota]
alice = 1024
[access.user_quota_periods]
alice = "monthly"
"#,
)
.unwrap();
let encoded = serde_json::to_value(&cfg).unwrap();
assert_eq!(encoded["access"]["user_quota_periods"]["alice"], "monthly");
}
#[test]
fn quota_period_requires_a_configured_positive_quota_in_both_validation_paths() {
for tail in [
"[access.user_quota_periods]\nghost = 'daily'\n",
"[access.user_quota_periods]\nalice = 'daily'\n",
"[access.user_data_quota]\nalice = 0\n[access.user_quota_periods]\nalice = 'daily'\n",
] {
let source = format!("[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n{tail}");
let cfg: ProxyConfig = toml::from_str(&source).unwrap();
assert!(
cfg.validate()
.unwrap_err()
.to_string()
.contains("user_quota_periods")
);
assert!(load_config_error_from_temp_toml(&source).contains("user_quota_periods"));
}
for period in ["daily", "weekly", "monthly"] {
let source = format!(
"[general]\nconfig_strict = true\n[access.users]\nalice = '0123456789abcdef0123456789abcdef'\n[access.user_data_quota]\nalice = 1024\n[access.user_quota_periods]\nalice = '{period}'\n"
);
assert_eq!(
load_config_from_temp_toml(&source)
.access
.user_quota_periods
.len(),
1
);
}
}
#[test]
fn quota_period_tokens_are_strict_and_inventory_is_bounded() {
for token in ["'Daily'", "'yearly'", "42", "true"] {
assert!(
toml::from_str::<ProxyConfig>(&format!("[access.user_quota_periods]\nalice = {token}"))
.is_err()
);
}
let mut cfg = ProxyConfig::default();
assert!(cfg.access.user_quota_periods.is_empty());
cfg.access.users = (0..65_537)
.map(|n| (format!("u{n}"), "0123456789abcdef0123456789abcdef".into()))
.collect();
cfg.access.user_data_quota.insert("u0".into(), 1);
cfg.access
.user_quota_periods
.insert("u0".into(), crate::config::UserQuotaPeriod::Daily);
assert!(cfg.validate().unwrap_err().to_string().contains("65536"));
}
+3
View File
@@ -22,6 +22,8 @@ mod links;
mod logging;
mod network;
mod policies;
// UTC quota calendar tokens have no process-local timezone dependency.
mod quota_period;
mod server;
mod web;
// WEB carrier tokens and fixed-slot policy helpers remain independent from bulky config types.
@@ -44,6 +46,7 @@ pub use policies::{
MeBindStaleMode, MeFloorMode, MeRouteNoWriterMode, MeSocksKdfPolicy, MeTelemetryLevel,
MeWriterPickMode, RstOnCloseMode, TelemetryConfig, UserMaxUniqueIpsMode,
};
pub use quota_period::UserQuotaPeriod;
#[allow(unused_imports)]
pub use server::{
CLIENT_MSS_2IN8, CLIENT_MSS_EXTREME_LOW, CLIENT_MSS_MAX, CLIENT_MSS_MIN, CLIENT_MSS_TSPU,
+27
View File
@@ -30,6 +30,10 @@ pub struct AccessConfig {
#[serde(default)]
pub user_data_quota: HashMap<String, u64>,
/// Optional UTC calendar for users with a positive data quota.
#[serde(default)]
pub user_quota_periods: HashMap<String, UserQuotaPeriod>,
/// Per-user transport rate limits in bits-per-second.
///
/// Each entry supports independent upload (`up_bps`) and download
@@ -90,6 +94,7 @@ impl Default for AccessConfig {
user_max_tcp_conns_global_each: default_user_max_tcp_conns_global_each(),
user_expirations: HashMap::new(),
user_data_quota: HashMap::new(),
user_quota_periods: HashMap::new(),
user_rate_limits: HashMap::new(),
cidr_rate_limits: HashMap::new(),
user_source_deny: HashMap::new(),
@@ -105,6 +110,28 @@ impl Default for AccessConfig {
}
impl AccessConfig {
/// Validates the opt-in calendar against the bounded durable quota inventory.
pub(crate) fn validate_quota_periods(&self) -> crate::error::Result<()> {
if !self.user_quota_periods.is_empty() && self.users.len() > 65_536 {
return Err(crate::error::ProxyError::Config(
"quota periods support at most 65536 configured users".to_string(),
));
}
for user in self.user_quota_periods.keys() {
if !self.users.contains_key(user)
|| !self
.user_data_quota
.get(user)
.is_some_and(|quota| *quota > 0)
{
return Err(crate::error::ProxyError::Config(format!(
"access.user_quota_periods.{user} requires an existing user with a positive data quota"
)));
}
}
Ok(())
}
pub fn is_user_enabled(&self, username: &str) -> bool {
self.user_enabled.get(username).copied().unwrap_or(true)
}
+114
View File
@@ -0,0 +1,114 @@
use serde::{Deserialize, Serialize};
/// Calendar used for recurring quota resets; all boundaries are UTC.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum UserQuotaPeriod {
/// Each UTC midnight.
Daily,
/// Monday at UTC midnight.
Weekly,
/// The first day of each UTC month.
Monthly,
}
impl UserQuotaPeriod {
/// Returns the UTC boundary at or before the supplied Unix timestamp.
pub(crate) fn start(self, epoch: u64) -> std::io::Result<u64> {
use chrono::Datelike;
let timestamp = i64::try_from(epoch).map_err(|_| invalid_time())?;
let date = chrono::DateTime::from_timestamp(timestamp, 0)
.ok_or_else(invalid_time)?
.date_naive();
let start = match self {
Self::Daily => Some(date),
Self::Weekly => date.checked_sub_days(chrono::Days::new(u64::from(
date.weekday().num_days_from_monday(),
))),
Self::Monthly => date.with_day(1),
}
.ok_or_else(invalid_time)?;
u64::try_from(
start
.and_hms_opt(0, 0, 0)
.ok_or_else(invalid_time)?
.and_utc()
.timestamp(),
)
.map_err(|_| invalid_time())
}
/// Returns the first UTC calendar boundary strictly after the cursor.
pub(crate) fn next(self, cursor: u64) -> std::io::Result<u64> {
let start = self.start(cursor)?;
let date = chrono::DateTime::from_timestamp(start as i64, 0)
.ok_or_else(invalid_time)?
.date_naive();
let next = match self {
Self::Daily => date.checked_add_days(chrono::Days::new(1)),
Self::Weekly => date.checked_add_days(chrono::Days::new(7)),
Self::Monthly => date.checked_add_months(chrono::Months::new(1)),
}
.ok_or_else(invalid_time)?;
u64::try_from(
next.and_hms_opt(0, 0, 0)
.ok_or_else(invalid_time)?
.and_utc()
.timestamp(),
)
.map_err(|_| invalid_time())
}
}
fn invalid_time() -> std::io::Error {
std::io::Error::new(
std::io::ErrorKind::InvalidData,
"quota calendar timestamp is outside the supported UTC range",
)
}
#[cfg(test)]
mod tests {
use super::*;
fn epoch(value: &str) -> u64 {
chrono::DateTime::parse_from_rfc3339(value)
.unwrap()
.timestamp() as u64
}
#[test]
fn quota_calendar_boundaries_are_strict_and_handle_leap_years() {
for (period, now, start, next) in [
(
UserQuotaPeriod::Daily,
"2024-02-29T23:59:59Z",
"2024-02-29T00:00:00Z",
"2024-03-01T00:00:00Z",
),
(
UserQuotaPeriod::Weekly,
"2024-12-31T12:00:00Z",
"2024-12-30T00:00:00Z",
"2025-01-06T00:00:00Z",
),
(
UserQuotaPeriod::Monthly,
"2024-02-15T12:00:00Z",
"2024-02-01T00:00:00Z",
"2024-03-01T00:00:00Z",
),
(
UserQuotaPeriod::Monthly,
"2023-12-31T23:59:59Z",
"2023-12-01T00:00:00Z",
"2024-01-01T00:00:00Z",
),
] {
assert_eq!(period.start(epoch(now)).unwrap(), epoch(start));
assert_eq!(period.next(epoch(now)).unwrap(), epoch(next));
assert_eq!(period.next(epoch(start)).unwrap(), epoch(next));
}
assert!(UserQuotaPeriod::Daily.next(u64::MAX).is_err());
}
}
+17 -3
View File
@@ -1,4 +1,3 @@
use std::collections::BTreeSet;
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
@@ -69,8 +68,7 @@ pub(super) async fn run_telemt_core(
let quota_state_path = config.general.quota_state_path.clone();
let quota_state =
crate::quota_state::QuotaStateOwner::new(quota_state_path, quota_store.clone());
let configured_quota_users = config.access.users.keys().cloned().collect::<BTreeSet<_>>();
quota_state.load(&configured_quota_users).await;
quota_state.load_for_config(&config).await?;
let upstream_manager = Arc::new(
UpstreamManager::new(
@@ -131,12 +129,28 @@ pub(super) async fn run_telemt_core(
traffic_limiter,
user_admission,
);
let quota_guard = quota_state.lock().await;
let _ = shared_state.activate_user_config_source(
1,
None,
&config.access.users,
&config.access.user_enabled,
);
shared_state.publish_quota_config(1, &Arc::new(config.clone()), true);
drop(quota_guard);
if let Err(error) = quota_state
.process_schedules(
crate::quota_state::now_epoch_secs(),
!config.access.user_quota_periods.is_empty(),
)
.await
{
if !config.access.user_quota_periods.is_empty() {
return Err(error.into());
}
warn!(error = %error, "Deferred quota calendar cleanup until persistence recovers");
}
quota_state.spawn_scheduler(&process_control_plane);
let max_connections_limit = if config.server.max_connections == 0 {
Semaphore::MAX_PERMITS
} else {
+8
View File
@@ -156,6 +156,7 @@ pub(crate) async fn spawn_runtime_tasks(
Some("spawn config hot-reload watcher".to_string()),
)
.await;
let quota_source_authority = shared_state.user_admission();
let (config_rx, log_level_rx, config_watcher_task) = spawn_config_watcher(
config_path.to_path_buf(),
config.clone(),
@@ -164,6 +165,10 @@ pub(crate) async fn spawn_runtime_tasks(
task_scope.cancellation_token(),
Some(upstream_manager.dns_resolver()),
config_watcher_activation,
Some(shared_state.quota_mutation_gate()),
Some(Arc::new(move || {
quota_source_authority.has_fresh_quota_schedules()
})),
);
task_scope.spawn(config_watcher_task);
startup_tracker
@@ -267,6 +272,7 @@ pub(crate) async fn spawn_runtime_tasks(
if config_rx_user_enabled.changed().await.is_err() {
break;
}
let quota_guard = shared_user_enabled.quota_mutation_gate().lock_owned().await;
let cfg = config_rx_user_enabled.borrow_and_update().clone();
let Some(cancelled_users) = shared_user_enabled.apply_user_config_from_source(
generation_id,
@@ -275,6 +281,8 @@ pub(crate) async fn spawn_runtime_tasks(
) else {
continue;
};
shared_user_enabled.publish_quota_config(generation_id, &cfg, true);
drop(quota_guard);
for (user, cancelled) in cancelled_users {
if cancelled > 0 {
info!(
+16
View File
@@ -238,6 +238,22 @@ impl ProxySharedState {
.apply_config_from_source(source_generation, users, user_enabled)
}
/// Serializes cold quota transitions without entering the relay hot path.
pub(crate) fn quota_mutation_gate(&self) -> Arc<tokio::sync::Mutex<()>> {
self.user_admission.quota_mutation_gate()
}
/// Reconciles the active source calendar under the process quota gate.
pub(crate) fn publish_quota_config(
&self,
source: u64,
config: &Arc<crate::config::ProxyConfig>,
accepted_after_mutation: bool,
) {
self.user_admission
.publish_quota_config(source, config, accepted_after_mutation);
}
/// Applies one persisted user mutation before asynchronous config reload.
pub(crate) fn stage_user(
&self,
+46
View File
@@ -121,6 +121,52 @@ impl UserAdmissionAuthority {
})
}
/// Requests source-evidence refresh only while a fresh account is fenced.
pub(crate) fn has_fresh_quota_schedules(&self) -> bool {
self.quota_store.has_fresh_schedules()
}
/// Shares the process quota gate with lifecycle publication callers.
pub(crate) fn quota_mutation_gate(&self) -> Arc<tokio::sync::Mutex<()>> {
self.quota_store.mutation_gate()
}
/// Publishes calendars only for credentials owned by the active config source.
/// The caller holds the process quota gate; disabled users retain their calendar.
pub(crate) fn publish_quota_config(
&self,
source: u64,
config: &Arc<crate::config::ProxyConfig>,
accepted_after_mutation: bool,
) {
let state = self.state.lock();
if state.active_config_source != Some(source) {
return;
}
let now = crate::quota_state::now_epoch_secs();
for (user, secret) in &config.access.users {
let Some(credential) = credential_id_from_hex(secret) else {
continue;
};
let Some(record) = state.users.get(user) else {
continue;
};
if record
.effective()
.is_some_and(|effective| effective.credential_id == credential)
{
self.quota_store.request_schedule(
user,
record.incarnation,
credential,
config,
now,
accepted_after_mutation,
);
}
}
}
/// Returns the mutation epoch used to reject stale candidate configuration.
pub(crate) fn epoch(&self) -> u64 {
self.state.lock().epoch
+135 -11
View File
@@ -7,7 +7,12 @@ use serde::{Deserialize, Serialize};
use tokio::sync::Mutex;
use tracing::{info, warn};
use crate::stats::{QuotaStore, UserQuotaSnapshot};
use crate::stats::{QuotaScheduleCheckpoint, QuotaStore, UserQuotaSnapshot};
// UTC period transactions share the existing bounded durable checkpoint owner.
mod periods;
#[cfg(test)]
mod test_io;
const QUOTA_STATE_MAX_BYTES: u64 = 16 * 1024 * 1024;
const QUOTA_STATE_MAX_USERS: usize = 65_536;
@@ -22,6 +27,9 @@ pub(crate) struct QuotaStateFile {
pub(crate) struct QuotaUserState {
pub(crate) used_bytes: u64,
pub(crate) last_reset_epoch_secs: u64,
/// Optional calendar metadata is compatible with legacy checkpoints.
#[serde(flatten)]
pub(crate) schedule: QuotaScheduleCheckpoint,
}
/// Serialized process owner for bounded durable quota checkpoints and resets.
@@ -29,6 +37,8 @@ pub(crate) struct QuotaStateOwner {
path: PathBuf,
store: Arc<QuotaStore>,
mutation: Arc<Mutex<()>>,
#[cfg(test)]
scheduler_iterations: std::sync::atomic::AtomicU64,
}
impl QuotaStateOwner {
@@ -36,8 +46,10 @@ impl QuotaStateOwner {
pub(crate) fn new(path: PathBuf, store: Arc<QuotaStore>) -> Arc<Self> {
Arc::new(Self {
path,
mutation: store.mutation_gate(),
store,
mutation: Arc::new(Mutex::new(())),
#[cfg(test)]
scheduler_iterations: std::sync::atomic::AtomicU64::new(0),
})
}
@@ -48,19 +60,46 @@ impl QuotaStateOwner {
/// Loads only quota entries owned by currently configured users.
pub(crate) async fn load(&self, configured_users: &BTreeSet<String>) {
let _ = self.load_inner(configured_users, false).await;
}
async fn load_inner(
&self,
configured_users: &BTreeSet<String>,
strict: bool,
) -> std::io::Result<()> {
let _guard = self.mutation.lock().await;
let state = match read_state_file(&self.path).await {
Ok(Some(state)) => state,
Ok(None) => return,
Ok(None) => return Ok(()),
Err(error) => {
if strict {
return Err(error);
}
warn!(
error = %error,
path = %self.path.display(),
"Failed to load quota state file"
);
return;
return Ok(());
}
};
if strict {
if state.users.len() > QUOTA_STATE_MAX_USERS {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"quota state exceeds 65536 users",
));
}
for entry in state.users.values() {
if let Some(period) = entry.schedule.quota_period {
period.next(entry.schedule.period_cursor_epoch_secs)?;
} else if entry.schedule.period_cursor_epoch_secs > 0 {
crate::config::UserQuotaPeriod::Daily
.next(entry.schedule.period_cursor_epoch_secs)?;
}
}
}
let persisted_users = state.users.len();
let mut loaded_users = 0usize;
@@ -70,6 +109,7 @@ impl QuotaStateOwner {
}
self.store
.load(&user, quota.used_bytes, quota.last_reset_epoch_secs);
self.store.load_schedule(&user, quota.schedule);
loaded_users = loaded_users.saturating_add(1);
}
info!(
@@ -78,6 +118,7 @@ impl QuotaStateOwner {
skipped_users = persisted_users.saturating_sub(loaded_users),
"Loaded bounded per-user quota state"
);
Ok(())
}
/// Persists a checkpoint filtered to the active configured user set.
@@ -99,18 +140,37 @@ impl QuotaStateOwner {
user: &str,
) -> std::io::Result<UserQuotaSnapshot> {
let guard = Arc::clone(&self.mutation).lock_owned().await;
let last_reset_epoch_secs = now_epoch_secs();
self.reset_user_locked(configured_users, user, now_epoch_secs(), guard)
.await
}
async fn reset_user_locked(
&self,
configured_users: &BTreeSet<String>,
user: &str,
last_reset_epoch_secs: u64,
guard: tokio::sync::OwnedMutexGuard<()>,
) -> std::io::Result<UserQuotaSnapshot> {
let prospective = UserQuotaSnapshot {
used_bytes: 0,
last_reset_epoch_secs,
};
let reset_target = self.store.current_or_legacy_handle(user);
let state = self.state_for_users(configured_users, Some((user, prospective.clone())));
let mut state = self.state_for_users(configured_users, Some((user, prospective.clone())));
let schedule = reset_target.manual_schedule_checkpoint(last_reset_epoch_secs);
if let Some(entry) = state.users.get_mut(user) {
entry.schedule = schedule;
}
let store = Arc::clone(&self.store);
let user_owned = user.to_owned();
let path = self.path.clone();
let task = tokio::task::spawn_blocking(move || {
let _guard = guard;
write_state_file_blocking(&path, &state)?;
Ok(reset_target.reset(last_reset_epoch_secs))
let snapshot = reset_target.reset(last_reset_epoch_secs);
reset_target.publish_manual_schedule(schedule);
store.refresh_schedule(&user_owned);
Ok(snapshot)
});
wait_for_blocking_io(task).await
}
@@ -122,6 +182,16 @@ impl QuotaStateOwner {
user: &str,
) -> std::io::Result<()> {
let guard = Arc::clone(&self.mutation).lock_owned().await;
self.remove_user_locked(configured_users, user, guard).await
}
/// Persists deletion while retaining the already-owned lifecycle gate.
pub(crate) async fn remove_user_locked(
&self,
configured_users: &BTreeSet<String>,
user: &str,
guard: tokio::sync::OwnedMutexGuard<()>,
) -> std::io::Result<()> {
debug_assert!(!configured_users.contains(user));
let state = self.state_for_users(configured_users, None);
let path = self.path.clone();
@@ -140,17 +210,53 @@ impl QuotaStateOwner {
let snapshot = self.store.snapshot();
let mut users = BTreeMap::new();
let mut last_reset_epoch_secs = 0;
for user in configured_users.iter().take(QUOTA_STATE_MAX_USERS) {
// Calendar watermarks and an explicit reset are mandatory. Legacy rows fill
// the remaining bounded space; no scheduled cursor can disappear by sorting.
let mut selected = configured_users
.iter()
.filter(|user| {
override_user
.as_ref()
.is_some_and(|(name, _)| name == &user.as_str())
|| self.store.current_handle(user).is_some_and(|target| {
target.schedule_checkpoint() != QuotaScheduleCheckpoint::default()
})
})
.cloned()
.collect::<BTreeSet<_>>();
for user in configured_users {
if selected.len() >= QUOTA_STATE_MAX_USERS {
break;
}
selected.insert(user.clone());
}
for user in &selected {
let quota = override_user
.as_ref()
.filter(|(override_name, _)| *override_name == user.as_str())
.map(|(_, quota)| quota.clone())
.or_else(|| snapshot.get(user).cloned());
let schedule = self
.store
.current_handle(user)
.map(|handle| handle.schedule_checkpoint())
.unwrap_or_default();
let quota = quota.or_else(|| {
(schedule != QuotaScheduleCheckpoint::default())
.then(|| {
self.store
.current_handle(user)
.map(|handle| handle.snapshot())
})
.flatten()
});
let Some(quota) = quota else {
continue;
};
last_reset_epoch_secs = last_reset_epoch_secs.max(quota.last_reset_epoch_secs);
users.insert(user.clone(), quota_user_state(quota));
let mut entry = quota_user_state(quota);
entry.schedule = schedule;
users.insert(user.clone(), entry);
}
QuotaStateFile {
last_reset_epoch_secs,
@@ -159,7 +265,8 @@ impl QuotaStateOwner {
}
}
fn now_epoch_secs() -> u64 {
/// Reads wall time only for cold quota administration.
pub(crate) fn now_epoch_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
@@ -226,6 +333,14 @@ async fn wait_for_blocking_io<T>(
}
fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Result<()> {
#[cfg(test)]
test_io::before(path);
if state.users.len() > QUOTA_STATE_MAX_USERS {
return Err(std::io::Error::new(
std::io::ErrorKind::InvalidData,
"quota state exceeds 65536 mandatory users",
));
}
let mut payload = serde_json::to_vec_pretty(state)?;
payload.push(b'\n');
if payload.len() as u64 > QUOTA_STATE_MAX_BYTES {
@@ -237,7 +352,10 @@ fn write_state_file_blocking(path: &Path, state: &QuotaStateFile) -> std::io::Re
#[cfg(unix)]
{
return crate::util::secure_fs::atomic_replace(path, &payload, 0o600);
crate::util::secure_fs::atomic_replace(path, &payload, 0o600)?;
#[cfg(test)]
test_io::after(path)?;
return Ok(());
}
#[cfg(not(unix))]
{
@@ -295,6 +413,7 @@ fn quota_user_state(quota: UserQuotaSnapshot) -> QuotaUserState {
QuotaUserState {
used_bytes: quota.used_bytes,
last_reset_epoch_secs: quota.last_reset_epoch_secs,
schedule: QuotaScheduleCheckpoint::default(),
}
}
@@ -400,6 +519,7 @@ mod tests {
(
"alice".to_string(),
QuotaUserState {
schedule: QuotaScheduleCheckpoint::default(),
used_bytes: 512,
last_reset_epoch_secs: 1,
},
@@ -407,6 +527,7 @@ mod tests {
(
"orphan".to_string(),
QuotaUserState {
schedule: QuotaScheduleCheckpoint::default(),
used_bytes: 1024,
last_reset_epoch_secs: 1,
},
@@ -424,3 +545,6 @@ mod tests {
assert!(!store.snapshot().contains_key("orphan"));
}
}
#[cfg(test)]
mod period_tests;
+297
View File
@@ -0,0 +1,297 @@
use super::*;
use crate::config::{ProxyConfig, UserQuotaPeriod};
use crate::proxy::user_admission::{UserAdmissionAuthority, credential_id_from_hex};
const SECRET: &str = "0123456789abcdef0123456789abcdef";
const DAY: u64 = 86400;
const START: u64 = 1_735_689_600;
fn config(period: Option<UserQuotaPeriod>) -> Arc<ProxyConfig> {
let mut config = ProxyConfig::default();
config.access.users.clear();
config.access.users.insert("alice".into(), SECRET.into());
config.access.user_data_quota.insert("alice".into(), 1024);
if let Some(period) = period {
config
.access
.user_quota_periods
.insert("alice".into(), period);
}
Arc::new(config)
}
fn request(store: &QuotaStore, config: &Arc<ProxyConfig>, now: u64) {
let target = store.current_handle("alice").unwrap();
store.request_schedule(
"alice",
target.incarnation(),
credential_id_from_hex(&config.access.users["alice"]).unwrap(),
config,
now,
true,
);
}
fn fixture() -> (tempfile::TempDir, Arc<QuotaStore>, Arc<QuotaStateOwner>) {
let dir = tempfile::tempdir().unwrap();
let store = Arc::new(QuotaStore::default());
store.activate_fresh("alice", 1);
let owner = QuotaStateOwner::new(dir.path().join("quota.json"), store.clone());
(dir, store, owner)
}
#[tokio::test]
async fn quota_period_activation_preserves_usage_and_boundaries_are_idempotent() {
let (_dir, store, owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Daily));
store.current_handle("alice").unwrap().charge(600);
request(&store, &cfg, START + 100);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), START)
.quota_schedule_state,
"pending"
);
owner.process_schedules(START + 100, false).await.unwrap();
assert_eq!(store.used("alice"), 600);
assert_eq!(store.schedule_deadline(), Some(START + DAY));
owner
.process_schedules(START + DAY - 1, false)
.await
.unwrap();
assert_eq!(store.used("alice"), 600);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(store.used("alice"), 0);
store.current_handle("alice").unwrap().charge(17);
for now in [START, START + DAY, START + DAY + 10] {
owner.process_schedules(now, false).await.unwrap();
assert_eq!(store.used("alice"), 17);
}
owner
.process_schedules(START + 50 * DAY + 1, false)
.await
.unwrap();
assert_eq!(store.used("alice"), 0);
assert_eq!(store.schedule_deadline(), Some(START + 51 * DAY));
let state = read_state_file(owner.path()).await.unwrap().unwrap();
assert_eq!(
state.users["alice"].schedule.period_cursor_epoch_secs,
START + 50 * DAY
);
}
#[tokio::test]
async fn quota_period_change_remove_reenable_and_rotation_preserve_watermark() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START + 100);
owner.process_schedules(START + 100, false).await.unwrap();
owner
.process_schedules(START + 2 * DAY, false)
.await
.unwrap();
store.current_handle("alice").unwrap().charge(600);
let before = store.current_handle("alice").unwrap().schedule_checkpoint();
store.advance_preserving_usage("alice", 2);
assert_eq!(
store.current_handle("alice").unwrap().schedule_checkpoint(),
before
);
for cfg in [
config(None),
config(Some(UserQuotaPeriod::Weekly)),
config(Some(UserQuotaPeriod::Daily)),
] {
request(&store, &cfg, START);
owner.process_schedules(START, false).await.unwrap();
assert_eq!(store.used("alice"), 600);
assert!(
store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.period_cursor_epoch_secs
>= START + 2 * DAY
);
}
assert_eq!(store.schedule_deadline(), Some(START + 3 * DAY));
}
#[tokio::test]
async fn quota_period_failure_keeps_usage_and_cursor_until_retry() {
let dir = tempfile::tempdir().unwrap();
let store = Arc::new(QuotaStore::default());
store.activate_fresh("alice", 1);
let path = dir.path().join("blocked");
std::fs::create_dir(&path).unwrap();
let owner = QuotaStateOwner::new(path.clone(), store.clone());
store.current_handle("alice").unwrap().charge(500);
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
assert!(owner.process_schedules(START, false).await.is_err());
assert_eq!(store.used("alice"), 500);
assert_eq!(
store.current_handle("alice").unwrap().schedule_checkpoint(),
QuotaScheduleCheckpoint::default()
);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), START)
.quota_schedule_state,
"error"
);
std::fs::remove_dir(path).unwrap();
owner.process_schedules(START, false).await.unwrap();
assert_eq!(store.used("alice"), 500);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(store.used("alice"), 0);
}
#[tokio::test]
async fn quota_period_old_reservations_cannot_refund_the_new_period() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
owner.process_schedules(START, false).await.unwrap();
let target = store.current_handle("alice").unwrap();
let reservation = target.try_reserve(700, 1024).unwrap();
owner.process_schedules(START + DAY, false).await.unwrap();
target.charge(200);
drop(reservation);
assert_eq!(target.used(), 200);
}
#[tokio::test]
async fn quota_period_restart_catches_up_once_and_legacy_migration_preserves_usage() {
let (dir, store, owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Monthly));
std::fs::write(owner.path(), r#"{"last_reset_epoch_secs":0,"users":{"alice":{"used_bytes":712,"last_reset_epoch_secs":0}}}"#).unwrap();
owner.load_for_config(&cfg).await.unwrap();
request(&store, &cfg, START + 200);
owner.process_schedules(START + 200, true).await.unwrap();
assert_eq!(store.used("alice"), 712);
let restarted = Arc::new(QuotaStore::default());
let second = QuotaStateOwner::new(dir.path().join("quota.json"), restarted.clone());
second.load_for_config(&cfg).await.unwrap();
restarted.activate_fresh("alice", 1);
request(&restarted, &cfg, START + 70 * DAY);
second
.process_schedules(START + 70 * DAY, true)
.await
.unwrap();
assert_eq!(restarted.used("alice"), 0);
restarted.current_handle("alice").unwrap().charge(55);
second
.process_schedules(START + 70 * DAY, false)
.await
.unwrap();
assert_eq!(restarted.used("alice"), 55);
}
#[tokio::test]
async fn quota_period_strict_startup_rejects_corruption_but_legacy_policy_does_not() {
let (_dir, _store, owner) = fixture();
std::fs::write(owner.path(), b"{broken").unwrap();
assert!(
owner
.load_for_config(&config(Some(UserQuotaPeriod::Daily)))
.await
.is_err()
);
assert!(owner.load_for_config(&config(None)).await.is_ok());
std::fs::write(owner.path(), r#"{"last_reset_epoch_secs":0,"users":{"alice":{"used_bytes":7,"last_reset_epoch_secs":0,"quota_period":"daily","period_cursor_epoch_secs":18446744073709551615}}}"#).unwrap();
assert!(
owner
.load_for_config(&config(Some(UserQuotaPeriod::Daily)))
.await
.is_err()
);
}
#[tokio::test]
async fn quota_period_fresh_account_rejects_stale_snapshot_and_does_not_inherit_cursor() {
let (_dir, store, owner) = fixture();
let old = config(Some(UserQuotaPeriod::Daily));
request(&store, &old, START);
owner.process_schedules(START, false).await.unwrap();
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
let credential = credential_id_from_hex(SECRET).unwrap();
store.fence_fresh_schedule(
"alice",
&old,
credential,
Some(UserQuotaPeriod::Monthly),
Some(1024),
);
request(&store, &old, START);
assert_eq!(
store.current_handle("alice").unwrap().schedule_checkpoint(),
QuotaScheduleCheckpoint::default()
);
assert_eq!(store.schedule_deadline(), None);
request(&store, &config(Some(UserQuotaPeriod::Monthly)), START + 1);
owner.process_schedules(START + 1, false).await.unwrap();
assert_eq!(store.schedule_deadline(), Some(START + 31 * DAY));
}
#[tokio::test]
async fn quota_period_manual_reset_advances_cursor_without_changing_cadence() {
let (_dir, store, owner) = fixture();
let now = now_epoch_secs();
let cfg = config(Some(UserQuotaPeriod::Daily));
request(&store, &cfg, now - 3 * DAY);
owner.process_schedules(now - 3 * DAY, false).await.unwrap();
store.current_handle("alice").unwrap().charge(250);
owner
.reset_user(&BTreeSet::from(["alice".into()]), "alice")
.await
.unwrap();
let next = store.schedule_deadline().unwrap();
assert_eq!(next, UserQuotaPeriod::Daily.next(now_epoch_secs()).unwrap());
store.current_handle("alice").unwrap().charge(80);
owner
.process_schedules(now_epoch_secs(), false)
.await
.unwrap();
assert_eq!(store.used("alice"), 80);
}
#[test]
fn quota_period_gate_is_shared_with_the_admission_authority() {
let (_dir, store, owner) = fixture();
let admission = UserAdmissionAuthority::new_with_quota_store(store.clone());
assert!(Arc::ptr_eq(
&owner.mutation,
&admission.quota_mutation_gate()
));
assert!(Arc::ptr_eq(&owner.mutation, &store.mutation_gate()));
}
#[tokio::test]
async fn quota_period_disabled_expired_account_keeps_calendar_without_becoming_enabled() {
let (_dir, store, owner) = fixture();
let authority = UserAdmissionAuthority::new_with_quota_store(store.clone());
let mut cfg = (*config(Some(UserQuotaPeriod::Daily))).clone();
cfg.access.user_enabled.insert("alice".into(), false);
cfg.access.user_expirations.insert(
"alice".into(),
chrono::DateTime::from_timestamp(0, 0).unwrap(),
);
let cfg = Arc::new(cfg);
authority.activate_config_source(1, None, &cfg.access.users, &cfg.access.user_enabled);
let guard = owner.lock().await;
authority.publish_quota_config(1, &cfg, true);
drop(guard);
owner
.process_schedules(now_epoch_secs(), false)
.await
.unwrap();
store.current_handle("alice").unwrap().charge(700);
let boundary = store.schedule_deadline().unwrap();
owner.process_schedules(boundary, false).await.unwrap();
assert_eq!(store.used("alice"), 0);
assert!(!authority.is_user_enabled("alice"));
assert!(cfg.access.user_expirations["alice"].timestamp() < now_epoch_secs() as i64);
}
mod concurrency;
mod resource_model;
+232
View File
@@ -0,0 +1,232 @@
use super::*;
use std::sync::atomic::Ordering;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn quota_period_metadata_commit_preserves_traffic_and_finishes_after_waiter_abort() {
let (_dir, store, owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Daily));
request(&store, &cfg, START);
store.current_handle("alice").unwrap().charge(100);
let (hook, started) = test_io::Hook::install(owner.path(), true);
let operation_owner = owner.clone();
let waiter = tokio::spawn(async move { operation_owner.process_schedules(START, false).await });
started.await.unwrap();
assert!(owner.mutation.try_lock().is_err());
store.current_handle("alice").unwrap().charge(77);
waiter.abort();
assert!(waiter.await.unwrap_err().is_cancelled());
assert!(owner.mutation.try_lock().is_err());
hook.release();
let _guard = tokio::time::timeout(std::time::Duration::from_secs(2), owner.lock())
.await
.unwrap();
assert_eq!(store.used("alice"), 177);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), START)
.quota_schedule_state,
"active"
);
assert_eq!(hook.writes.load(Ordering::Relaxed), 1);
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn quota_period_durable_reset_fences_recreation_until_commit() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
owner.process_schedules(START, false).await.unwrap();
let old = store.current_handle("alice").unwrap();
old.charge(900);
let (hook, started) = test_io::Hook::install(owner.path(), true);
let reset_owner = owner.clone();
let reset =
tokio::spawn(async move { reset_owner.process_schedules(START + DAY, false).await });
started.await.unwrap();
let mutate_store = store.clone();
let mutation = tokio::spawn(async move {
let _guard = mutate_store.mutation_gate().lock_owned().await;
mutate_store.retire_through("alice", 2);
mutate_store.activate_fresh("alice", 3);
mutate_store.current_handle("alice").unwrap().charge(400);
});
tokio::task::yield_now().await;
assert!(!mutation.is_finished());
assert_eq!(old.used(), 900);
hook.release();
reset.await.unwrap().unwrap();
mutation.await.unwrap();
assert_eq!(old.used(), 0);
assert_eq!(store.used("alice"), 400);
assert_eq!(
store.current_handle("alice").unwrap().schedule_checkpoint(),
QuotaScheduleCheckpoint::default()
);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(store.used("alice"), 400);
let persisted = read_state_file(owner.path()).await.unwrap().unwrap();
assert_eq!(persisted.users["alice"].used_bytes, 400);
assert_eq!(
persisted.users["alice"].schedule,
QuotaScheduleCheckpoint::default()
);
}
#[tokio::test]
async fn quota_period_restart_after_ambiguous_commit_never_replays_completed_boundary() {
let (dir, store, owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Daily));
request(&store, &cfg, START);
owner.process_schedules(START, false).await.unwrap();
store.current_handle("alice").unwrap().charge(900);
let (hook, _) = test_io::Hook::install(owner.path(), false);
hook.fail_after.store(true, Ordering::Relaxed);
assert!(owner.process_schedules(START + DAY, false).await.is_err());
assert_eq!(store.used("alice"), 900);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), START + DAY)
.quota_schedule_state,
"error"
);
drop(hook);
let restarted = Arc::new(QuotaStore::default());
let restart = QuotaStateOwner::new(dir.path().join("quota.json"), restarted.clone());
restart.load_for_config(&cfg).await.unwrap();
restarted.activate_fresh("alice", 1);
request(&restarted, &cfg, START + DAY);
restarted.current_handle("alice").unwrap().charge(77);
restart.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(restarted.used("alice"), 77);
}
#[tokio::test]
async fn quota_period_unchanged_reload_is_io_free_and_pending_reversal_is_not_a_reset() {
let (_dir, store, owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Daily));
request(&store, &cfg, START);
owner.process_schedules(START, false).await.unwrap();
let (hook, _) = test_io::Hook::install(owner.path(), false);
for _ in 0..1000 {
request(&store, &cfg, START + 500);
owner.process_schedules(START + 500, false).await.unwrap();
}
request(&store, &config(Some(UserQuotaPeriod::Weekly)), START + 500);
request(&store, &cfg, START + 500);
owner.process_schedules(START + 500, false).await.unwrap();
assert_eq!(hook.writes.load(Ordering::Relaxed), 0);
assert_eq!(store.schedule_deadline(), Some(START + DAY));
}
#[tokio::test]
async fn quota_period_manual_recovery_clears_error_without_postponing_calendar() {
let (_dir, store, owner) = fixture();
let now = now_epoch_secs();
request(&store, &config(Some(UserQuotaPeriod::Daily)), now - 3 * DAY);
owner.process_schedules(now - 3 * DAY, false).await.unwrap();
let (hook, _) = test_io::Hook::install(owner.path(), false);
hook.fail_after.store(true, Ordering::Relaxed);
assert!(owner.process_schedules(now, false).await.is_err());
hook.fail_after.store(false, Ordering::Relaxed);
owner
.reset_user(&BTreeSet::from(["alice".into()]), "alice")
.await
.unwrap();
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), now)
.quota_schedule_state,
"active"
);
assert!(!store.schedule_ready(now));
}
#[tokio::test]
async fn quota_period_stale_candidate_cannot_attach_retired_policy_to_fresh_account() {
let (_dir, store, owner) = fixture();
let old = config(Some(UserQuotaPeriod::Daily));
let stale_candidate = Arc::new((*old).clone());
let credential = credential_id_from_hex(SECRET).unwrap();
store.fence_fresh_schedule(
"alice",
&old,
credential,
Some(UserQuotaPeriod::Monthly),
Some(1024),
);
store.request_schedule("alice", 1, credential, &stale_candidate, START, false);
assert_eq!(store.schedule_deadline(), None);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Monthly), START)
.quota_schedule_state,
"pending"
);
request(&store, &config(Some(UserQuotaPeriod::Monthly)), START);
owner.process_schedules(START, false).await.unwrap();
assert_eq!(store.schedule_deadline(), Some(START + 31 * DAY));
}
#[tokio::test]
async fn quota_period_worker_drops_obsolete_retry_and_shuts_down_cooperatively() {
let (_dir, store, owner) = fixture();
let now = now_epoch_secs();
request(&store, &config(Some(UserQuotaPeriod::Daily)), now - 3 * DAY);
owner.process_schedules(now - 3 * DAY, false).await.unwrap();
let (hook, started) = test_io::Hook::install(owner.path(), false);
hook.fail_after.store(true, Ordering::Relaxed);
let control = crate::maestro::control_plane::ProcessControlPlane::new();
owner.spawn_scheduler(&control);
started.await.unwrap();
let guard = owner.lock().await;
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), now)
.quota_schedule_state,
"error"
);
hook.fail_after.store(false, Ordering::Relaxed);
owner
.reset_user_locked(&BTreeSet::from(["alice".into()]), "alice", now, guard)
.await
.unwrap();
tokio::time::pause();
let before = owner.scheduler_iterations.load(Ordering::Relaxed);
tokio::time::advance(std::time::Duration::from_secs(65)).await;
for _ in 0..100 {
tokio::task::yield_now().await;
}
let iterations = owner.scheduler_iterations.load(Ordering::Relaxed) - before;
assert!(iterations < 20, "idle scheduler spun {iterations} times");
assert!(!store.schedule_ready(now));
assert!(control.shutdown(std::time::Duration::from_secs(2)).await);
tokio::time::resume();
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn quota_period_manual_commit_crossing_midnight_handles_new_boundary_once() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
owner.process_schedules(START, false).await.unwrap();
let (hook, started) = test_io::Hook::install(owner.path(), true);
let manual_owner = owner.clone();
let manual = tokio::spawn(async move {
let guard = manual_owner.lock().await;
manual_owner
.reset_user_locked(
&BTreeSet::from(["alice".into()]),
"alice",
START + DAY - 1,
guard,
)
.await
});
started.await.unwrap();
hook.release();
manual.await.unwrap().unwrap();
store.current_handle("alice").unwrap().charge(13);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(store.used("alice"), 0);
store.current_handle("alice").unwrap().charge(19);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(store.used("alice"), 19);
}
@@ -0,0 +1,140 @@
use super::*;
use std::sync::atomic::Ordering;
#[tokio::test]
async fn quota_period_retained_cursor_survives_legacy_inventory_growth() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
owner.process_schedules(START, false).await.unwrap();
let mut users = (0..QUOTA_STATE_MAX_USERS)
.map(|i| format!("a{i:05}"))
.collect::<BTreeSet<_>>();
for user in &users {
store.current_or_legacy_handle(user).charge(1);
}
users.insert("alice".into());
request(&store, &config(None), START + 1);
owner.process_schedules(START + 1, false).await.unwrap();
owner.save(&users).await.unwrap();
let state = read_state_file(owner.path()).await.unwrap().unwrap();
assert_eq!(state.users.len(), QUOTA_STATE_MAX_USERS);
assert_eq!(state.users["alice"].schedule.quota_period, None);
assert_eq!(
state.users["alice"].schedule.period_cursor_epoch_secs,
START + 1
);
assert!(!state.users.contains_key("a65535"));
}
#[tokio::test]
async fn quota_period_oversized_checkpoint_cannot_publish_reset() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), START);
owner.process_schedules(START, false).await.unwrap();
let before = std::fs::read(owner.path()).unwrap();
let huge_name = "x".repeat(QUOTA_STATE_MAX_BYTES as usize);
store.current_or_legacy_handle(&huge_name).charge(1);
store.current_handle("alice").unwrap().charge(77);
assert!(owner.process_schedules(START + DAY, false).await.is_err());
assert_eq!(store.used("alice"), 77);
assert_eq!(std::fs::read(owner.path()).unwrap(), before);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), START + DAY)
.quota_schedule_state,
"error"
);
}
#[tokio::test]
async fn quota_period_many_due_users_share_one_durable_batch() {
let (_dir, store, owner) = fixture();
let mut cfg = (*config(Some(UserQuotaPeriod::Daily))).clone();
for i in 0..128 {
let user = format!("user{i}");
store.activate_fresh(&user, 1);
store.current_handle(&user).unwrap().charge(123);
cfg.access.users.insert(user.clone(), SECRET.into());
cfg.access.user_data_quota.insert(user.clone(), 1024);
cfg.access
.user_quota_periods
.insert(user, UserQuotaPeriod::Daily);
}
let cfg = Arc::new(cfg);
for user in cfg.access.users.keys() {
store.request_schedule(
user,
1,
credential_id_from_hex(SECRET).unwrap(),
&cfg,
START,
true,
);
}
let (hook, _) = super::super::test_io::Hook::install(owner.path(), false);
owner.process_schedules(START, false).await.unwrap();
assert_eq!(hook.writes.load(Ordering::Relaxed), 1);
owner.process_schedules(START + DAY, false).await.unwrap();
assert_eq!(hook.writes.load(Ordering::Relaxed), 2);
for user in cfg.access.users.keys() {
assert_eq!(store.used(user), 0);
}
assert_eq!(store.schedule_deadline(), Some(START + 2 * DAY));
}
#[test]
fn quota_period_daily_model_preserves_monotonic_cursor_under_clock_pressure() {
let (_dir, store, _owner) = fixture();
let cfg = config(Some(UserQuotaPeriod::Daily));
request(&store, &cfg, START + 123);
for change in store.schedule_changes(START + 123).unwrap() {
store.publish_schedule(&change, START + 123);
}
let mut cursor = START + 123;
let mut used = 0;
let mut seed = 19u64;
for step in 0..10_000 {
seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
let now = START + (seed % 90) * DAY + seed % DAY;
let amount = seed % 100;
store.current_handle("alice").unwrap().charge(amount);
used += amount;
let boundary = now / DAY * DAY;
if boundary > cursor {
cursor = boundary;
used = 0;
}
for change in store.schedule_changes(now).unwrap() {
store.publish_schedule(&change, now);
}
if step % 100 == 0 {
store.advance_preserving_usage("alice", step + 2);
}
assert_eq!(store.used("alice"), used);
assert_eq!(
store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.period_cursor_epoch_secs,
cursor
);
assert_eq!(store.schedule_deadline(), Some((cursor / DAY + 1) * DAY));
}
}
#[tokio::test]
async fn quota_period_invalid_clock_reports_error_without_publishing_metadata() {
let (_dir, store, owner) = fixture();
request(&store, &config(Some(UserQuotaPeriod::Daily)), u64::MAX);
store.current_handle("alice").unwrap().charge(5);
assert!(owner.process_schedules(u64::MAX, false).await.is_err());
assert_eq!(store.used("alice"), 5);
assert_eq!(
store
.schedule_status("alice", Some(UserQuotaPeriod::Daily), u64::MAX)
.quota_schedule_state,
"error"
);
assert!(!owner.path().exists());
}
+174
View File
@@ -0,0 +1,174 @@
use std::sync::Arc;
use std::time::Duration;
use tracing::{info, warn};
use super::{QuotaStateOwner, now_epoch_secs, wait_for_blocking_io, write_state_file_blocking};
use crate::config::ProxyConfig;
use crate::maestro::control_plane::ProcessControlPlane;
use crate::stats::{QuotaScheduleStatus, UserQuotaSnapshot};
impl QuotaStateOwner {
/// Exposes the same control-plane gate used by account lifecycle transitions.
pub(crate) async fn lock(&self) -> tokio::sync::OwnedMutexGuard<()> {
Arc::clone(&self.mutation).lock_owned().await
}
/// Loads legacy files tolerantly; scheduled policies require trustworthy state.
pub(crate) async fn load_for_config(&self, config: &ProxyConfig) -> std::io::Result<()> {
self.load_inner(
&config.access.users.keys().cloned().collect(),
!config.access.user_quota_periods.is_empty(),
)
.await
}
/// Fences first calendar attachment for an API-created account.
pub(crate) fn fence_fresh_schedule(
&self,
user: &str,
source: &Arc<ProxyConfig>,
credential: crate::proxy::user_admission::UserCredentialId,
period: Option<crate::config::UserQuotaPeriod>,
quota: Option<u64>,
) {
self.store
.fence_fresh_schedule(user, source, credential, period, quota);
}
/// Retains first-attachment fencing through API edits and secret rotation.
pub(crate) fn refresh_fresh_schedule(&self, user: &str, config: &ProxyConfig) {
self.store.refresh_fresh_schedule(user, config);
}
/// Captures API status under the caller's mutation guard.
pub(crate) fn schedule_status(
&self,
user: &str,
desired: Option<crate::config::UserQuotaPeriod>,
) -> QuotaScheduleStatus {
self.store.schedule_status(user, desired, now_epoch_secs())
}
/// Checkpoints a bounded batch before publishing any reset or calendar metadata.
pub(crate) async fn process_schedules(
&self,
now: u64,
force_save: bool,
) -> std::io::Result<()> {
let guard = self.lock().await;
self.process_schedules_locked(now, force_save, guard).await
}
async fn process_schedules_locked(
&self,
now: u64,
force_save: bool,
guard: tokio::sync::OwnedMutexGuard<()>,
) -> std::io::Result<()> {
if !force_save && !self.store.schedule_ready(now) {
return Ok(());
}
let changes = self.store.schedule_changes(now)?;
let configured_users = self.store.configured_quota_users();
let mut state = self.state_for_users(&configured_users, None);
for change in &changes {
let quota = if change.reset {
UserQuotaSnapshot {
used_bytes: 0,
last_reset_epoch_secs: now,
}
} else {
change.target.snapshot()
};
let mut entry = super::quota_user_state(quota);
entry.schedule = change.checkpoint;
state.users.insert(change.user.clone(), entry);
}
state.last_reset_epoch_secs = state
.users
.values()
.map(|entry| entry.last_reset_epoch_secs)
.max()
.unwrap_or(0);
let path = self.path.clone();
let store = Arc::clone(&self.store);
let task = tokio::task::spawn_blocking(move || {
let _guard = guard;
if let Err(error) = write_state_file_blocking(&path, &state) {
store.schedule_failed(&changes);
return Err(error);
}
let reset_count = changes.iter().filter(|change| change.reset).count();
for change in &changes {
store.publish_schedule(change, now);
}
store.schedule_saved();
if reset_count > 0 {
info!(users = reset_count, "Completed scheduled quota reset batch");
}
Ok(())
});
wait_for_blocking_io(task).await
}
/// Owns one worker across runtime generations; accepted commits finish on shutdown.
pub(crate) fn spawn_scheduler(self: &Arc<Self>, control: &ProcessControlPlane) {
let owner = Arc::clone(self);
let _ = control.spawn_cooperative(move |cancel| async move {
let mut retry = 1u64;
let mut retry_at = None;
let mut last_warning: Option<tokio::time::Instant> = None;
loop {
#[cfg(test)]
owner.scheduler_iterations.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
if cancel.is_cancelled() { break; }
let now = now_epoch_secs();
let can_retry = retry_at.is_none_or(|at| tokio::time::Instant::now() >= at);
if can_retry && owner.store.schedule_ready(now) {
// Cancellation can discard waiting work, but never an accepted commit.
let guard = tokio::select! {
biased;
_ = cancel.cancelled() => break,
guard = owner.lock() => guard,
};
if let Err(error) = owner.process_schedules_locked(now_epoch_secs(), false, guard).await {
let instant = tokio::time::Instant::now();
if last_warning.is_none_or(|at| instant.duration_since(at) >= Duration::from_secs(60)) {
warn!(error = %error, "Quota schedule checkpoint failed; usage remains unchanged");
last_warning = Some(instant);
}
retry_at = Some(instant + Duration::from_secs(retry));
retry = (retry * 2).min(60);
} else {
if retry_at.is_some() { info!("Quota schedule persistence recovered"); }
retry_at = None;
retry = 1;
last_warning = None;
}
continue;
}
// A manual reset or policy removal may resolve the failed work.
// An expired retry must never turn an idle worker into sleep(0) polling.
if can_retry && !owner.store.schedule_ready(now) {
retry_at = None;
retry = 1;
last_warning = None;
}
let delay = retry_at.map(|at| at.saturating_duration_since(tokio::time::Instant::now()))
.or_else(|| owner.store.schedule_deadline().map(|at| Duration::from_secs(at.saturating_sub(now).min(1))));
tokio::select! {
biased;
_ = cancel.cancelled() => break,
_ = owner.store.schedule_notified() => {},
_ = async {
match delay {
Some(duration) => tokio::time::sleep(duration).await,
None => std::future::pending::<()>().await,
}
} => {},
}
}
});
}
}
+97
View File
@@ -0,0 +1,97 @@
//! Path-scoped deterministic I/O barriers; absent from production builds.
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, Condvar, Mutex, OnceLock};
fn hooks() -> &'static Mutex<HashMap<PathBuf, Arc<Hook>>> {
static HOOKS: OnceLock<Mutex<HashMap<PathBuf, Arc<Hook>>>> = OnceLock::new();
HOOKS.get_or_init(Mutex::default)
}
pub(super) struct Hook {
path: PathBuf,
started: Mutex<Option<tokio::sync::oneshot::Sender<()>>>,
released: Mutex<bool>,
wake: Condvar,
pub(super) writes: AtomicUsize,
pub(super) fail_after: AtomicBool,
}
pub(super) struct HookGuard(Arc<Hook>);
impl std::ops::Deref for HookGuard {
type Target = Hook;
fn deref(&self) -> &Hook {
&self.0
}
}
impl Drop for HookGuard {
fn drop(&mut self) {
self.0.remove();
}
}
impl Hook {
pub(super) fn install(
path: &Path,
block: bool,
) -> (HookGuard, tokio::sync::oneshot::Receiver<()>) {
let (tx, rx) = tokio::sync::oneshot::channel();
let hook = Arc::new(Self {
path: path.into(),
started: Mutex::new(Some(tx)),
released: Mutex::new(!block),
wake: Condvar::new(),
writes: AtomicUsize::new(0),
fail_after: AtomicBool::new(false),
});
assert!(
hooks()
.lock()
.unwrap()
.insert(path.into(), hook.clone())
.is_none()
);
(HookGuard(hook), rx)
}
pub(super) fn release(&self) {
*self.released.lock().unwrap() = true;
self.wake.notify_all();
}
pub(super) fn remove(&self) {
self.release();
hooks().lock().unwrap().remove(&self.path);
}
}
pub(super) fn before(path: &Path) {
let hook = hooks().lock().unwrap().get(path).cloned();
if let Some(hook) = hook {
hook.writes.fetch_add(1, Ordering::Relaxed);
if let Some(tx) = hook.started.lock().unwrap().take() {
let _ = tx.send(());
}
let mut released = hook.released.lock().unwrap();
while !*released {
released = hook.wake.wait(released).unwrap();
}
}
}
pub(super) fn after(path: &Path) -> std::io::Result<()> {
if hooks()
.lock()
.unwrap()
.get(path)
.is_some_and(|hook| hook.fail_after.load(Ordering::Relaxed))
{
return Err(std::io::Error::other(
"injected failure after durable replacement",
));
}
Ok(())
}
+3 -1
View File
@@ -21,7 +21,9 @@ use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
use std::time::Instant;
pub(crate) use self::quota_store::{QuotaReservation, QuotaStore, UserQuotaHandle};
pub(crate) use self::quota_store::{
QuotaReservation, QuotaScheduleCheckpoint, QuotaScheduleStatus, QuotaStore, UserQuotaHandle,
};
#[allow(unused_imports)]
pub use self::replay::{ReplayChecker, ReplayStats};
use self::telemetry::TelemetryPolicy;
+42 -134
View File
@@ -9,10 +9,18 @@ use parking_lot::Mutex;
use super::{QuotaReserveError, UserQuotaSnapshot};
use crate::proxy::user_admission::UserIncarnation;
// Cold calendar metadata and its bounded work index never enter reserve/charge.
mod schedule;
use schedule::{QuotaSchedule, ScheduleWork};
pub(crate) use schedule::{QuotaScheduleCheckpoint, QuotaScheduleStatus};
/// Process-scoped per-user quota accounting shared by runtime generations.
#[derive(Default)]
pub struct QuotaStore {
users: DashMap<String, Arc<QuotaUserSlot>>,
mutation: Arc<tokio::sync::Mutex<()>>,
schedule_work: Mutex<ScheduleWork>,
schedule_notify: tokio::sync::Notify,
}
struct QuotaUserSlot {
@@ -41,6 +49,7 @@ pub(crate) struct UserQuotaHandle {
/// Atomically replaceable quota state for one configured user.
pub(crate) struct UserQuotaCounters {
generation: ArcSwap<QuotaGeneration>,
schedule: Mutex<QuotaSchedule>,
}
struct QuotaGeneration {
@@ -124,6 +133,15 @@ impl QuotaStore {
if incarnation <= state.high_water {
return;
}
let schedule = if state.high_water == 0 {
state
.current
.as_ref()
.map(|account| account.counters.schedule.lock().clone())
.unwrap_or_default()
} else {
QuotaSchedule::default()
};
let snapshot = if state.high_water == 0 {
state
.current
@@ -139,9 +157,11 @@ impl QuotaStore {
last_reset_epoch_secs: 0,
});
state.high_water = state.high_water.max(incarnation);
let counters = Arc::new(UserQuotaCounters::from_snapshot(&snapshot));
*counters.schedule.lock() = schedule;
state.current = Some(QuotaAccount {
incarnation,
counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)),
counters,
});
}
@@ -152,6 +172,11 @@ impl QuotaStore {
if incarnation <= state.high_water {
return;
}
let schedule = state
.current
.as_ref()
.map(|account| account.counters.schedule.lock().clone())
.unwrap_or_default();
let snapshot = state
.current
.as_ref()
@@ -162,9 +187,11 @@ impl QuotaStore {
last_reset_epoch_secs: 0,
});
state.high_water = incarnation;
let counters = Arc::new(UserQuotaCounters::from_snapshot(&snapshot));
*counters.schedule.lock() = schedule;
state.current = Some(QuotaAccount {
incarnation,
counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)),
counters,
});
}
@@ -176,8 +203,13 @@ impl QuotaStore {
return;
}
state.high_water = incarnation;
let retained_metadata = state.current.as_ref().is_some_and(|account| {
account.counters.schedule.lock().checkpoint != QuotaScheduleCheckpoint::default()
});
state.current = None;
state.startup_seed = None;
drop(state);
self.schedule_retired(user, retained_metadata);
}
pub(crate) fn used(&self, user: &str) -> u64 {
@@ -216,8 +248,13 @@ impl QuotaStore {
let slot = self.slot(user);
let mut state = slot.state.lock();
state.high_water = state.high_water.saturating_add(1);
let retained_metadata = state.current.as_ref().is_some_and(|account| {
account.counters.schedule.lock().checkpoint != QuotaScheduleCheckpoint::default()
});
state.current = None;
state.startup_seed = None;
drop(state);
self.schedule_retired(user, retained_metadata);
}
pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
@@ -243,6 +280,7 @@ impl QuotaStore {
impl Default for UserQuotaCounters {
fn default() -> Self {
Self {
schedule: Mutex::new(QuotaSchedule::default()),
generation: ArcSwap::from_pointee(QuotaGeneration {
used_bytes: AtomicU64::new(0),
last_reset_epoch_secs: 0,
@@ -254,6 +292,7 @@ impl Default for UserQuotaCounters {
impl UserQuotaCounters {
fn from_snapshot(snapshot: &UserQuotaSnapshot) -> Self {
Self {
schedule: Mutex::new(QuotaSchedule::default()),
generation: ArcSwap::from_pointee(QuotaGeneration {
used_bytes: AtomicU64::new(snapshot.used_bytes),
last_reset_epoch_secs: snapshot.last_reset_epoch_secs,
@@ -410,135 +449,4 @@ fn refund_generation(generation: &QuotaGeneration, bytes: u64) {
}
#[cfg(test)]
mod tests {
use super::*;
use crate::stats::Stats;
use std::sync::Barrier;
#[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);
}
#[test]
fn quota_snapshot_never_combines_different_generations() {
const ITERATIONS: u64 = 10_000;
let store = Arc::new(QuotaStore::default());
store.load("alice", 1, 1);
let barrier = Arc::new(Barrier::new(2));
let writer_store = Arc::clone(&store);
let writer_barrier = Arc::clone(&barrier);
let writer = std::thread::spawn(move || {
writer_barrier.wait();
for generation in 2..=ITERATIONS {
writer_store.load("alice", generation, generation);
}
});
barrier.wait();
for _ in 0..ITERATIONS {
let snapshot = store.snapshot().remove("alice").unwrap();
assert_eq!(snapshot.used_bytes, snapshot.last_reset_epoch_secs);
}
writer.join().unwrap();
}
#[test]
fn repeated_old_generation_refunds_leave_new_usage_intact() {
const ITERATIONS: u64 = 10_000;
let store = QuotaStore::default();
let state = store.user("alice");
for generation in 1..=ITERATIONS {
let reservation = state.try_reserve(80, 100).unwrap();
store.load("alice", 40, generation);
drop(reservation);
assert_eq!(store.used("alice"), 40);
store.reset("alice", generation);
}
}
#[test]
fn retired_incarnation_cannot_charge_recreated_username() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let retired = store.handle_exact("alice", 1).unwrap();
retired.charge(40);
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
retired.charge(20);
assert_eq!(retired.used(), 60);
assert_eq!(store.handle_exact("alice", 3).unwrap().used(), 0);
}
#[test]
fn credential_rotation_preserves_usage_without_sharing_future_charges() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let old = store.handle_exact("alice", 1).unwrap();
old.charge(40);
store.advance_preserving_usage("alice", 2);
let current = store.handle_exact("alice", 2).unwrap();
old.charge(20);
assert_eq!(old.used(), 60);
assert_eq!(current.used(), 40);
}
#[test]
fn captured_reset_handle_cannot_reset_a_new_incarnation() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let reset_target = store.handle_exact("alice", 1).unwrap();
reset_target.charge(40);
store.advance_preserving_usage("alice", 2);
let current = store.handle_exact("alice", 2).unwrap();
current.charge(20);
reset_target.reset(7);
assert_eq!(reset_target.used(), 0);
assert_eq!(current.used(), 60);
}
#[test]
fn stale_retirement_cannot_remove_newer_quota_owner() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
store.retire_through("alice", 2);
assert!(store.handle_exact("alice", 3).is_some());
}
#[test]
fn old_reservation_refund_does_not_debit_recreated_username() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let old = store.handle_exact("alice", 1).unwrap();
let reservation = old.try_reserve(80, 100).unwrap();
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
let current = store.handle_exact("alice", 3).unwrap();
current.charge(50);
drop(reservation);
assert_eq!(old.used(), 0);
assert_eq!(current.used(), 50);
}
}
mod tests;
+443
View File
@@ -0,0 +1,443 @@
use std::collections::{BTreeMap, BTreeSet};
use std::sync::{Arc, Weak};
use serde::{Deserialize, Serialize};
use super::{QuotaStore, UserQuotaHandle};
use crate::config::{ProxyConfig, UserQuotaPeriod};
use crate::proxy::user_admission::UserCredentialId;
/// Durable calendar watermark, independent from the last manual reset timestamp.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub(crate) struct QuotaScheduleCheckpoint {
/// Effective calendar; absent leaves the existing manual-only quota policy.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub(crate) quota_period: Option<UserQuotaPeriod>,
/// Monotonic processed time; activation and manual reset need not be midnight.
#[serde(default, skip_serializing_if = "is_zero")]
pub(crate) period_cursor_epoch_secs: u64,
}
fn is_zero(value: &u64) -> bool {
*value == 0
}
/// First policy attachment fence for an API-created account.
#[derive(Clone)]
pub(super) struct FreshSchedule {
/// Snapshot visible at creation; preparation can publish another older snapshot later.
pub(super) source: Weak<ProxyConfig>,
/// Monotonic commit time, independent from quota calendar wall time.
pub(super) created_at: std::time::Instant,
/// Latest committed credential until first policy attachment.
pub(super) credential: UserCredentialId,
/// Latest committed desired calendar.
pub(super) period: Option<UserQuotaPeriod>,
/// Latest committed quota amount.
pub(super) quota: Option<u64>,
}
/// Account-owned cold state serialized by the process quota mutation gate.
#[derive(Clone, Default)]
pub(super) struct QuotaSchedule {
/// Last successful durable calendar publication.
pub(super) checkpoint: QuotaScheduleCheckpoint,
/// Calendar accepted from the active runtime configuration.
pub(super) requested: Option<UserQuotaPeriod>,
/// Wall time captured when the requested calendar changed.
pub(super) accepted_at: u64,
/// Suspends the previous cadence while a policy change awaits persistence.
pub(super) pending: bool,
/// Reports the last unsuccessful attempt without modifying the durable cursor.
pub(super) failed: bool,
/// Rejects pre-creation policies until a matching or newer source arrives.
pub(super) fresh: Option<FreshSchedule>,
}
/// Bounded replacement index: a username owns at most one deadline and dirty entry.
#[derive(Default)]
pub(super) struct ScheduleWork {
due: BTreeSet<(u64, String)>,
by_user: BTreeMap<String, u64>,
dirty: BTreeSet<String>,
removed: bool,
}
/// Consistent control-plane status; no calendar work runs on the relay path.
#[derive(Debug, Clone, Default, Serialize)]
pub(crate) struct QuotaScheduleStatus {
/// Effective calendar, absent while a policy change is awaiting its checkpoint.
pub(crate) active_quota_period: Option<UserQuotaPeriod>,
/// UTC boundary containing the durable cursor.
pub(crate) quota_period_start_epoch_secs: Option<u64>,
/// First boundary strictly after the durable cursor.
pub(crate) quota_next_reset_epoch_secs: Option<u64>,
/// Bounded API state token; never includes internal error details.
pub(crate) quota_schedule_state: &'static str,
}
/// A captured account and policy transition owned by the process mutation gate.
pub(crate) struct QuotaScheduleChange {
/// Current account name used to replace one bounded work-index entry.
pub(crate) user: String,
/// Exact accounting incarnation captured under the process gate.
pub(crate) target: UserQuotaHandle,
/// Prospective metadata written before live publication.
pub(crate) checkpoint: QuotaScheduleCheckpoint,
/// Distinguishes a due reset from metadata-only activation.
pub(crate) reset: bool,
}
impl QuotaStore {
/// Shares the control-plane gate across API, scheduler, and runtime sources.
pub(crate) fn mutation_gate(&self) -> Arc<tokio::sync::Mutex<()>> {
Arc::clone(&self.mutation)
}
/// Waits for an account or policy notification; its payload is never authoritative.
pub(crate) async fn schedule_notified(&self) {
self.schedule_notify.notified().await;
}
/// Detects accounts still awaiting first policy attachment on a cold reload.
pub(crate) fn has_fresh_schedules(&self) -> bool {
self.users.iter().any(|entry| {
entry
.value()
.state
.lock()
.current
.as_ref()
.is_some_and(|account| account.counters.schedule.lock().fresh.is_some())
})
}
/// Returns whether reconciliation or a due boundary requires work.
pub(crate) fn schedule_ready(&self, now: u64) -> bool {
let work = self.schedule_work.lock();
work.removed || !work.dirty.is_empty() || work.due.first().is_some_and(|(at, _)| *at <= now)
}
/// Returns the earliest deadline without scanning accounts.
pub(crate) fn schedule_deadline(&self) -> Option<u64> {
self.schedule_work.lock().due.first().map(|(at, _)| *at)
}
fn schedule_index(&self, user: &str, target: Option<&UserQuotaHandle>) {
let schedule = target.map(|handle| handle.counters.schedule.lock().clone());
let mut work = self.schedule_work.lock();
if let Some(old) = work.by_user.remove(user) {
work.due.remove(&(old, user.to_owned()));
}
work.dirty.remove(user);
if let Some(state) = schedule {
if state.pending {
work.dirty.insert(user.to_owned());
} else if let Some(period) = state.checkpoint.quota_period {
match period.next(state.checkpoint.period_cursor_epoch_secs) {
Ok(at) => {
work.by_user.insert(user.to_owned(), at);
work.due.insert((at, user.to_owned()));
}
Err(_) => {
work.dirty.insert(user.to_owned());
}
}
}
}
drop(work);
self.schedule_notify.notify_one();
}
/// Discards retired work; the next checkpoint removes its durable entry.
pub(super) fn schedule_retired(&self, user: &str, retained_metadata: bool) {
let had_schedule = {
let work = self.schedule_work.lock();
work.by_user.contains_key(user) || work.dirty.contains(user)
};
if had_schedule || retained_metadata {
self.schedule_work.lock().removed = true;
}
self.schedule_index(user, None);
}
/// Attaches legacy/startup metadata without changing the loaded byte counter.
pub(crate) fn load_schedule(&self, user: &str, checkpoint: QuotaScheduleCheckpoint) {
let target = self.current_or_legacy_handle(user);
*target.counters.schedule.lock() = QuotaSchedule {
checkpoint,
requested: checkpoint.quota_period,
..QuotaSchedule::default()
};
}
/// Records the snapshot fence for a newly created API account.
pub(crate) fn fence_fresh_schedule(
&self,
user: &str,
source: &Arc<ProxyConfig>,
credential: UserCredentialId,
period: Option<UserQuotaPeriod>,
quota: Option<u64>,
) {
let target = self.current_or_legacy_handle(user);
target.counters.schedule.lock().fresh = Some(FreshSchedule {
source: Arc::downgrade(source),
created_at: std::time::Instant::now(),
credential,
period,
quota,
});
}
/// Updates an unconsumed creation fence after another committed API mutation.
pub(crate) fn refresh_fresh_schedule(&self, user: &str, config: &ProxyConfig) {
let Some(target) = self.current_handle(user) else {
return;
};
let mut schedule = target.counters.schedule.lock();
let Some(fresh) = schedule.fresh.as_mut() else {
return;
};
if let Some(credential) = config
.access
.users
.get(user)
.and_then(|secret| crate::proxy::user_admission::credential_id_from_hex(secret))
{
fresh.credential = credential;
fresh.period = config.access.user_quota_periods.get(user).copied();
fresh.quota = config.access.user_data_quota.get(user).copied();
fresh.created_at = std::time::Instant::now();
}
}
/// Accepts policy only from a current credential and an eligible config snapshot.
pub(crate) fn request_schedule(
&self,
user: &str,
incarnation: u64,
credential: UserCredentialId,
config: &Arc<ProxyConfig>,
now: u64,
accepted_after_mutation: bool,
) {
let Some(target) = self.handle_exact(user, incarnation) else {
return;
};
let period = config.access.user_quota_periods.get(user).copied();
let mut schedule = target.counters.schedule.lock();
if let Some(fresh) = &schedule.fresh {
let same_source = fresh.source.ptr_eq(&Arc::downgrade(config));
let read_after_creation = config
.quota_source_read_at
.is_some_and(|at| at > fresh.created_at);
if (same_source || !accepted_after_mutation || !read_after_creation)
&& (credential != fresh.credential
|| period != fresh.period
|| config.access.user_data_quota.get(user).copied() != fresh.quota)
{
return;
}
schedule.fresh = None;
}
if schedule.requested != period {
schedule.requested = period;
schedule.accepted_at = now;
schedule.pending = period != schedule.checkpoint.quota_period;
schedule.failed = false;
}
drop(schedule);
self.schedule_index(user, Some(&target));
}
/// Captures only dirty or due accounts after the caller owns the mutation gate.
pub(crate) fn schedule_changes(&self, now: u64) -> std::io::Result<Vec<QuotaScheduleChange>> {
let names = {
let work = self.schedule_work.lock();
work.dirty
.iter()
.cloned()
.chain(
work.due
.iter()
.take_while(|(at, _)| *at <= now)
.map(|(_, user)| user.clone()),
)
.collect::<BTreeSet<_>>()
};
let mut changes = Vec::with_capacity(names.len());
for user in names {
let Some(target) = self.current_handle(&user) else {
continue;
};
let state = target.counters.schedule.lock().clone();
let mut checkpoint = state.checkpoint;
let reset = if state.pending {
checkpoint.quota_period = state.requested;
checkpoint.period_cursor_epoch_secs =
checkpoint.period_cursor_epoch_secs.max(state.accepted_at);
if let Some(period) = checkpoint.quota_period {
period
.next(checkpoint.period_cursor_epoch_secs)
.inspect_err(|_| {
target.counters.schedule.lock().failed = true;
})?;
}
false
} else if let Some(period) = checkpoint.quota_period {
let start = period.start(now).inspect_err(|_| {
target.counters.schedule.lock().failed = true;
})?;
if start <= checkpoint.period_cursor_epoch_secs {
continue;
}
checkpoint.period_cursor_epoch_secs = start;
period.next(start).inspect_err(|_| {
target.counters.schedule.lock().failed = true;
})?;
true
} else {
continue;
};
changes.push(QuotaScheduleChange {
user,
target,
checkpoint,
reset,
});
}
Ok(changes)
}
/// Publishes durable metadata without overwriting traffic accumulated during I/O.
pub(crate) fn publish_schedule(&self, change: &QuotaScheduleChange, now: u64) {
if change.reset {
change.target.reset(now);
}
let mut schedule = change.target.counters.schedule.lock();
schedule.checkpoint = change.checkpoint;
schedule.pending = false;
schedule.failed = false;
drop(schedule);
self.schedule_index(&change.user, Some(&change.target));
}
/// Marks the attempted batch without changing its accounting or cursor.
pub(crate) fn schedule_failed(&self, changes: &[QuotaScheduleChange]) {
for change in changes {
change.target.counters.schedule.lock().failed = true;
}
}
/// Reindexes a manual reset after its durable cursor publication.
pub(crate) fn refresh_schedule(&self, user: &str) {
self.schedule_index(user, self.current_handle(user).as_ref());
}
/// A successful full checkpoint has reconciled retired accounts as well.
pub(crate) fn schedule_saved(&self) {
self.schedule_work.lock().removed = false;
}
/// Lists current ownership, excluding deletion tombstones and stale generations.
pub(crate) fn configured_quota_users(&self) -> BTreeSet<String> {
self.users
.iter()
.filter_map(|entry| {
entry
.value()
.state
.lock()
.current
.as_ref()
.map(|_| entry.key().clone())
})
.collect()
}
/// Resolves current ownership without creating a legacy account for a tombstone.
pub(crate) fn current_handle(&self, user: &str) -> Option<UserQuotaHandle> {
let slot = self.users.get(user)?;
let state = slot.state.lock();
let account = state.current.as_ref()?;
Some(UserQuotaHandle {
incarnation: account.incarnation,
counters: Arc::clone(&account.counters),
})
}
/// Reads status while the caller owns the process mutation gate.
pub(crate) fn schedule_status(
&self,
user: &str,
desired: Option<UserQuotaPeriod>,
now: u64,
) -> QuotaScheduleStatus {
let Some(target) = self.current_handle(user) else {
return QuotaScheduleStatus {
quota_schedule_state: if desired.is_some() {
"pending"
} else {
"disabled"
},
..Default::default()
};
};
let schedule = target.counters.schedule.lock();
let period = (!schedule.pending)
.then_some(schedule.checkpoint.quota_period)
.flatten();
let cursor = schedule.checkpoint.period_cursor_epoch_secs;
let start = period.and_then(|p| p.start(cursor).ok());
let next = period.and_then(|p| p.next(cursor).ok());
let state = if schedule.failed {
"error"
} else if schedule.pending || schedule.requested != desired || schedule.fresh.is_some() {
"pending"
} else if next.is_some_and(|at| at <= now) {
"overdue"
} else if period.is_some() {
"active"
} else {
"disabled"
};
QuotaScheduleStatus {
active_quota_period: period,
quota_period_start_epoch_secs: start,
quota_next_reset_epoch_secs: next,
quota_schedule_state: state,
}
}
}
impl UserQuotaHandle {
/// Captures the current accounting generation for a cold checkpoint.
pub(crate) fn snapshot(&self) -> super::UserQuotaSnapshot {
self.counters.snapshot()
}
/// Captures only cold calendar state; reserve/charge never take this lock.
pub(crate) fn schedule_checkpoint(&self) -> QuotaScheduleCheckpoint {
self.counters.schedule.lock().checkpoint
}
/// Keeps manual-only legacy checkpoints unchanged until a calendar is configured.
pub(crate) fn manual_schedule_checkpoint(&self, now: u64) -> QuotaScheduleCheckpoint {
let schedule = self.counters.schedule.lock();
let mut checkpoint = schedule.checkpoint;
if schedule.requested.is_some() || checkpoint != QuotaScheduleCheckpoint::default() {
checkpoint.period_cursor_epoch_secs = checkpoint.period_cursor_epoch_secs.max(now);
}
checkpoint
}
/// Publishes a durable manual watermark without shifting the UTC cadence.
pub(crate) fn publish_manual_schedule(&self, checkpoint: QuotaScheduleCheckpoint) {
let mut schedule = self.counters.schedule.lock();
schedule.checkpoint = checkpoint;
if !schedule.pending {
schedule.failed = false;
}
}
}
+130
View File
@@ -0,0 +1,130 @@
use super::*;
use crate::stats::Stats;
use std::sync::Barrier;
#[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);
}
#[test]
fn quota_snapshot_never_combines_different_generations() {
const ITERATIONS: u64 = 10_000;
let store = Arc::new(QuotaStore::default());
store.load("alice", 1, 1);
let barrier = Arc::new(Barrier::new(2));
let writer_store = Arc::clone(&store);
let writer_barrier = Arc::clone(&barrier);
let writer = std::thread::spawn(move || {
writer_barrier.wait();
for generation in 2..=ITERATIONS {
writer_store.load("alice", generation, generation);
}
});
barrier.wait();
for _ in 0..ITERATIONS {
let snapshot = store.snapshot().remove("alice").unwrap();
assert_eq!(snapshot.used_bytes, snapshot.last_reset_epoch_secs);
}
writer.join().unwrap();
}
#[test]
fn repeated_old_generation_refunds_leave_new_usage_intact() {
const ITERATIONS: u64 = 10_000;
let store = QuotaStore::default();
let state = store.user("alice");
for generation in 1..=ITERATIONS {
let reservation = state.try_reserve(80, 100).unwrap();
store.load("alice", 40, generation);
drop(reservation);
assert_eq!(store.used("alice"), 40);
store.reset("alice", generation);
}
}
#[test]
fn retired_incarnation_cannot_charge_recreated_username() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let retired = store.handle_exact("alice", 1).unwrap();
retired.charge(40);
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
retired.charge(20);
assert_eq!(retired.used(), 60);
assert_eq!(store.handle_exact("alice", 3).unwrap().used(), 0);
}
#[test]
fn credential_rotation_preserves_usage_without_sharing_future_charges() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let old = store.handle_exact("alice", 1).unwrap();
old.charge(40);
store.advance_preserving_usage("alice", 2);
let current = store.handle_exact("alice", 2).unwrap();
old.charge(20);
assert_eq!(old.used(), 60);
assert_eq!(current.used(), 40);
}
#[test]
fn captured_reset_handle_cannot_reset_a_new_incarnation() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let reset_target = store.handle_exact("alice", 1).unwrap();
reset_target.charge(40);
store.advance_preserving_usage("alice", 2);
let current = store.handle_exact("alice", 2).unwrap();
current.charge(20);
reset_target.reset(7);
assert_eq!(reset_target.used(), 0);
assert_eq!(current.used(), 60);
}
#[test]
fn stale_retirement_cannot_remove_newer_quota_owner() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
store.retire_through("alice", 2);
assert!(store.handle_exact("alice", 3).is_some());
}
#[test]
fn old_reservation_refund_does_not_debit_recreated_username() {
let store = QuotaStore::default();
store.activate_fresh("alice", 1);
let old = store.handle_exact("alice", 1).unwrap();
let reservation = old.try_reserve(80, 100).unwrap();
store.retire_through("alice", 2);
store.activate_fresh("alice", 3);
let current = store.handle_exact("alice", 3).unwrap();
current.charge(50);
drop(reservation);
assert_eq!(old.used(), 0);
assert_eq!(current.used(), 50);
}