Fixes in tests

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/proxy/relay.rs

@@ -669,3 +669,7 @@ mod relay_quota_extended_attack_surface_security_tests;
 #[cfg(test)]
 #[path = "tests/relay_watchdog_delta_security_tests.rs"]
 mod relay_watchdog_delta_security_tests;
+
+#[cfg(test)]
+#[path = "tests/relay_atomic_quota_invariant_tests.rs"]
+mod relay_atomic_quota_invariant_tests;

src/proxy/tests/relay_quota_boundary_blackhat_tests.rs

@@ -29,6 +29,11 @@ async fn read_available<R: AsyncRead + Unpin>(reader: &mut R, budget: Duration)
     total
 }
 
+fn preload_user_quota(stats: &Stats, user: &str, bytes: u64) {
+    let user_stats = stats.get_or_create_user_stats_handle(user);
+    stats.quota_charge_post_write(user_stats.as_ref(), bytes);
+}
+
 #[tokio::test]
 async fn integration_full_duplex_exact_budget_then_hard_cutoff() {
     let stats = Arc::new(Stats::new());
@@ -102,14 +107,14 @@ async fn integration_full_duplex_exact_budget_then_hard_cutoff() {
         relay_result,
         Err(ProxyError::DataQuotaExceeded { ref user }) if user == "quota-full-duplex-boundary-user"
     ));
-    assert!(stats.get_user_total_octets(user) <= 10);
+    assert!(stats.get_user_quota_used(user) <= 10);
 }
 
 #[tokio::test]
 async fn negative_preloaded_quota_blocks_both_directions_immediately() {
     let stats = Arc::new(Stats::new());
     let user = "quota-preloaded-cutoff-user";
-    stats.add_user_octets_from(user, 5);
+    preload_user_quota(stats.as_ref(), user, 5);
 
     let (mut client_peer, relay_client) = duplex(2048);
     let (relay_server, mut server_peer) = duplex(2048);
@@ -154,7 +159,7 @@ async fn negative_preloaded_quota_blocks_both_directions_immediately() {
         relay_result,
         Err(ProxyError::DataQuotaExceeded { .. })
     ));
-    assert!(stats.get_user_total_octets(user) <= 5);
+    assert!(stats.get_user_quota_used(user) <= 5);
 }
 
 #[tokio::test]
@@ -212,7 +217,7 @@ async fn edge_quota_one_bidirectional_race_allows_at_most_one_forwarded_octet()
         relay_result,
         Err(ProxyError::DataQuotaExceeded { .. })
     ));
-    assert!(stats.get_user_total_octets(user) <= 1);
+    assert!(stats.get_user_quota_used(user) <= 1);
 }
 
 #[tokio::test]
@@ -277,7 +282,7 @@ async fn adversarial_blackhat_alternating_fragmented_jitter_never_overshoots_glo
         delivered_to_server + delivered_to_client <= quota as usize,
         "combined forwarded bytes must never exceed configured quota"
     );
-    assert!(stats.get_user_total_octets(user) <= quota);
+    assert!(stats.get_user_quota_used(user) <= quota);
 }
 
 #[tokio::test]
@@ -356,7 +361,7 @@ async fn light_fuzz_randomized_schedule_preserves_quota_and_forwarded_byte_invar
             "fuzz case {case}: forwarded bytes must not exceed quota"
         );
         assert!(
-            stats.get_user_total_octets(&user) <= quota,
+            stats.get_user_quota_used(&user) <= quota,
             "fuzz case {case}: accounted bytes must not exceed quota"
         );
     }
@@ -451,7 +456,7 @@ async fn stress_multi_relay_same_user_mixed_direction_jitter_respects_global_quo
     }
 
     assert!(
-        stats.get_user_total_octets(user) <= quota,
+        stats.get_user_quota_used(user) <= quota,
         "global per-user quota must hold under concurrent mixed-direction relay stress"
     );
     assert!(

src/stats/mod.rs

@@ -2580,6 +2580,56 @@ mod tests {
         assert_eq!(user_stats.quota_used(), limit);
     }
 
+    #[test]
+    fn test_quota_reserve_200x_1k_reaches_100k_without_overshoot() {
+        let user_stats = Arc::new(UserStats::default());
+        let successes = Arc::new(AtomicU64::new(0));
+        let failures = Arc::new(AtomicU64::new(0));
+        let attempts = 200usize;
+        let reserve_bytes = 1_024u64;
+        let limit = 100 * 1_024u64;
+        let mut workers = Vec::with_capacity(attempts);
+
+        for _ in 0..attempts {
+            let user_stats = user_stats.clone();
+            let successes = successes.clone();
+            let failures = failures.clone();
+            workers.push(std::thread::spawn(move || {
+                loop {
+                    match user_stats.quota_try_reserve(reserve_bytes, limit) {
+                        Ok(_) => {
+                            successes.fetch_add(1, Ordering::Relaxed);
+                            return;
+                        }
+                        Err(QuotaReserveError::LimitExceeded) => {
+                            failures.fetch_add(1, Ordering::Relaxed);
+                            return;
+                        }
+                        Err(QuotaReserveError::Contended) => {
+                            std::hint::spin_loop();
+                        }
+                    }
+                }
+            }));
+        }
+
+        for worker in workers {
+            worker.join().expect("reservation worker must finish");
+        }
+
+        assert_eq!(
+            successes.load(Ordering::Relaxed),
+            100,
+            "exactly 100 reservations of 1 KiB must fit into a 100 KiB quota"
+        );
+        assert_eq!(
+            failures.load(Ordering::Relaxed),
+            100,
+            "remaining workers must fail once quota is fully reserved"
+        );
+        assert_eq!(user_stats.quota_used(), limit);
+    }
+
     #[test]
     fn test_quota_used_is_authoritative_and_independent_from_octets_telemetry() {
         let stats = Stats::new();
@@ -2594,6 +2644,33 @@ mod tests {
         assert_eq!(stats.get_user_total_octets(user), 5);
         assert_eq!(stats.get_user_quota_used(user), 7);
     }
+
+    #[test]
+    fn test_cached_handle_survives_map_cleanup_until_last_drop() {
+        let stats = Stats::new();
+        let user = "quota-handle-lifetime-user";
+        let user_stats = stats.get_or_create_user_stats_handle(user);
+        let weak = Arc::downgrade(&user_stats);
+
+        stats.user_stats.remove(user);
+        assert!(
+            stats.user_stats.get(user).is_none(),
+            "map cleanup should remove idle entry"
+        );
+        assert!(
+            weak.upgrade().is_some(),
+            "cached handle must keep user stats object alive after map removal"
+        );
+
+        stats.quota_charge_post_write(user_stats.as_ref(), 3);
+        assert_eq!(user_stats.quota_used(), 3);
+
+        drop(user_stats);
+        assert!(
+            weak.upgrade().is_none(),
+            "user stats object must be dropped after the last cached handle is released"
+        );
+    }
 }
 
 #[cfg(test)]