From 6fdc2a5ba47e26bc638497bbc982590e5d678101 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Mon, 19 Oct 2026 00:15:27 +0300 Subject: [PATCH] Rustfmt --- src/api/config_edit.rs | 3 + src/api/config_edit/quota_period_tests.rs | 19 + src/api/config_store.rs | 3 + src/api/config_store/persistence.rs | 5 + src/api/handler/read_routes.rs | 8 +- src/api/model/users.rs | 12 + src/api/users.rs | 3 + src/api/users/create.rs | 18 + src/api/users/lifecycle.rs | 9 +- src/api/users/period_tests.rs | 445 ++++++++++++++++++ src/api/users/update.rs | 21 + src/api/users/view.rs | 8 +- src/config/hot_reload.rs | 3 + src/config/hot_reload/fields.rs | 5 + src/config/hot_reload/quota_period_tests.rs | 33 ++ src/config/hot_reload/reporting.rs | 6 + src/config/hot_reload/tests.rs | 2 + src/config/hot_reload/watcher.rs | 21 +- .../hot_reload/watcher/decoy_dns_tests.rs | 10 +- src/config/hot_reload/watcher/preparation.rs | 22 + .../hot_reload/watcher/quota_period_tests.rs | 105 +++++ src/config/load.rs | 16 +- src/config/load/strict_keys.rs | 1 + src/config/load/validate_runtime.rs | 1 + src/config/tests/load_basic_tests.rs | 2 + .../load_basic_tests/quota_period_tests.rs | 69 +++ src/config/types.rs | 3 + src/config/types/access.rs | 27 ++ src/config/types/quota_period.rs | 114 +++++ src/maestro/orchestrator.rs | 20 +- src/maestro/runtime_tasks.rs | 8 + src/proxy/shared_state.rs | 16 + src/proxy/user_admission.rs | 46 ++ src/quota_state.rs | 146 +++++- src/quota_state/period_tests.rs | 297 ++++++++++++ src/quota_state/period_tests/concurrency.rs | 232 +++++++++ .../period_tests/resource_model.rs | 140 ++++++ src/quota_state/periods.rs | 174 +++++++ src/quota_state/test_io.rs | 97 ++++ src/stats/mod.rs | 4 +- src/stats/quota_store.rs | 176 ++----- src/stats/quota_store/schedule.rs | 443 +++++++++++++++++ src/stats/quota_store/tests.rs | 130 +++++ 43 files changed, 2764 insertions(+), 159 deletions(-) create mode 100644 src/api/config_edit/quota_period_tests.rs create mode 100644 src/api/users/period_tests.rs create mode 100644 src/config/hot_reload/quota_period_tests.rs create mode 100644 src/config/hot_reload/watcher/quota_period_tests.rs create mode 100644 src/config/tests/load_basic_tests/quota_period_tests.rs create mode 100644 src/config/types/quota_period.rs create mode 100644 src/quota_state/period_tests.rs create mode 100644 src/quota_state/period_tests/concurrency.rs create mode 100644 src/quota_state/period_tests/resource_model.rs create mode 100644 src/quota_state/periods.rs create mode 100644 src/quota_state/test_io.rs create mode 100644 src/stats/quota_store/schedule.rs create mode 100644 src/stats/quota_store/tests.rs diff --git a/src/api/config_edit.rs b/src/api/config_edit.rs index 6c34dbb..362df0b 100644 --- a/src/api/config_edit.rs +++ b/src/api/config_edit.rs @@ -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; diff --git a/src/api/config_edit/quota_period_tests.rs b/src/api/config_edit/quota_period_tests.rs new file mode 100644 index 0000000..d3688bf --- /dev/null +++ b/src/api/config_edit/quota_period_tests.rs @@ -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); +} diff --git a/src/api/config_store.rs b/src/api/config_store.rs index 01f6876..7813859 100644 --- a/src/api/config_store.rs +++ b/src/api/config_store.rs @@ -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", } diff --git a/src/api/config_store/persistence.rs b/src/api/config_store/persistence.rs index 0a3219c..c49b598 100644 --- a/src/api/config_store/persistence.rs +++ b/src/api/config_store/persistence.rs @@ -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 = 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(), } diff --git a/src/api/handler/read_routes.rs b/src/api/handler/read_routes.rs index 8c0ef4a..b389a12 100644 --- a/src/api/handler/read_routes.rs +++ b/src/api/handler/read_routes.rs @@ -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)) } diff --git a/src/api/model/users.rs b/src/api/model/users.rs index 8d6390b..d2facfe 100644 --- a/src/api/model/users.rs +++ b/src/api/model/users.rs @@ -23,6 +23,8 @@ pub(in crate::api) struct UserInfo { pub(in crate::api) max_tcp_conns: Option, pub(in crate::api) expiration_rfc3339: Option, pub(in crate::api) data_quota_bytes: Option, + /// Desired UTC period; runtime activation is reported by the quota endpoint. + pub(in crate::api) quota_period: Option, pub(in crate::api) rate_limit_up_bps: Option, pub(in crate::api) rate_limit_down_bps: Option, pub(in crate::api) max_unique_ips: Option, @@ -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, + /// 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, pub(in crate::api) expiration_rfc3339: Option, pub(in crate::api) data_quota_bytes: Option, + /// Desired UTC period; runtime activation is reported by the quota endpoint. + pub(in crate::api) quota_period: Option, pub(in crate::api) rate_limit_up_bps: Option, pub(in crate::api) rate_limit_down_bps: Option, pub(in crate::api) max_unique_ips: Option, @@ -98,6 +107,9 @@ pub(in crate::api) struct PatchUserRequest { pub(in crate::api) expiration_rfc3339: Patch, #[serde(default, deserialize_with = "patch_field")] pub(in crate::api) data_quota_bytes: Patch, + /// Omitted preserves the period; null removes it. + #[serde(default, deserialize_with = "patch_field")] + pub(in crate::api) quota_period: Patch, #[serde(default, deserialize_with = "patch_field")] pub(in crate::api) rate_limit_up_bps: Patch, #[serde(default, deserialize_with = "patch_field")] diff --git a/src/api/users.rs b/src/api/users.rs index b16f4bb..bca151d 100644 --- a/src/api/users.rs +++ b/src/api/users.rs @@ -35,3 +35,6 @@ pub(super) use view::{build_user_quota_list, users_from_config}; #[cfg(test)] mod tests; + +#[cfg(test)] +mod period_tests; diff --git a/src/api/users/create.rs b/src/api/users/create.rs index 509b1ca..be33008 100644 --- a/src/api/users/create.rs +++ b/src/api/users/create.rs @@ -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, diff --git a/src/api/users/lifecycle.rs b/src/api/users/lifecycle.rs index b5a595a..4f945f2 100644 --- a/src/api/users/lifecycle.rs +++ b/src/api/users/lifecycle.rs @@ -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!( diff --git a/src/api/users/period_tests.rs b/src/api/users/period_tests.rs new file mode 100644 index 0000000..b34d1f1 --- /dev/null +++ b/src/api/users/period_tests.rs @@ -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, + config_tx: watch::Sender>, +} + +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::(serde_json::json!({"quota_period":"Weekly"})) + .is_err() + ); +} diff --git a/src/api/users/update.rs b/src/api/users/update.rs index 713dd53..6111118 100644 --- a/src/api/users/update.rs +++ b/src/api/users/update.rs @@ -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(); diff --git a/src/api/users/view.rs b/src/api/users/view.rs index 4c89d4d..98094a2 100644 --- a/src/api/users/view.rs +++ b/src/api/users/view.rs @@ -51,6 +51,7 @@ pub(in crate::api) async fn users_from_config( .get(&username) .map(chrono::DateTime::::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, diff --git a/src/config/hot_reload.rs b/src/config/hot_reload.rs index 9b1268a..4acf8b7 100644 --- a/src/config/hot_reload.rs +++ b/src/config/hot_reload.rs @@ -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; diff --git a/src/config/hot_reload/fields.rs b/src/config/hot_reload/fields.rs index 450583d..c3033ae 100644 --- a/src/config/hot_reload/fields.rs +++ b/src/config/hot_reload/fields.rs @@ -84,6 +84,8 @@ pub struct HotFields { pub user_max_tcp_conns_global_each: usize, pub user_expirations: std::collections::HashMap>, pub user_data_quota: std::collections::HashMap, + /// Recurring UTC calendars accepted by the active runtime source. + pub user_quota_periods: std::collections::HashMap, pub user_rate_limits: std::collections::HashMap, pub cidr_rate_limits: std::collections::HashMap, pub user_max_unique_ips: std::collections::HashMap, @@ -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(); diff --git a/src/config/hot_reload/quota_period_tests.rs b/src/config/hot_reload/quota_period_tests.rs new file mode 100644 index 0000000..771f1dd --- /dev/null +++ b/src/config/hot_reload/quota_period_tests.rs @@ -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)); + } + } +} diff --git a/src/config/hot_reload/reporting.rs b/src/config/hot_reload/reporting.rs index f532ee1..f7c289d 100644 --- a/src/config/hot_reload/reporting.rs +++ b/src/config/hot_reload/reporting.rs @@ -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)", diff --git a/src/config/hot_reload/tests.rs b/src/config/hot_reload/tests.rs index f81bbd7..128d025 100644 --- a/src/config/hot_reload/tests.rs +++ b/src/config/hot_reload/tests.rs @@ -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); diff --git a/src/config/hot_reload/watcher.rs b/src/config/hot_reload/watcher.rs index 7eaed74..80baefb 100644 --- a/src/config/hot_reload/watcher.rs +++ b/src/config/hot_reload/watcher.rs @@ -156,6 +156,7 @@ fn reload_config_with_resolver( detected_ip_v6: Option, reload_state: &mut ReloadState, dns_resolver: Option<&crate::network::dns_overrides::GenerationDnsResolver>, + refresh_quota_source: bool, ) -> Option { 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>, mut activation: Option>, + quota_mutation: Option>>, + quota_source_refresh: Option bool + Send + Sync>>, ) -> ( watch::Receiver>, watch::Receiver, @@ -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; diff --git a/src/config/hot_reload/watcher/decoy_dns_tests.rs b/src/config/hot_reload/watcher/decoy_dns_tests.rs index bb30aa4..6085606 100644 --- a/src/config/hot_reload/watcher/decoy_dns_tests.rs +++ b/src/config/hot_reload/watcher/decoy_dns_tests.rs @@ -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(); diff --git a/src/config/hot_reload/watcher/preparation.rs b/src/config/hot_reload/watcher/preparation.rs index 62062eb..72f21ba 100644 --- a/src/config/hot_reload/watcher/preparation.rs +++ b/src/config/hot_reload/watcher/preparation.rs @@ -34,9 +34,30 @@ pub(super) async fn reload( reload_state: &mut ReloadState, dns_resolver: Option>, force: bool, + quota_mutation: Option<&Arc>>, + quota_source_refresh: Option<&(dyn Fn() -> bool + Send + Sync)>, ) -> Option { + // 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!( diff --git a/src/config/hot_reload/watcher/quota_period_tests.rs b/src/config/hot_reload/watcher/quota_period_tests.rs new file mode 100644 index 0000000..199bb3c --- /dev/null +++ b/src/config/hot_reload/watcher/quota_period_tests.rs @@ -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(); +} diff --git a/src/config/load.rs b/src/config/load.rs index 424bb1c..98a005a 100644 --- a/src/config/load.rs +++ b/src/config/load.rs @@ -156,6 +156,10 @@ pub struct ProxyConfig { #[serde(default)] pub default_dc: Option, + /// Monotonic start of source acquisition, used only for cold quota policy fencing. + #[serde(skip)] + pub(crate) quota_source_read_at: Option, + /// Precomputed authentication snapshot for handshake hot paths. #[serde(skip)] pub(crate) runtime_user_auth: Option>, @@ -177,8 +181,11 @@ impl ProxyConfig { path: P, source_overrides: &BTreeMap, ) -> Result { + 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, ) -> Result { - 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())); } diff --git a/src/config/load/strict_keys.rs b/src/config/load/strict_keys.rs index 598620d..eda1a8a 100644 --- a/src/config/load/strict_keys.rs +++ b/src/config/load/strict_keys.rs @@ -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", diff --git a/src/config/load/validate_runtime.rs b/src/config/load/validate_runtime.rs index 6cae0ed..2dc43ef 100644 --- a/src/config/load/validate_runtime.rs +++ b/src/config/load/validate_runtime.rs @@ -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(), diff --git a/src/config/tests/load_basic_tests.rs b/src/config/tests/load_basic_tests.rs index 82c7da4..f6a025c 100644 --- a/src/config/tests/load_basic_tests.rs +++ b/src/config/tests/load_basic_tests.rs @@ -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"] diff --git a/src/config/tests/load_basic_tests/quota_period_tests.rs b/src/config/tests/load_basic_tests/quota_period_tests.rs new file mode 100644 index 0000000..c2f24bb --- /dev/null +++ b/src/config/tests/load_basic_tests/quota_period_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::(&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")); +} diff --git a/src/config/types.rs b/src/config/types.rs index 88849fc..51d8efd 100644 --- a/src/config/types.rs +++ b/src/config/types.rs @@ -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, diff --git a/src/config/types/access.rs b/src/config/types/access.rs index ac35dff..7db7900 100644 --- a/src/config/types/access.rs +++ b/src/config/types/access.rs @@ -30,6 +30,10 @@ pub struct AccessConfig { #[serde(default)] pub user_data_quota: HashMap, + /// Optional UTC calendar for users with a positive data quota. + #[serde(default)] + pub user_quota_periods: HashMap, + /// 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) } diff --git a/src/config/types/quota_period.rs b/src/config/types/quota_period.rs new file mode 100644 index 0000000..fee73ff --- /dev/null +++ b/src/config/types/quota_period.rs @@ -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 { + 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 { + 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()); + } +} diff --git a/src/maestro/orchestrator.rs b/src/maestro/orchestrator.rs index aa80973..7859ab3 100644 --- a/src/maestro/orchestrator.rs +++ b/src/maestro/orchestrator.rs @@ -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::>(); - 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 { diff --git a/src/maestro/runtime_tasks.rs b/src/maestro/runtime_tasks.rs index 891e6d6..b8a1970 100644 --- a/src/maestro/runtime_tasks.rs +++ b/src/maestro/runtime_tasks.rs @@ -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!( diff --git a/src/proxy/shared_state.rs b/src/proxy/shared_state.rs index 0bc500d..c532a0b 100644 --- a/src/proxy/shared_state.rs +++ b/src/proxy/shared_state.rs @@ -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> { + 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, + 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, diff --git a/src/proxy/user_admission.rs b/src/proxy/user_admission.rs index 8e8724c..c913fd9 100644 --- a/src/proxy/user_admission.rs +++ b/src/proxy/user_admission.rs @@ -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> { + 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, + 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 diff --git a/src/quota_state.rs b/src/quota_state.rs index 2db71f0..fdd2265 100644 --- a/src/quota_state.rs +++ b/src/quota_state.rs @@ -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, mutation: Arc>, + #[cfg(test)] + scheduler_iterations: std::sync::atomic::AtomicU64, } impl QuotaStateOwner { @@ -36,8 +46,10 @@ impl QuotaStateOwner { pub(crate) fn new(path: PathBuf, store: Arc) -> Arc { 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) { + let _ = self.load_inner(configured_users, false).await; + } + + async fn load_inner( + &self, + configured_users: &BTreeSet, + 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 { 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, + user: &str, + last_reset_epoch_secs: u64, + guard: tokio::sync::OwnedMutexGuard<()>, + ) -> std::io::Result { 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, + 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::>(); + 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( } 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; diff --git a/src/quota_state/period_tests.rs b/src/quota_state/period_tests.rs new file mode 100644 index 0000000..df52987 --- /dev/null +++ b/src/quota_state/period_tests.rs @@ -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) -> Arc { + 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, 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, Arc) { + 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; diff --git a/src/quota_state/period_tests/concurrency.rs b/src/quota_state/period_tests/concurrency.rs new file mode 100644 index 0000000..f23761a --- /dev/null +++ b/src/quota_state/period_tests/concurrency.rs @@ -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); +} diff --git a/src/quota_state/period_tests/resource_model.rs b/src/quota_state/period_tests/resource_model.rs new file mode 100644 index 0000000..cd6dc4b --- /dev/null +++ b/src/quota_state/period_tests/resource_model.rs @@ -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::>(); + 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()); +} diff --git a/src/quota_state/periods.rs b/src/quota_state/periods.rs new file mode 100644 index 0000000..de9ee6b --- /dev/null +++ b/src/quota_state/periods.rs @@ -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, + credential: crate::proxy::user_admission::UserCredentialId, + period: Option, + quota: Option, + ) { + 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, + ) -> 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, 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 = 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, + } + } => {}, + } + } + }); + } +} diff --git a/src/quota_state/test_io.rs b/src/quota_state/test_io.rs new file mode 100644 index 0000000..2f0d496 --- /dev/null +++ b/src/quota_state/test_io.rs @@ -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>> { + static HOOKS: OnceLock>>> = OnceLock::new(); + HOOKS.get_or_init(Mutex::default) +} + +pub(super) struct Hook { + path: PathBuf, + started: Mutex>>, + released: Mutex, + wake: Condvar, + pub(super) writes: AtomicUsize, + pub(super) fail_after: AtomicBool, +} + +pub(super) struct HookGuard(Arc); + +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(()) +} diff --git a/src/stats/mod.rs b/src/stats/mod.rs index 036bec9..0a85b6b 100644 --- a/src/stats/mod.rs +++ b/src/stats/mod.rs @@ -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; diff --git a/src/stats/quota_store.rs b/src/stats/quota_store.rs index 48e95b0..5369e21 100644 --- a/src/stats/quota_store.rs +++ b/src/stats/quota_store.rs @@ -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>, + mutation: Arc>, + schedule_work: Mutex, + 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, + schedule: Mutex, } 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 { @@ -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; diff --git a/src/stats/quota_store/schedule.rs b/src/stats/quota_store/schedule.rs new file mode 100644 index 0000000..d86e9f0 --- /dev/null +++ b/src/stats/quota_store/schedule.rs @@ -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, + /// 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, + /// 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, + /// Latest committed quota amount. + pub(super) quota: Option, +} + +/// 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, + /// 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, +} + +/// 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, + dirty: BTreeSet, + 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, + /// UTC boundary containing the durable cursor. + pub(crate) quota_period_start_epoch_secs: Option, + /// First boundary strictly after the durable cursor. + pub(crate) quota_next_reset_epoch_secs: Option, + /// 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> { + 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 { + 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, + credential: UserCredentialId, + period: Option, + quota: Option, + ) { + 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, + 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> { + 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::>() + }; + 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 { + 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 { + 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, + 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; + } + } +} diff --git a/src/stats/quota_store/tests.rs b/src/stats/quota_store/tests.rs new file mode 100644 index 0000000..37a1c55 --- /dev/null +++ b/src/stats/quota_store/tests.rs @@ -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); +}