Process-wide concurrency + Cancellation ownership fixes

src/proxy/traffic_limiter/lease.rs

@@ -2,20 +2,16 @@ use super::*;
 
 impl TrafficLease {
     fn current_binding(&self) -> Arc<TrafficLeaseBinding> {
-        let published_revision = self.limiter.published_revision.load(Ordering::Acquire);
+        let policy = self.limiter.policy.load();
         let current = self.binding.load_full();
-        if current.revision == published_revision {
+        if current.revision == policy.revision {
             return current;
         }
+        drop(policy);
 
-        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 _refresh = self.refresh.lock();
         let policy = self.limiter.policy.load_full();
+        let current = self.binding.load_full();
         if current.revision == policy.revision {
             return current;
         }
@@ -23,9 +19,6 @@ impl TrafficLease {
             .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
     }
 

src/proxy/traffic_limiter/limiter.rs

@@ -6,7 +6,6 @@ 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(),
@@ -15,16 +14,31 @@ impl TrafficLimiter {
         })
     }
 
+    #[cfg(test)]
     pub fn apply_policy(
         &self,
         user_limits: HashMap<String, RateLimitBps>,
         cidr_limits: HashMap<CidrRateLimitKey, RateLimitBps>,
     ) {
-        let policy_update = self.policy_update.lock();
-        // Revision wrap could otherwise let an old lease restore stale rates.
-        let Some(revision) = self.policy.load().revision.checked_add(1) else {
-            return;
-        };
+        let _ = self.apply_policy_inner(None, user_limits, cidr_limits);
+    }
+
+    /// Publishes policy only when the source runtime is not older than the active source.
+    pub(crate) fn apply_policy_from_source(
+        &self,
+        source_generation: u64,
+        user_limits: HashMap<String, RateLimitBps>,
+        cidr_limits: HashMap<CidrRateLimitKey, RateLimitBps>,
+    ) -> bool {
+        self.apply_policy_inner(Some(source_generation), user_limits, cidr_limits)
+    }
+
+    fn apply_policy_inner(
+        &self,
+        source_generation: Option<u64>,
+        user_limits: HashMap<String, RateLimitBps>,
+        cidr_limits: HashMap<CidrRateLimitKey, RateLimitBps>,
+    ) -> bool {
         let filtered_users = user_limits
             .into_iter()
             .filter(|(_, limit)| limit.up_bps > 0 || limit.down_bps > 0)
@@ -77,6 +91,17 @@ impl TrafficLimiter {
         let cidr_policy_entries =
             cidr_rule_keys.len() + cidr_auto_rules_v4.len() + cidr_auto_rules_v6.len();
 
+        let policy_update = self.policy_update.lock();
+        let current = self.policy.load_full();
+        if source_generation.is_some_and(|source| source < current.source_generation) {
+            return false;
+        }
+        // Revision wrap could otherwise let an old lease restore stale rates.
+        let Some(revision) = current.revision.checked_add(1) else {
+            return false;
+        };
+        let source_generation = source_generation.unwrap_or(current.source_generation);
+
         self.user_scope
             .policy_entries
             .store(filtered_users.len() as u64, Ordering::Relaxed);
@@ -86,6 +111,7 @@ impl TrafficLimiter {
 
         self.policy.store(Arc::new(PolicySnapshot {
             revision,
+            source_generation,
             user_limits: filtered_users,
             cidr_rules_v4,
             cidr_rules_v6,
@@ -93,10 +119,10 @@ impl TrafficLimiter {
             cidr_auto_rules_v6,
             cidr_rule_keys,
         }));
-        self.published_revision.store(revision, Ordering::Release);
 
         drop(policy_update);
         self.maybe_cleanup();
+        true
     }
 
     pub fn acquire_lease(
@@ -104,11 +130,8 @@ 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(),

src/proxy/traffic_limiter/tests.rs

@@ -5,6 +5,79 @@ fn rate(up_bps: u64, down_bps: u64) -> RateLimitBps {
     RateLimitBps { up_bps, down_bps }
 }
 
+#[test]
+fn stale_runtime_cannot_overwrite_newer_rate_policy() {
+    let limiter = TrafficLimiter::new();
+    let mut newer = HashMap::new();
+    newer.insert("alice".to_string(), rate(2_000, 3_000));
+    assert!(limiter.apply_policy_from_source(2, newer, HashMap::new()));
+
+    let mut stale = HashMap::new();
+    stale.insert("alice".to_string(), rate(1_000, 1_000));
+    assert!(!limiter.apply_policy_from_source(1, stale, HashMap::new()));
+
+    let policy = limiter.policy.load_full();
+    assert_eq!(policy.source_generation, 2);
+    assert_eq!(policy.user_limits["alice"].up_bps, 2_000);
+    assert_eq!(policy.user_limits["alice"].down_bps, 3_000);
+}
+
+#[test]
+fn active_runtime_can_publish_same_generation_rate_update() {
+    let limiter = TrafficLimiter::new();
+    assert!(limiter.apply_policy_from_source(4, HashMap::new(), HashMap::new()));
+    let mut updated = HashMap::new();
+    updated.insert("alice".to_string(), rate(4_000, 5_000));
+
+    assert!(limiter.apply_policy_from_source(4, updated, HashMap::new()));
+
+    let policy = limiter.policy.load_full();
+    assert_eq!(policy.source_generation, 4);
+    assert_eq!(policy.user_limits["alice"].up_bps, 4_000);
+}
+
+#[test]
+fn lease_acquisition_and_refresh_do_not_wait_for_policy_publication_lock() {
+    let limiter = TrafficLimiter::new();
+    let mut initial = HashMap::new();
+    initial.insert("alice".to_string(), rate(1_000, 1_000));
+    limiter.apply_policy(initial, HashMap::new());
+    let lease = limiter
+        .acquire_lease("alice", "203.0.113.7".parse().unwrap())
+        .unwrap();
+
+    let mut updated = HashMap::new();
+    updated.insert("alice".to_string(), rate(2_000, 2_000));
+    limiter.apply_policy(updated, HashMap::new());
+
+    let publication = limiter.policy_update.lock();
+    let (completed_tx, completed_rx) = std::sync::mpsc::channel();
+    let acquire_limiter = Arc::clone(&limiter);
+    let acquire_tx = completed_tx.clone();
+    let acquire = std::thread::spawn(move || {
+        let _lease = acquire_limiter
+            .acquire_lease("bob", "203.0.113.8".parse().unwrap())
+            .unwrap();
+        acquire_tx.send(()).unwrap();
+    });
+    let refresh = std::thread::spawn(move || {
+        let _ = lease.try_consume(RateDirection::Up, 1);
+        completed_tx.send(()).unwrap();
+    });
+
+    let first_completed = completed_rx
+        .recv_timeout(std::time::Duration::from_secs(1))
+        .is_ok();
+    let second_completed = completed_rx
+        .recv_timeout(std::time::Duration::from_secs(1))
+        .is_ok();
+    drop(publication);
+    acquire.join().unwrap();
+    refresh.join().unwrap();
+
+    assert!(first_completed && second_completed);
+}
+
 #[test]
 fn explicit_cidr_rule_wins_over_auto_template() {
     let limiter = TrafficLimiter::new();