Hardswap Invariants in tests + Quota fixes

src/stats/mod.rs

@@ -21,7 +21,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
 use std::time::Instant;
 
-pub(crate) use self::quota_store::{QuotaReservation, QuotaStore};
+pub(crate) use self::quota_store::{QuotaReservation, QuotaStore, UserQuotaHandle};
 #[allow(unused_imports)]
 pub use self::replay::{ReplayChecker, ReplayStats};
 use self::telemetry::TelemetryPolicy;
@@ -430,6 +430,11 @@ impl Stats {
         *stats.start_time.write() = Some(Instant::now());
         stats
     }
+
+    #[cfg(test)]
+    pub(crate) fn quota_store(&self) -> Arc<QuotaStore> {
+        Arc::clone(&self.quota_store)
+    }
 }
 
 #[cfg(test)]

src/stats/quota_store.rs

@@ -4,13 +4,38 @@ use std::sync::atomic::{AtomicU64, Ordering};
 
 use arc_swap::ArcSwap;
 use dashmap::DashMap;
+use parking_lot::Mutex;
 
 use super::{QuotaReserveError, UserQuotaSnapshot};
+use crate::proxy::user_admission::UserIncarnation;
 
 /// Process-scoped per-user quota accounting shared by runtime generations.
 #[derive(Default)]
 pub struct QuotaStore {
-    users: DashMap<String, Arc<UserQuotaCounters>>,
+    users: DashMap<String, Arc<QuotaUserSlot>>,
+}
+
+struct QuotaUserSlot {
+    state: Mutex<QuotaSlotState>,
+}
+
+#[derive(Default)]
+struct QuotaSlotState {
+    high_water: UserIncarnation,
+    current: Option<QuotaAccount>,
+    startup_seed: Option<UserQuotaSnapshot>,
+}
+
+struct QuotaAccount {
+    incarnation: UserIncarnation,
+    counters: Arc<UserQuotaCounters>,
+}
+
+/// Exact quota ownership pinned to one authenticated user incarnation.
+#[derive(Clone)]
+pub(crate) struct UserQuotaHandle {
+    incarnation: UserIncarnation,
+    counters: Arc<UserQuotaCounters>,
 }
 
 /// Atomically replaceable quota state for one configured user.
