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