Process-wide concurrency + Cancellation ownership fixes

src/stats/mod.rs

@@ -22,9 +22,11 @@ use std::sync::atomic::{AtomicBool, AtomicU8, AtomicU64, AtomicUsize, Ordering};
 use std::time::Instant;
 
 pub(crate) use self::quota_store::{QuotaReservation, QuotaStore, UserQuotaHandle};
+pub(crate) use self::users::UserConnectionObservation;
 #[allow(unused_imports)]
 pub use self::replay::{ReplayChecker, ReplayStats};
 use self::telemetry::TelemetryPolicy;
+use crate::proxy::user_connection_authority::UserConnectionAuthority;
 pub use self::tls_fingerprints::TlsFingerprintSnapshotRow;
 use crate::config::MeWriterPickMode;
 
@@ -351,6 +353,7 @@ pub struct Stats {
     tls_fingerprints: tls_fingerprints::TlsFingerprintCollector,
     user_stats: DashMap<String, Arc<UserStats>>,
     quota_store: Arc<QuotaStore>,
+    connection_authority: Arc<UserConnectionAuthority>,
     user_stats_last_cleanup_epoch_secs: AtomicU64,
     start_time: parking_lot::RwLock<Option<Instant>>,
 }
@@ -417,12 +420,28 @@ impl UserStats {
 
 impl Stats {
     pub fn new() -> Self {
-        Self::with_quota_store(Arc::new(QuotaStore::default()))
+        Self::with_process_authorities(
+            Arc::new(QuotaStore::default()),
+            Arc::new(UserConnectionAuthority::default()),
+        )
     }
 
+    #[cfg(test)]
     pub(crate) fn with_quota_store(quota_store: Arc<QuotaStore>) -> Self {
+        Self::with_process_authorities(
+            quota_store,
+            Arc::new(UserConnectionAuthority::default()),
+        )
+    }
+
+    /// Creates generation telemetry around process-owned enforcement authorities.
+    pub(crate) fn with_process_authorities(
+        quota_store: Arc<QuotaStore>,
+        connection_authority: Arc<UserConnectionAuthority>,
+    ) -> Self {
         let stats = Self {
             quota_store,
+            connection_authority,
             ..Self::default()
         };
         stats.apply_telemetry_policy(TelemetryPolicy::default());
@@ -431,10 +450,16 @@ impl Stats {
         stats
     }
 
+    /// Returns the process-scoped quota authority for test runtime construction.
     #[cfg(test)]
     pub(crate) fn quota_store(&self) -> Arc<QuotaStore> {
         Arc::clone(&self.quota_store)
     }
+
+    /// Returns process-owned per-user connection admission.
+    pub(crate) fn connection_authority(&self) -> Arc<UserConnectionAuthority> {
+        Arc::clone(&self.connection_authority)
+    }
 }
 
 #[cfg(test)]

src/stats/tests.rs

@@ -13,6 +13,32 @@ fn test_stats_shared_counters() {
     assert_eq!(stats.get_connects_all(), 3);
 }
 
+#[test]
+fn runtime_stats_share_process_connection_admission_authority() {
+    let quota_store = Arc::new(QuotaStore::default());
+    let authority = Arc::new(UserConnectionAuthority::default());
+    let first = Stats::with_process_authorities(
+        Arc::clone(&quota_store),
+        Arc::clone(&authority),
+    );
+    let second = Stats::with_process_authorities(quota_store, authority);
+
+    let permit = first
+        .connection_authority()
+        .try_acquire("alice", Some(1))
+        .unwrap();
+    assert!(second
+        .connection_authority()
+        .try_acquire("alice", Some(1))
+        .is_none());
+
+    drop(permit);
+    assert!(second
+        .connection_authority()
+        .try_acquire("alice", Some(1))
+        .is_some());
+}
+
 #[test]
 fn test_telemetry_policy_disables_core_and_user_counters() {
     let stats = Stats::new();

src/stats/users.rs

@@ -1,5 +1,32 @@
 use super::*;
 
+/// Mirrors one live connection into optional per-user telemetry.
+#[must_use = "the observation must be retained for the connection lifetime"]
+pub(crate) struct UserConnectionObservation {
+    stats: Arc<UserStats>,
+}
+
+impl Drop for UserConnectionObservation {
+    fn drop(&mut self) {
+        decrement_current_connections(&self.stats.curr_connects);
+    }
+}
+
+fn decrement_current_connections(counter: &AtomicU64) {
+    let mut current = counter.load(Ordering::Relaxed);
+    while current != 0 {
+        match counter.compare_exchange_weak(
+            current,
+            current - 1,
+            Ordering::Relaxed,
+            Ordering::Relaxed,
+        ) {
+            Ok(_) => return,
+            Err(actual) => current = actual,
+        }
+    }
+}
+
 impl Stats {
     pub fn increment_user_connects(&self, user: &str) {
         if !self.telemetry_user_enabled() {
@@ -19,6 +46,25 @@ impl Stats {
         stats.curr_connects.fetch_add(1, Ordering::Relaxed);
     }
 
+    /// Starts an optional telemetry observation without owning admission policy.
+    pub(crate) fn observe_user_current_connection(
+        &self,
+        user: &str,
+    ) -> Option<UserConnectionObservation> {
+        if !self.telemetry_user_enabled() {
+            return None;
+        }
+        let stats = self.get_or_create_user_stats_handle(user);
+        self.touch_user_stats(stats.as_ref());
+        stats.curr_connects.fetch_add(1, Ordering::Relaxed);
+        Some(UserConnectionObservation { stats })
+    }
+
+    /// Returns the process-scoped count used by connection admission.
+    pub(crate) fn get_process_user_curr_connects(&self, user: &str) -> u64 {
+        self.connection_authority.active(user)
+    }
+
     pub fn try_acquire_user_curr_connects(&self, user: &str, limit: Option<u64>) -> bool {
         if !self.telemetry_user_enabled() {
             return true;
@@ -50,22 +96,7 @@ impl Stats {
     pub fn decrement_user_curr_connects(&self, user: &str) {
         if let Some(stats) = self.user_stats.get(user) {
             self.touch_user_stats(stats.value().as_ref());
-            let counter = &stats.curr_connects;
-            let mut current = counter.load(Ordering::Relaxed);
-            loop {
-                if current == 0 {
-                    break;
-                }
-                match counter.compare_exchange_weak(
-                    current,
-                    current - 1,
-                    Ordering::Relaxed,
-                    Ordering::Relaxed,
-                ) {
-                    Ok(_) => break,
-                    Err(actual) => current = actual,
-                }
-            }
+            decrement_current_connections(&stats.curr_connects);
         }
     }