@@ -32,30 +57,164 @@ pub(crate) struct QuotaReservation {
 }
 
 impl QuotaStore {
-    pub(crate) fn user(&self, user: &str) -> Arc<UserQuotaCounters> {
+    fn slot(&self, user: &str) -> Arc<QuotaUserSlot> {
         if let Some(existing) = self.users.get(user) {
             return Arc::clone(existing.value());
         }
         Arc::clone(
             self.users
                 .entry(user.to_string())
-                .or_insert_with(|| Arc::new(UserQuotaCounters::default()))
+                .or_insert_with(|| {
+                    Arc::new(QuotaUserSlot {
+                        state: Mutex::new(QuotaSlotState::default()),
+                    })
+                })
                 .value(),
         )
     }
 
+    pub(crate) fn current_or_legacy_handle(&self, user: &str) -> UserQuotaHandle {
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        if let Some(account) = &state.current {
+            return UserQuotaHandle {
+                incarnation: account.incarnation,
+                counters: Arc::clone(&account.counters),
+            };
+        }
+        let seed = state.startup_seed.take().unwrap_or(UserQuotaSnapshot {
+            used_bytes: 0,
+            last_reset_epoch_secs: 0,
+        });
+        let counters = Arc::new(UserQuotaCounters::from_snapshot(&seed));
+        let incarnation = state.high_water;
+        state.current = Some(QuotaAccount {
+            incarnation,
+            counters: Arc::clone(&counters),
+        });
+        UserQuotaHandle {
+            incarnation,
+            counters,
+        }
+    }
+
+    pub(crate) fn user(&self, user: &str) -> Arc<UserQuotaCounters> {
+        self.current_or_legacy_handle(user).counters
+    }
+
+    /// Returns the quota account owned by the exact current incarnation.
+    pub(crate) fn handle_exact(
+        &self,
+        user: &str,
+        incarnation: UserIncarnation,
+    ) -> Option<UserQuotaHandle> {
+        let slot = self.users.get(user)?;
+        let state = slot.state.lock();
+        let account = state.current.as_ref()?;
+        (account.incarnation == incarnation).then(|| UserQuotaHandle {
+            incarnation,
+            counters: Arc::clone(&account.counters),
+        })
+    }
+
+    /// Creates a quota account for a new username lifetime without inheriting a retired account.
+    pub(crate) fn activate_fresh(&self, user: &str, incarnation: UserIncarnation) {
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        if incarnation <= state.high_water {
+            return;
+        }
+        let snapshot = if state.high_water == 0 {
+            state
+                .current
+                .as_ref()
+                .map(|account| account.counters.snapshot())
+                .or_else(|| state.startup_seed.take())
+        } else {
+            state.startup_seed = None;
+            None
+        }
+        .unwrap_or(UserQuotaSnapshot {
+            used_bytes: 0,
+            last_reset_epoch_secs: 0,
+        });
+        state.high_water = state.high_water.max(incarnation);
+        state.current = Some(QuotaAccount {
+            incarnation,
+            counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)),
+        });
+    }
+
+    /// Advances a credential incarnation while preserving usage captured at the transition.
+    pub(crate) fn advance_preserving_usage(
+        &self,
+        user: &str,
+        incarnation: UserIncarnation,
+    ) {
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        if incarnation <= state.high_water {
+            return;
+        }
+        let snapshot = state
+            .current
+            .as_ref()
+            .map(|account| account.counters.snapshot())
+            .or_else(|| state.startup_seed.take())
+            .unwrap_or(UserQuotaSnapshot {
+                used_bytes: 0,
+                last_reset_epoch_secs: 0,
+            });
+        state.high_water = incarnation;
+        state.current = Some(QuotaAccount {
+            incarnation,
+            counters: Arc::new(UserQuotaCounters::from_snapshot(&snapshot)),
+        });
+    }
+
+    /// Retires quota ownership without affecting a newer incarnation.
+    pub(crate) fn retire_through(&self, user: &str, incarnation: UserIncarnation) {
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        if incarnation < state.high_water {
+            return;
+        }
+        state.high_water = incarnation;
+        state.current = None;
+        state.startup_seed = None;
+    }
+
     pub(crate) fn used(&self, user: &str) -> u64 {
-        self.users.get(user).map(|state| state.used()).unwrap_or(0)
+        self.users
+            .get(user)
+            .and_then(|slot| {
+                let state = slot.state.lock();
+                state
+                    .current
+                    .as_ref()
+                    .map(|account| account.counters.used())
+                    .or_else(|| state.startup_seed.as_ref().map(|seed| seed.used_bytes))
+            })
+            .unwrap_or(0)
     }
 
     pub(crate) fn load(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
-        let state = self.user(user);
-        state.replace(used_bytes, last_reset_epoch_secs);
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        let snapshot = UserQuotaSnapshot {
+            used_bytes,
+            last_reset_epoch_secs,
+        };
+        if let Some(account) = &state.current {
+            account.replace_from_snapshot(&snapshot);
+        } else {
+            state.startup_seed = Some(snapshot);
+        }
     }
 
     pub(crate) fn reset(&self, user: &str, now_epoch_secs: u64) -> UserQuotaSnapshot {
-        let state = self.user(user);
-        state.replace(0, now_epoch_secs);
+        let state = self.current_or_legacy_handle(user);
+        state.counters.replace(0, now_epoch_secs);
         UserQuotaSnapshot {
             used_bytes: 0,
             last_reset_epoch_secs: now_epoch_secs,
@@ -63,26 +222,28 @@ impl QuotaStore {
     }
 
     pub(crate) fn remove(&self, user: &str) {
-        self.users.remove(user);
+        let slot = self.slot(user);
+        let mut state = slot.state.lock();
+        state.high_water = state.high_water.saturating_add(1);
+        state.current = None;
+        state.startup_seed = None;
     }
 
     pub(crate) fn snapshot(&self) -> HashMap<String, UserQuotaSnapshot> {
         let mut out = HashMap::new();
         for entry in self.users.iter() {
-            let state = entry.value();
-            let generation = state.generation.load_full();
-            let used_bytes = generation.used_bytes.load(Ordering::Relaxed);
-            let last_reset_epoch_secs = generation.last_reset_epoch_secs;
-            if used_bytes == 0 && last_reset_epoch_secs == 0 {
+            let state = entry.value().state.lock();
+            let snapshot = if let Some(account) = state.current.as_ref() {
+                account.counters.snapshot()
+            } else if let Some(seed) = state.startup_seed.as_ref() {
+                seed.clone()
+            } else {
+                continue;
+            };
+            if snapshot.used_bytes == 0 && snapshot.last_reset_epoch_secs == 0 {
                 continue;
             }
-            out.insert(
-                entry.key().clone(),
-                UserQuotaSnapshot {
-                    used_bytes,
-                    last_reset_epoch_secs,
-                },
-            );
+            out.insert(entry.key().clone(), snapshot);
         }
         out
     }
@@ -100,6 +261,23 @@ impl Default for UserQuotaCounters {
 }
 
 impl UserQuotaCounters {
+    fn from_snapshot(snapshot: &UserQuotaSnapshot) -> Self {
+        Self {
+            generation: ArcSwap::from_pointee(QuotaGeneration {
+                used_bytes: AtomicU64::new(snapshot.used_bytes),
+                last_reset_epoch_secs: snapshot.last_reset_epoch_secs,
+            }),
+        }
+    }
+
+    fn snapshot(&self) -> UserQuotaSnapshot {
+        let generation = self.generation.load_full();
+        UserQuotaSnapshot {
+            used_bytes: generation.used_bytes.load(Ordering::Relaxed),
+            last_reset_epoch_secs: generation.last_reset_epoch_secs,
+        }
+    }
+
     fn replace(&self, used_bytes: u64, last_reset_epoch_secs: u64) {
         self.generation.store(Arc::new(QuotaGeneration {
             used_bytes: AtomicU64::new(used_bytes),
@@ -150,6 +328,39 @@ impl UserQuotaCounters {
     }
 }
 
+impl QuotaAccount {
+    fn replace_from_snapshot(&self, snapshot: &UserQuotaSnapshot) {
+        self.counters
+            .replace(snapshot.used_bytes, snapshot.last_reset_epoch_secs);
+    }
+}
+
+impl UserQuotaHandle {
+    /// Returns the immutable incarnation owned by this handle.
+    pub(crate) fn incarnation(&self) -> UserIncarnation {
+        self.incarnation
+    }
+
+    #[inline]
+    pub(crate) fn used(&self) -> u64 {
+        self.counters.used()
+    }
+
+    #[inline]
+    pub(crate) fn charge(&self, bytes: u64) -> u64 {
+        self.counters.charge(bytes)
+    }
+
+    #[inline]
+    pub(crate) fn try_reserve(
+        &self,
+        bytes: u64,
+        limit: u64,
+    ) -> Result<QuotaReservation, QuotaReserveError> {
+        self.counters.try_reserve(bytes, limit)
+    }
+}
+
 impl QuotaReservation {
     /// Returns the number of bytes held by this reservation.
     pub(crate) fn reserved_bytes(&self) -> u64 {
@@ -255,4 +466,63 @@ mod tests {
             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 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);
+    }
 }

src/stats/replay.rs

@@ -73,6 +73,7 @@ pub struct ReplayChecker {
     checks: AtomicU64,
     hits: AtomicU64,
     additions: AtomicU64,
+    capacity_rejections: AtomicU64,
     cleanups: AtomicU64,
     next_claim_token: AtomicU64,
 }
@@ -141,21 +142,25 @@ impl ReplayShard {
         self.cache.get(key).is_some() || self.pending.contains_key(key)
     }
 
-    fn add_owned(&mut self, key: ReplayKey, now: Instant, window: Duration) {
+    fn add_owned(&mut self, key: ReplayKey, now: Instant, window: Duration) -> bool {
         if window.is_zero() {
-            return;
+            return true;
         }
         self.cleanup(now, window);
         if self.cache.peek(key.as_slice()).is_some() || self.pending.contains_key(key.as_slice()) {
-            return;
+            return true;
         }
-        while self.queue.len() >= self.capacity {
+        while self.cache.len().saturating_add(self.pending.len()) >= self.capacity {
+            if self.queue.is_empty() {
+                return false;
+            }
             self.evict_queue_front();
         }
 
         let seq = self.next_seq();
         self.cache.put(key.clone(), ReplayEntry { seq });
         self.queue.push_back((now, key, seq));
+        true
     }
 
     fn claim_owned(
@@ -218,7 +223,12 @@ impl TlsReplayClaim<'_> {
         if !shard.remove_pending(key.as_slice(), self.token) {
             return false;
         }
-        shard.add_owned(key, Instant::now(), self.checker.tls_window);
+        if !shard.add_owned(key, Instant::now(), self.checker.tls_window) {
+            self.checker
+                .capacity_rejections
+                .fetch_add(1, Ordering::Relaxed);
+            return false;
+        }
         self.checker.additions.fetch_add(1, Ordering::Relaxed);
         self.reserved = false;
         true
@@ -261,6 +271,7 @@ impl ReplayChecker {
             checks: AtomicU64::new(0),
             hits: AtomicU64::new(0),
             additions: AtomicU64::new(0),
+            capacity_rejections: AtomicU64::new(0),
             cleanups: AtomicU64::new(0),
             next_claim_token: AtomicU64::new(1),
         }
@@ -304,11 +315,15 @@ impl ReplayChecker {
         let found = shard.check(data, now, window);
         if found {
             self.hits.fetch_add(1, Ordering::Relaxed);
-        } else {
-            shard.add_owned(owned_key, now, window);
-            self.additions.fetch_add(1, Ordering::Relaxed);
+            return true;
+        }
+        if shard.add_owned(owned_key, now, window) {
+            self.additions.fetch_add(1, Ordering::Relaxed);
+            false
+        } else {
+            self.capacity_rejections.fetch_add(1, Ordering::Relaxed);
+            true
         }
-        found
     }
 
     fn check_only_internal(
@@ -328,11 +343,14 @@ impl ReplayChecker {
     }
 
     fn add_only(&self, data: &[u8], shards: &[Mutex<ReplayShard>], window: Duration) {
-        self.additions.fetch_add(1, Ordering::Relaxed);
         let idx = self.get_shard_idx(data);
         let owned_key = ReplayKey::from_slice(data);
         let mut shard = shards[idx].lock();
-        shard.add_owned(owned_key, Instant::now(), window);
+        if shard.add_owned(owned_key, Instant::now(), window) {
+            self.additions.fetch_add(1, Ordering::Relaxed);
+        } else {
+            self.capacity_rejections.fetch_add(1, Ordering::Relaxed);
+        }
     }
 
     pub fn check_and_add_handshake(&self, data: &[u8]) -> bool {
@@ -394,12 +412,12 @@ impl ReplayChecker {
         let mut total_queue_len = 0;
         for shard in &self.handshake_shards {
             let s = shard.lock();
-            total_entries += s.cache.len();
+            total_entries += s.len();
             total_queue_len += s.queue.len();
         }
         for shard in &self.tls_shards {
             let s = shard.lock();
-            total_entries += s.cache.len();
+            total_entries += s.len();
             total_queue_len += s.queue.len();
         }
 
@@ -409,6 +427,7 @@ impl ReplayChecker {
             total_checks: self.checks.load(Ordering::Relaxed),
             total_hits: self.hits.load(Ordering::Relaxed),
             total_additions: self.additions.load(Ordering::Relaxed),
+            total_capacity_rejections: self.capacity_rejections.load(Ordering::Relaxed),
             total_cleanups: self.cleanups.load(Ordering::Relaxed),
             num_shards: self.handshake_shards.len() + self.tls_shards.len(),
             window_secs: self.window.as_secs(),
@@ -459,6 +478,7 @@ pub struct ReplayStats {
     pub total_checks: u64,
     pub total_hits: u64,
     pub total_additions: u64,
+    pub total_capacity_rejections: u64,
     pub total_cleanups: u64,
     pub num_shards: usize,
     pub window_secs: u64,
@@ -481,3 +501,24 @@ impl ReplayStats {
         }
     }
 }
+
+#[cfg(test)]
+mod capacity_tests {
+    use super::*;
+
+    #[test]
+    fn committed_and_pending_entries_share_one_shard_capacity() {
+        let capacity = NonZeroUsize::new(2).unwrap();
+        let mut shard = ReplayShard::new(capacity);
+        let now = Instant::now();
+        let window = Duration::from_secs(60);
+        assert!(shard.claim_owned(ReplayKey::from_slice(b"pending-a"), now, window, 1));
+        assert!(shard.claim_owned(ReplayKey::from_slice(b"pending-b"), now, window, 2));
+
+        shard.add_owned(ReplayKey::from_slice(b"committed"), now, window);
+
+        assert!(shard.len() <= capacity.get());
+        assert!(shard.pending.contains_key(b"pending-a".as_slice()));
+        assert!(shard.pending.contains_key(b"pending-b".as_slice()));
+    }
+}

src/stats/users.rs

@@ -122,6 +122,27 @@ impl Stats {
         self.quota_store.used(user)
     }
 
+    /// Returns quota ownership for the exact authenticated user incarnation.
+    pub(crate) fn quota_handle_for_incarnation(
+        &self,
+        user: &str,
+        incarnation: crate::proxy::user_admission::UserIncarnation,
+    ) -> Option<UserQuotaHandle> {
+        if let Some(handle) = self.quota_store.handle_exact(user, incarnation) {
+            return Some(handle);
+        }
+        if incarnation != 0 {
+            return None;
+        }
+        let handle = self.quota_store.current_or_legacy_handle(user);
+        (handle.incarnation() == incarnation).then_some(handle)
+    }
+
+    /// Returns the currently published quota owner for compatibility relay entrypoints.
+    pub(crate) fn current_user_quota_handle(&self, user: &str) -> UserQuotaHandle {
+        self.quota_store.current_or_legacy_handle(user)
+    }
+
     pub fn load_user_quota_state(&self, user: &str, used_bytes: u64, last_reset_epoch_secs: u64) {
         self.quota_store
             .load(user, used_bytes, last_reset_epoch_secs);