Hardswap Invariants in tests + Quota fixes

src/proxy/traffic_limiter/buckets.rs

@@ -71,11 +71,11 @@ impl DirectionBucket {
     }
 
     pub(super) fn try_reserve_at(
-        &self,
+        self: &Arc<Self>,
         epoch: u64,
         cap: u64,
         requested: u64,
-    ) -> Option<DirectionDebit<'_>> {
+    ) -> Option<DirectionDebit> {
         if requested == 0 || cap == 0 || epoch > PACKED_EPOCH_MAX {
             return None;
         }
@@ -109,7 +109,7 @@ impl DirectionBucket {
             ) {
                 Ok(_) => {
                     return Some(DirectionDebit {
-                        bucket: self,
+                        bucket: Arc::clone(self),
                         epoch,
                         refundable: grant,
                     });
@@ -144,7 +144,7 @@ impl DirectionBucket {
     }
 }
 
-impl DirectionDebit<'_> {
+impl DirectionDebit {
     fn granted(&self) -> u64 {
         self.refundable
     }
@@ -168,7 +168,7 @@ impl DirectionDebit<'_> {
     }
 }
 
-impl Drop for DirectionDebit<'_> {
+impl Drop for DirectionDebit {
     fn drop(&mut self) {
         self.bucket.refund_at(self.epoch, self.refundable);
     }
@@ -178,8 +178,8 @@ impl UserBucket {
     pub(super) fn new(revision: u64, limits: RateLimitBps) -> Self {
         Self {
             rates: AtomicRatePair::new(revision, limits),
-            up: DirectionBucket::default(),
-            down: DirectionBucket::default(),
+            up: Arc::new(DirectionBucket::default()),
+            down: Arc::new(DirectionBucket::default()),
             active_leases: AtomicU64::new(0),
         }
     }
@@ -192,7 +192,7 @@ impl UserBucket {
         &self,
         direction: RateDirection,
         requested: u64,
-    ) -> (u64, Option<DirectionDebit<'_>>) {
+    ) -> (u64, Option<DirectionDebit>) {
         let cap_bps = self.rates.get(direction);
         if cap_bps == 0 {
             return (requested, None);
@@ -208,12 +208,12 @@ impl UserBucket {
 }
 
 impl CidrDirectionBucket {
-    pub(super) fn try_reserve<'a>(
-        &'a self,
-        user_state: &'a CidrUserDirectionState,
+    pub(super) fn try_reserve(
+        &self,
+        user_state: &CidrUserDirectionState,
         cap_epoch: u64,
         requested: u64,
-    ) -> (u64, Option<DirectionDebit<'a>>, Option<DirectionDebit<'a>>) {
+    ) -> (u64, Option<DirectionDebit>, Option<DirectionDebit>) {
         if requested == 0 || cap_epoch == 0 {
             return (0, None, None);
         }
@@ -260,7 +260,7 @@ impl CidrDirectionBucket {
 }
 
 impl CidrUserDirectionState {
-    pub(super) fn ensure_active(&self, epoch: u64, active_users: &DirectionBucket) -> bool {
+    pub(super) fn ensure_active(&self, epoch: u64, active_users: &Arc<DirectionBucket>) -> bool {
         if epoch > PACKED_EPOCH_MAX {
             return false;
         }
@@ -340,12 +340,12 @@ impl CidrBucket {
         });
     }
 
-    pub(super) fn try_reserve_for_user<'a>(
-        &'a self,
+    pub(super) fn try_reserve_for_user(
+        &self,
         direction: RateDirection,
-        share: &'a CidrUserShare,
+        share: &CidrUserShare,
         requested: u64,
-    ) -> (u64, Option<DirectionDebit<'a>>, Option<DirectionDebit<'a>>) {
+    ) -> (u64, Option<DirectionDebit>, Option<DirectionDebit>) {
         let cap_bps = self.rates.get(direction);
         if cap_bps == 0 {
             return (requested, None, None);

src/proxy/traffic_limiter/lease.rs

@@ -1,12 +1,41 @@
 use super::*;
 
 impl TrafficLease {
+    fn current_binding(&self) -> Arc<TrafficLeaseBinding> {
+        let published_revision = self.limiter.published_revision.load(Ordering::Acquire);
+        let current = self.binding.load_full();
+        if current.revision == published_revision {
+            return current;
+        }
+
+        let refresh = self.refresh.lock();
+        let published_revision = self.limiter.published_revision.load(Ordering::Acquire);
+        let current = self.binding.load_full();
+        if current.revision == published_revision {
+            return current;
+        }
+        let policy_update = self.limiter.policy_update.lock();
+        let policy = self.limiter.policy.load_full();
+        if current.revision == policy.revision {
+            return current;
+        }
+        let next = self
+            .limiter
+            .build_binding(&self.user, self.client_ip, &policy);
+        self.binding.store(Arc::clone(&next));
+        drop(policy_update);
+        drop(refresh);
+        self.limiter.maybe_cleanup();
+        next
+    }
+
     /// Reserves shaping budget until the associated I/O result is settled.
     pub(crate) fn try_reserve(
         &self,
         direction: RateDirection,
         requested: u64,
-    ) -> TrafficReservation<'_> {
+    ) -> TrafficReservation {
+        let binding = self.current_binding();
         if requested == 0 {
             return TrafficReservation {
                 result: TrafficConsumeResult {
@@ -14,6 +43,7 @@ impl TrafficLease {
                     blocked_user: false,
                     blocked_cidr: false,
                 },
+                _binding: binding,
                 user: None,
                 cidr: None,
                 cidr_user: None,
@@ -22,7 +52,7 @@ impl TrafficLease {
 
         let mut granted = requested;
         let mut user_debit = None;
-        if let Some(user_bucket) = self.user_bucket.as_ref() {
+        if let Some(user_bucket) = binding.user_bucket.as_ref() {
             let (user_granted, debit) = user_bucket.try_reserve(direction, granted);
             user_debit = debit;
             if user_granted == 0 {
@@ -33,6 +63,7 @@ impl TrafficLease {
                         blocked_user: true,
                         blocked_cidr: false,
                     },
+                    _binding: binding,
                     user: user_debit,
                     cidr: None,
                     cidr_user: None,
@@ -44,7 +75,7 @@ impl TrafficLease {
         let mut cidr_debit = None;
         let mut cidr_user_debit = None;
         if let (Some(cidr_bucket), Some(cidr_user_share)) =
-            (self.cidr_bucket.as_ref(), self.cidr_user_share.as_ref())
+            (binding.cidr_bucket.as_ref(), binding.cidr_user_share.as_ref())
         {
             let (cidr_granted, aggregate_debit, share_debit) =
                 cidr_bucket.try_reserve_for_user(direction, cidr_user_share, granted);
@@ -63,6 +94,7 @@ impl TrafficLease {
                         blocked_user: false,
                         blocked_cidr: true,
                     },
+                    _binding: binding,
                     user: user_debit,
                     cidr: cidr_debit,
                     cidr_user: cidr_user_debit,
@@ -77,6 +109,7 @@ impl TrafficLease {
                 blocked_user: false,
                 blocked_cidr: false,
             },
+            _binding: binding,
             user: user_debit,
             cidr: cidr_debit,
             cidr_user: cidr_user_debit,
@@ -105,7 +138,7 @@ impl TrafficLease {
     }
 }
 
-impl TrafficReservation<'_> {
+impl TrafficReservation {
     /// Returns the shaping decision associated with this reservation.
     pub(crate) fn result(&self) -> TrafficConsumeResult {
         self.result
@@ -126,7 +159,7 @@ impl TrafficReservation<'_> {
     }
 }
 
-impl Drop for TrafficLease {
+impl Drop for TrafficLeaseBinding {
     fn drop(&mut self) {
         if let Some(bucket) = self.user_bucket.as_ref() {
             decrement_atomic_saturating(&bucket.active_leases, 1);

src/proxy/traffic_limiter/limiter.rs

@@ -6,6 +6,7 @@ impl TrafficLimiter {
         Arc::new(Self {
             policy: ArcSwap::from_pointee(PolicySnapshot::default()),
             policy_update: ParkingMutex::new(()),
+            published_revision: AtomicU64::new(0),
             user_buckets: ShardedRegistry::new(REGISTRY_SHARDS),
             cidr_buckets: ShardedRegistry::new(REGISTRY_SHARDS),
             user_scope: ScopeMetrics::default(),
@@ -92,6 +93,7 @@ impl TrafficLimiter {
             cidr_auto_rules_v6,
             cidr_rule_keys,
         }));
+        self.published_revision.store(revision, Ordering::Release);
 
         drop(policy_update);
         self.maybe_cleanup();
@@ -102,7 +104,26 @@ impl TrafficLimiter {
         user: &str,
         client_ip: IpAddr,
     ) -> Option<Arc<TrafficLease>> {
+        let policy_update = self.policy_update.lock();
         let policy = self.policy.load_full();
+        let binding = self.build_binding(user, client_ip, &policy);
+        drop(policy_update);
+        self.maybe_cleanup();
+        Some(Arc::new(TrafficLease {
+            limiter: Arc::clone(self),
+            user: user.to_string(),
+            client_ip,
+            binding: ArcSwap::from(binding),
+            refresh: ParkingMutex::new(()),
+        }))
+    }
+
+    pub(super) fn build_binding(
+        self: &Arc<Self>,
+        user: &str,
+        client_ip: IpAddr,
+        policy: &PolicySnapshot,
+    ) -> Arc<TrafficLeaseBinding> {
         let mut user_bucket = None;
         if let Some(limit) = policy.user_limits.get(user).copied() {
             let bucket = self.user_buckets.get_or_insert_with(
@@ -144,18 +165,14 @@ impl TrafficLimiter {
             cidr_bucket = Some(bucket);
         }
 
-        if user_bucket.is_none() && cidr_bucket.is_none() {
-            return None;
-        }
-
-        self.maybe_cleanup();
-        Some(Arc::new(TrafficLease {
+        Arc::new(TrafficLeaseBinding {
             limiter: Arc::clone(self),
+            revision: policy.revision,
             user_bucket,
             cidr_bucket,
             cidr_user_key,
             cidr_user_share,
-        }))
+        })
     }
 
     pub fn metrics_snapshot(&self) -> TrafficLimiterMetricsSnapshot {

src/proxy/traffic_limiter/tests.rs

@@ -77,7 +77,7 @@ fn auto_cidr_bucket_key_canonicalizes_network_address() {
 
 #[test]
 fn refund_from_an_old_epoch_does_not_reduce_the_current_epoch() {
-    let bucket = DirectionBucket::default();
+    let bucket = Arc::new(DirectionBucket::default());
     let old_debit = bucket.try_reserve_at(7, 100, 80).unwrap();
     let current_debit = bucket.try_reserve_at(8, 100, 60).unwrap();
 
@@ -166,7 +166,7 @@ fn stale_policy_revision_cannot_restore_an_old_rate() {
 
 #[test]
 fn dropped_debit_refunds_only_its_packed_epoch() {
-    let bucket = DirectionBucket::default();
+    let bucket = Arc::new(DirectionBucket::default());
     let debit = bucket.try_reserve_at(11, 100, 80).unwrap();
 
     drop(debit);
@@ -229,9 +229,10 @@ fn dropped_traffic_reservation_refunds_user_and_cidr_debits() {
     let epoch = reservation.user.as_ref().unwrap().epoch;
     drop(reservation);
 
-    let user_bucket = lease.user_bucket.as_ref().unwrap();
-    let cidr_bucket = lease.cidr_bucket.as_ref().unwrap();
-    let cidr_user = lease.cidr_user_share.as_ref().unwrap();
+    let binding = lease.binding.load_full();
+    let user_bucket = binding.user_bucket.as_ref().unwrap();
+    let cidr_bucket = binding.cidr_bucket.as_ref().unwrap();
+    let cidr_user = binding.cidr_user_share.as_ref().unwrap();
     assert_eq!(user_bucket.down.used_at(epoch), Some(0));
     assert_eq!(cidr_bucket.down.used.used_at(epoch), Some(0));
     assert_eq!(cidr_user.down.used.used_at(epoch), Some(0));
@@ -252,7 +253,31 @@ fn partial_traffic_settlement_charges_only_committed_bytes() {
     reservation.settle_written(300);
 
     assert_eq!(
-        lease.user_bucket.as_ref().unwrap().down.used_at(epoch),
+        lease
+            .binding
+            .load_full()
+            .user_bucket
+            .as_ref()
+            .unwrap()
+            .down
+            .used_at(epoch),
         Some(300)
     );
 }
+
+#[test]
+fn active_lease_observes_policy_removal() {
+    let limiter = TrafficLimiter::new();
+    let mut user_limits = HashMap::new();
+    user_limits.insert("alice".to_string(), rate(1, 0));
+    limiter.apply_policy(user_limits, HashMap::new());
+    let lease = limiter
+        .acquire_lease("alice", "203.0.113.7".parse().unwrap())
+        .unwrap();
+    assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
+    assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 0);
+
+    limiter.apply_policy(HashMap::new(), HashMap::new());
+
+    assert_eq!(lease.try_consume(RateDirection::Up, 1).granted, 1);
+}