This commit is contained in:
Alexey
2026-10-30 22:34:37 +03:00
parent e7e588ebe5
commit 5a5009174a
5 changed files with 242 additions and 66 deletions
+94 -23
View File
@@ -2,34 +2,66 @@ use super::*;
#[tokio::test]
async fn quota_period_failed_delete_recreate_restart_must_not_inherit_retired_calendar() {
for secret in ["0123456789abcdef0123456789abcdef", "fedcba9876543210fedcba9876543210"] {
for secret in [
"0123456789abcdef0123456789abcdef",
"fedcba9876543210fedcba9876543210",
] {
quota_period_delete_recreate_restart(true, secret).await;
}
}
#[tokio::test]
async fn quota_period_successful_delete_recreate_restart_has_fresh_calendar() {
for secret in ["0123456789abcdef0123456789abcdef", "fedcba9876543210fedcba9876543210"] {
for secret in [
"0123456789abcdef0123456789abcdef",
"fedcba9876543210fedcba9876543210",
] {
quota_period_delete_recreate_restart(false, secret).await;
}
}
async fn quota_period_delete_recreate_restart(fail_deletion: bool, secret: &str) {
let fixture = Fixture::new();
let create = |secret: &str| serde_json::from_value(serde_json::json!({
"username":"alice", "secret":secret,
"data_quota_bytes":1024, "quota_period":"monthly"
})).unwrap();
create_user(create("0123456789abcdef0123456789abcdef"), None, &fixture.shared).await.unwrap();
let create = |secret: &str| {
serde_json::from_value(serde_json::json!({
"username":"alice", "secret":secret,
"data_quota_bytes":1024, "quota_period":"monthly"
}))
.unwrap()
};
create_user(
create("0123456789abcdef0123456789abcdef"),
None,
&fixture.shared,
)
.await
.unwrap();
fixture.publish().await;
let now = crate::quota_state::now_epoch_secs();
// Model a previously processed future boundary followed by clock rollback.
let next = UserQuotaPeriod::Monthly.next(now).unwrap();
let future = UserQuotaPeriod::Monthly.next(next).unwrap();
fixture.shared.quota_state.process_schedules(future, false).await.unwrap();
fixture
.shared
.quota_state
.process_schedules(future, false)
.await
.unwrap();
fixture.store.current_handle("alice").unwrap().charge(700);
fixture.shared.quota_state.save(&std::collections::BTreeSet::from(["alice".into(), "admin".into()])).await.unwrap();
let old_checkpoint = fixture.store.current_handle("alice").unwrap().schedule_checkpoint();
fixture
.shared
.quota_state
.save(&std::collections::BTreeSet::from([
"alice".into(),
"admin".into(),
]))
.await
.unwrap();
let old_checkpoint = fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint();
let quota_path = fixture.shared.quota_state.path().to_owned();
let retired_file = quota_path.with_extension("retired");
if fail_deletion {
@@ -41,30 +73,60 @@ async fn quota_period_delete_recreate_restart(fail_deletion: bool, secret: &str)
assert!(fixture.store.current_handle("alice").is_none());
if fail_deletion {
let config_before = std::fs::read(&fixture.shared.config_path).unwrap();
assert!(create_user(create(secret), None, &fixture.shared).await.is_err());
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), config_before);
assert!(
create_user(create(secret), None, &fixture.shared)
.await
.is_err()
);
assert_eq!(
std::fs::read(&fixture.shared.config_path).unwrap(),
config_before
);
assert!(fixture.store.current_handle("alice").is_none());
// Restore the old durable bytes to model storage recovery before retry.
std::fs::remove_dir(&quota_path).unwrap();
std::fs::rename(&retired_file, &quota_path).unwrap();
}
create_user(create(secret), None, &fixture.shared).await.unwrap();
create_user(create(secret), None, &fixture.shared)
.await
.unwrap();
assert_eq!(fixture.store.used("alice"), 0);
assert_eq!(fixture.store.current_handle("alice").unwrap().schedule_checkpoint().period_cursor_epoch_secs, 0);
assert_eq!(
fixture
.store
.current_handle("alice")
.unwrap()
.schedule_checkpoint()
.period_cursor_epoch_secs,
0
);
let cfg = Arc::new(ProxyConfig::load(&fixture.shared.config_path).unwrap());
let restart_store = Arc::new(QuotaStore::default());
let restart = QuotaStateOwner::new(quota_path, restart_store.clone());
restart.load_for_config(&cfg).await.unwrap();
let authority = UserAdmissionAuthority::new_with_quota_store(restart_store.clone());
let gate = restart.lock().await;
authority.activate_config_source(1, None, &cfg.access.users, &cfg.access.user_enabled).unwrap();
authority
.activate_config_source(1, None, &cfg.access.users, &cfg.access.user_enabled)
.unwrap();
authority.publish_quota_config(1, &cfg, true);
drop(gate);
restart.process_schedules(now, true).await.unwrap();
let actual = restart_store.current_handle("alice").unwrap().schedule_checkpoint();
let actual = restart_store
.current_handle("alice")
.unwrap()
.schedule_checkpoint();
assert!(actual.period_cursor_epoch_secs < old_checkpoint.period_cursor_epoch_secs);
assert_eq!(restart_store.schedule_deadline(), Some(next), "recreated account inherited a retired calendar watermark across restart");
assert_eq!(restart_store.used("alice"), 0, "recreated account inherited retired usage");
assert_eq!(
restart_store.schedule_deadline(),
Some(next),
"recreated account inherited a retired calendar watermark across restart"
);
assert_eq!(
restart_store.used("alice"),
0,
"recreated account inherited retired usage"
);
fixture.finish().await;
}
@@ -82,11 +144,16 @@ async fn quota_period_create_rejects_unreadable_checkpoint_without_publishing_us
let result = create_user(
serde_json::from_value(serde_json::json!({
"username":"alice", "data_quota_bytes":1024, "quota_period":"monthly"
})).unwrap(),
}))
.unwrap(),
None,
&fixture.shared,
).await;
assert!(result.is_err(), "CREATE must reject untrusted durable quota state");
)
.await;
assert!(
result.is_err(),
"CREATE must reject untrusted durable quota state"
);
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), before);
assert!(fixture.store.current_handle("alice").is_none());
fixture.finish().await;
@@ -103,10 +170,14 @@ async fn quota_period_create_revision_conflict_does_not_touch_retired_checkpoint
serde_json::from_value(serde_json::json!({"username":"alice"})).unwrap(),
Some("outdated-revision".into()),
&fixture.shared,
).await;
)
.await;
assert_eq!(result.err().unwrap().code, "revision_conflict");
assert_eq!(std::fs::read(&fixture.shared.config_path).unwrap(), before);
assert_eq!(std::fs::read(fixture.shared.quota_state.path()).unwrap(), retired);
assert_eq!(
std::fs::read(fixture.shared.quota_state.path()).unwrap(),
retired
);
assert!(fixture.store.current_handle("alice").is_none());
fixture.finish().await;
}
+14 -5
View File
@@ -13,10 +13,10 @@ pub(super) fn me_connection_cleanup_capacity(max_connections: u32) -> usize {
#[cfg(test)]
mod tests {
use super::*;
use crate::transport::middle_proxy::{ConnRegistry, MeResponse};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Semaphore;
use crate::transport::middle_proxy::{ConnRegistry, MeResponse};
#[test]
fn me_cleanup_capacity_preserves_zero_and_finite_limits() {
@@ -25,7 +25,12 @@ mod tests {
(1, 129),
(128, 256),
(10_000, 10_128),
(u32::MAX, (u32::MAX as usize).saturating_add(128).min(Semaphore::MAX_PERMITS)),
(
u32::MAX,
(u32::MAX as usize)
.saturating_add(128)
.min(Semaphore::MAX_PERMITS),
),
] {
let capacity = me_connection_cleanup_capacity(configured);
assert_eq!(capacity, expected);
@@ -36,17 +41,21 @@ mod tests {
#[tokio::test]
async fn me_cleanup_capacity_zero_admits_512_simultaneous_leased_routes() {
let registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(
8, me_connection_cleanup_capacity(0),
8,
me_connection_cleanup_capacity(0),
));
let mut clients = Vec::new();
for index in 0..512 {
let client = tokio::time::timeout(Duration::from_secs(1), registry.register_leased())
.await.unwrap_or_else(|_| panic!("registration {} blocked", index + 1))
.await
.unwrap_or_else(|_| panic!("registration {} blocked", index + 1))
.unwrap();
clients.push(client);
}
for (lease, mut receiver) in clients {
registry.route_nowait(lease.conn_id(), MeResponse::Ack(1)).await;
registry
.route_nowait(lease.conn_id(), MeResponse::Ack(1))
.await;
assert!(matches!(receiver.try_recv(), Ok(MeResponse::Ack(1))));
lease.unregister().await;
}
+96 -26
View File
@@ -3,21 +3,30 @@ use std::sync::atomic::Ordering;
async fn fixture() -> (tempfile::TempDir, Arc<QuotaStateOwner>) {
let dir = tempfile::tempdir().unwrap();
let owner = QuotaStateOwner::new(dir.path().join("quota.json"), Arc::new(QuotaStore::default()));
let owner = QuotaStateOwner::new(
dir.path().join("quota.json"),
Arc::new(QuotaStore::default()),
);
let mut state = QuotaStateFile::default();
state.last_reset_epoch_secs = 123;
state.users.insert("alice".into(), quota_user_state(UserQuotaSnapshot {
used_bytes: 700,
last_reset_epoch_secs: 123,
}));
state.users.insert("bob".into(), QuotaUserState {
used_bytes: 321,
last_reset_epoch_secs: 111,
schedule: QuotaScheduleCheckpoint {
quota_period: Some(crate::config::UserQuotaPeriod::Monthly),
period_cursor_epoch_secs: 1_793_491_200,
state.users.insert(
"alice".into(),
quota_user_state(UserQuotaSnapshot {
used_bytes: 700,
last_reset_epoch_secs: 123,
}),
);
state.users.insert(
"bob".into(),
QuotaUserState {
used_bytes: 321,
last_reset_epoch_secs: 111,
schedule: QuotaScheduleCheckpoint {
quota_period: Some(crate::config::UserQuotaPeriod::Monthly),
period_cursor_epoch_secs: 1_793_491_200,
},
},
});
);
write_state_file(owner.path(), state).await.unwrap();
(dir, owner)
}
@@ -27,12 +36,17 @@ async fn creation_removes_only_retired_row_and_returns_owned_gate() {
let (_dir, owner) = fixture().await;
let before = read_state_file(owner.path()).await.unwrap().unwrap();
let guard = owner.lock().await;
let guard = owner.prepare_user_creation_locked("alice", guard).await.unwrap();
let guard = owner
.prepare_user_creation_locked("alice", guard)
.await
.unwrap();
assert!(owner.mutation.try_lock().is_err());
let after = read_state_file(owner.path()).await.unwrap().unwrap();
assert!(!after.users.contains_key("alice"));
assert_eq!(serde_json::to_value(&after.users["bob"]).unwrap(),
serde_json::to_value(&before.users["bob"]).unwrap());
assert_eq!(
serde_json::to_value(&after.users["bob"]).unwrap(),
serde_json::to_value(&before.users["bob"]).unwrap()
);
assert_eq!(after.last_reset_epoch_secs, before.last_reset_epoch_secs);
drop(guard);
assert!(owner.mutation.try_lock().is_ok());
@@ -44,7 +58,10 @@ async fn creation_without_checkpoint_does_not_require_a_write() {
let path = dir.path().join("missing/quota.json");
let owner = QuotaStateOwner::new(path.clone(), Arc::new(QuotaStore::default()));
let guard = owner.lock().await;
let _guard = owner.prepare_user_creation_locked("alice", guard).await.unwrap();
let _guard = owner
.prepare_user_creation_locked("alice", guard)
.await
.unwrap();
assert!(!path.exists());
assert!(owner.mutation.try_lock().is_err());
}
@@ -55,15 +72,35 @@ async fn creation_retries_durability_after_ambiguous_removal() {
let (hook, _) = test_io::Hook::install(owner.path(), false);
hook.fail_after.store(true, Ordering::Relaxed);
let guard = owner.lock().await;
assert!(owner.prepare_user_creation_locked("alice", guard).await.is_err());
assert!(!read_state_file(owner.path()).await.unwrap().unwrap().users.contains_key("alice"));
assert!(
owner
.prepare_user_creation_locked("alice", guard)
.await
.is_err()
);
assert!(
!read_state_file(owner.path())
.await
.unwrap()
.unwrap()
.users
.contains_key("alice")
);
// The absent row does not prove that the previous directory sync succeeded.
let guard = owner.lock().await;
assert!(owner.prepare_user_creation_locked("alice", guard).await.is_err());
assert!(
owner
.prepare_user_creation_locked("alice", guard)
.await
.is_err()
);
assert_eq!(hook.writes.load(Ordering::Relaxed), 2);
hook.fail_after.store(false, Ordering::Relaxed);
let guard = owner.lock().await;
let _guard = owner.prepare_user_creation_locked("alice", guard).await.unwrap();
let _guard = owner
.prepare_user_creation_locked("alice", guard)
.await
.unwrap();
assert_eq!(hook.writes.load(Ordering::Relaxed), 3);
}
@@ -80,10 +117,26 @@ async fn creation_waiter_abort_keeps_gate_until_durable_removal_finishes() {
waiter.abort();
assert!(waiter.await.unwrap_err().is_cancelled());
assert!(owner.mutation.try_lock().is_err());
assert!(read_state_file(owner.path()).await.unwrap().unwrap().users.contains_key("alice"));
assert!(
read_state_file(owner.path())
.await
.unwrap()
.unwrap()
.users
.contains_key("alice")
);
hook.release();
let _guard = tokio::time::timeout(std::time::Duration::from_secs(2), owner.lock()).await.unwrap();
assert!(!read_state_file(owner.path()).await.unwrap().unwrap().users.contains_key("alice"));
let _guard = tokio::time::timeout(std::time::Duration::from_secs(2), owner.lock())
.await
.unwrap();
assert!(
!read_state_file(owner.path())
.await
.unwrap()
.unwrap()
.users
.contains_key("alice")
);
assert_eq!(hook.writes.load(Ordering::Relaxed), 1);
}
@@ -104,8 +157,25 @@ async fn creation_write_failure_keeps_retired_checkpoint_for_safe_retry() {
assert!(writer.await.unwrap().is_err());
std::fs::remove_dir(owner.path()).unwrap();
std::fs::rename(retired, owner.path()).unwrap();
assert!(read_state_file(owner.path()).await.unwrap().unwrap().users.contains_key("alice"));
assert!(
read_state_file(owner.path())
.await
.unwrap()
.unwrap()
.users
.contains_key("alice")
);
let guard = owner.lock().await;
let _guard = owner.prepare_user_creation_locked("alice", guard).await.unwrap();
assert!(!read_state_file(owner.path()).await.unwrap().unwrap().users.contains_key("alice"));
let _guard = owner
.prepare_user_creation_locked("alice", guard)
.await
.unwrap();
assert!(
!read_state_file(owner.path())
.await
.unwrap()
.unwrap()
.users
.contains_key("alice")
);
}
@@ -1,19 +1,28 @@
use super::super::ConnRegistry;
use super::super::pool_writer_security_tests::make_pool;
use super::*;
use std::future::{Future, poll_fn};
use std::task::Poll;
use super::super::pool_writer_security_tests::make_pool;
use super::super::ConnRegistry;
#[tokio::test]
async fn shutdown_cancels_registration_waiting_for_cleanup_capacity() {
let mut pool = make_pool().await;
Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().routing).unwrap().registry =
Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1));
Arc::get_mut(&mut Arc::get_mut(&mut pool).unwrap().routing)
.unwrap()
.registry = Arc::new(ConnRegistry::with_route_and_cleanup_capacity(8, 1));
let (lease, _receiver) = pool.register_connection().await.unwrap();
let mut waiting = Box::pin(pool.register_connection());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
pool.begin_shutdown();
assert!(tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().is_none());
assert!(
tokio::time::timeout(Duration::from_secs(1), waiting)
.await
.unwrap()
.is_none()
);
lease.unregister().await;
assert!(pool.shutdown_until(Duration::from_secs(1)).await);
assert!(pool.register_connection().await.is_none());
@@ -13,16 +13,28 @@ async fn finite_cleanup_capacity_waits_until_dropped_route_is_received() {
}
assert_eq!(registry.routing.map.len(), 128);
assert_eq!(registry.cleanup_tx.capacity(), 0);
assert!(matches!(cleanup.try_recv(), Err(mpsc::error::TryRecvError::Empty)));
assert!(matches!(
cleanup.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
));
let mut waiting = Box::pin(registry.register_leased());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
let retired = clients.last().unwrap().0.conn_id();
drop(clients.pop().unwrap());
assert!(!registry.routing.map.contains_key(&retired));
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
assert_eq!(cleanup.recv().await, Some(retired));
registry.unregister(retired).await;
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), waiting).await.unwrap().unwrap();
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), waiting)
.await
.unwrap()
.unwrap();
lease.unregister().await;
for (lease, _) in clients {
lease.unregister().await;
@@ -39,7 +51,10 @@ async fn cancelled_registration_leaks_neither_route_nor_cleanup_permit() {
let (lease, _receiver) = registry.register_leased().await.unwrap();
for _ in 0..64 {
let mut waiting = Box::pin(registry.register_leased());
assert!(matches!(poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await, Poll::Pending));
assert!(matches!(
poll_fn(|cx| Poll::Ready(waiting.as_mut().poll(cx))).await,
Poll::Pending
));
drop(waiting);
assert_eq!(registry.routing.map.len(), 1);
assert_eq!(registry.cleanup_tx.capacity(), 0);
@@ -48,7 +63,9 @@ async fn cancelled_registration_leaks_neither_route_nor_cleanup_permit() {
assert!(registry.routing.map.is_empty());
assert_eq!(registry.cleanup_tx.capacity(), 1);
let (lease, _) = tokio::time::timeout(Duration::from_secs(1), registry.register_leased())
.await.unwrap().unwrap();
.await
.unwrap()
.unwrap();
lease.unregister().await;
assert_eq!(registry.cleanup_tx.capacity(), 1);
}