From c1ee43fbac29c7fcf11c276286036be61ecc646a Mon Sep 17 00:00:00 2001 From: David Osipov Date: Sat, 21 Mar 2026 15:54:14 +0400 Subject: [PATCH] Add stress testing for quota-lock and refactor test guard usage --- .github/workflows/rust.yml | 12 ++++ src/proxy/relay.rs | 7 ++ ...y_quota_lock_pressure_adversarial_tests.rs | 64 ++++++++++--------- ...ay_quota_wake_liveness_regression_tests.rs | 6 +- ...lay_quota_waker_storm_adversarial_tests.rs | 6 +- src/proxy/tests/relay_security_tests.rs | 12 +--- 6 files changed, 60 insertions(+), 47 deletions(-) diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index effe3ea..799f2ce 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -45,6 +45,18 @@ jobs: - name: Run tests run: cargo test --verbose + - name: Stress quota-lock suites (PR only) + if: github.event_name == 'pull_request' + env: + RUST_TEST_THREADS: 16 + run: | + set -euo pipefail + for i in $(seq 1 12); do + echo "[quota-lock-stress] iteration ${i}/12" + cargo test quota_lock_ --bin telemt -- --nocapture --test-threads 16 + cargo test relay_quota_wake --bin telemt -- --nocapture --test-threads 16 + done + # clippy dont fail on warnings because of active development of telemt # and many warnings - name: Run clippy diff --git a/src/proxy/relay.rs b/src/proxy/relay.rs index 6b71ace..88a8bd5 100644 --- a/src/proxy/relay.rs +++ b/src/proxy/relay.rs @@ -316,6 +316,13 @@ fn quota_user_lock_test_guard() -> &'static Mutex<()> { TEST_LOCK.get_or_init(|| Mutex::new(())) } +#[cfg(test)] +fn quota_user_lock_test_scope() -> std::sync::MutexGuard<'static, ()> { + quota_user_lock_test_guard() + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) +} + fn quota_overflow_user_lock(user: &str) -> Arc> { let stripes = QUOTA_USER_OVERFLOW_LOCKS.get_or_init(|| { (0..QUOTA_OVERFLOW_LOCK_STRIPES) diff --git a/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs b/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs index fd8fb2f..4add5f0 100644 --- a/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs +++ b/src/proxy/tests/relay_quota_lock_pressure_adversarial_tests.rs @@ -12,9 +12,7 @@ use tokio::time::Instant; #[test] fn quota_lock_same_user_returns_same_arc_instance() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -25,9 +23,7 @@ fn quota_lock_same_user_returns_same_arc_instance() { #[test] fn quota_lock_parallel_same_user_reuses_single_lock() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -51,9 +47,7 @@ fn quota_lock_parallel_same_user_reuses_single_lock() { #[test] fn quota_lock_unique_users_materialize_distinct_entries() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -74,9 +68,7 @@ fn quota_lock_unique_users_materialize_distinct_entries() { #[test] fn quota_lock_unique_churn_stress_keeps_all_inserted_keys_addressable() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -94,9 +86,7 @@ fn quota_lock_unique_churn_stress_keeps_all_inserted_keys_addressable() { #[test] fn quota_lock_saturation_returns_stable_overflow_lock_without_cache_growth() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -135,17 +125,19 @@ fn quota_lock_saturation_returns_stable_overflow_lock_without_cache_growth() { #[test] fn quota_lock_reclaims_unreferenced_entries_before_ephemeral_fallback() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); - // Fill and immediately drop strong references, leaving only map-owned Arcs. + // Saturate with retained strong references first so parallel tests cannot + // reclaim our fixture entries before we validate the reclaim path. + let prefix = format!("quota-reclaim-drop-{}", std::process::id()); + let mut retained = Vec::with_capacity(QUOTA_USER_LOCKS_MAX); for idx in 0..QUOTA_USER_LOCKS_MAX { - let _ = quota_user_lock(&format!("quota-reclaim-drop-{}-{idx}", std::process::id())); + retained.push(quota_user_lock(&format!("{prefix}-{idx}"))); } - assert_eq!(map.len(), QUOTA_USER_LOCKS_MAX); + + drop(retained); let overflow_user = format!("quota-reclaim-overflow-{}", std::process::id()); let overflow = quota_user_lock(&overflow_user); @@ -162,9 +154,7 @@ fn quota_lock_reclaims_unreferenced_entries_before_ephemeral_fallback() { #[test] fn quota_lock_saturated_same_user_must_not_return_distinct_locks() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -187,9 +177,7 @@ fn quota_lock_saturated_same_user_must_not_return_distinct_locks() { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn quota_lock_saturation_concurrent_same_user_never_overshoots_quota() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -240,9 +228,7 @@ async fn quota_lock_saturation_concurrent_same_user_never_overshoots_quota() { #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn quota_lock_saturation_stress_same_user_never_overshoots_quota() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); map.clear(); @@ -322,6 +308,24 @@ fn quota_error_classifier_rejects_plain_permission_denied() { assert!(!is_quota_io_error(&err)); } +#[test] +fn quota_lock_test_scope_recovers_after_guard_poison() { + let poison_result = std::thread::spawn(|| { + let _guard = super::quota_user_lock_test_scope(); + panic!("intentional test-only guard poison"); + }) + .join(); + assert!(poison_result.is_err(), "poison setup thread must panic"); + + let _guard = super::quota_user_lock_test_scope(); + let map = QUOTA_USER_LOCKS.get_or_init(DashMap::new); + map.clear(); + + let a = quota_user_lock("quota-lock-poison-recovery-user"); + let b = quota_user_lock("quota-lock-poison-recovery-user"); + assert!(Arc::ptr_eq(&a, &b)); +} + #[tokio::test] async fn quota_lock_integration_zero_quota_cuts_off_without_forwarding() { let stats = Arc::new(Stats::new()); diff --git a/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs b/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs index f68609a..1cd5920 100644 --- a/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs +++ b/src/proxy/tests/relay_quota_wake_liveness_regression_tests.rs @@ -18,10 +18,8 @@ fn saturate_lock_cache() -> Vec>> { retained } -fn quota_test_guard() -> std::sync::MutexGuard<'static, ()> { - super::quota_user_lock_test_guard() - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) +fn quota_test_guard() -> impl Drop { + super::quota_user_lock_test_scope() } #[tokio::test] diff --git a/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs b/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs index 666d90c..2dabaa3 100644 --- a/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs +++ b/src/proxy/tests/relay_quota_waker_storm_adversarial_tests.rs @@ -23,10 +23,8 @@ impl std::task::Wake for WakeCounter { } } -fn quota_test_guard() -> std::sync::MutexGuard<'static, ()> { - super::quota_user_lock_test_guard() - .lock() - .unwrap_or_else(|poisoned| poisoned.into_inner()) +fn quota_test_guard() -> impl Drop { + super::quota_user_lock_test_scope() } fn saturate_quota_user_locks() -> Vec>> { diff --git a/src/proxy/tests/relay_security_tests.rs b/src/proxy/tests/relay_security_tests.rs index 8f51cf3..b9b3478 100644 --- a/src/proxy/tests/relay_security_tests.rs +++ b/src/proxy/tests/relay_security_tests.rs @@ -31,9 +31,7 @@ impl std::task::Wake for WakeCounter { #[tokio::test] async fn quota_lock_contention_does_not_self_wake_pending_writer() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); map.clear(); @@ -72,9 +70,7 @@ async fn quota_lock_contention_does_not_self_wake_pending_writer() { #[tokio::test] async fn quota_lock_contention_writer_schedules_single_deferred_wake_until_lock_acquired() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); map.clear(); @@ -145,9 +141,7 @@ async fn quota_lock_contention_writer_schedules_single_deferred_wake_until_lock_ #[tokio::test] async fn quota_lock_contention_read_path_schedules_deferred_wake_for_liveness() { - let _guard = super::quota_user_lock_test_guard() - .lock() - .expect("quota lock test guard must be available"); + let _guard = super::quota_user_lock_test_scope(); let map = super::QUOTA_USER_LOCKS.get_or_init(dashmap::DashMap::new); map.clear();