Hardswap Invariants in tests + Quota fixes

src/maestro/control_plane.rs

@@ -120,6 +120,19 @@ impl ProcessControlPlane {
         Ok(())
     }
 
+    /// Registers work that must finish once accepted, even after shutdown cancellation starts.
+    pub(crate) fn spawn_completion<F>(&self, future: F) -> Result<(), F>
+    where
+        F: Future<Output = ()> + Send + 'static,
+    {
+        let Some(registration) = self.inner.admission.try_register() else {
+            return Err(future);
+        };
+        self.inner.tasks.spawn(future);
+        drop(registration);
+        Ok(())
+    }
+
     /// Closes task admission, cancels all owned work, and joins it within the deadline.
     pub(crate) async fn shutdown(&self, timeout: Duration) -> bool {
         let deadline = tokio::time::Instant::now() + timeout;
@@ -214,4 +227,31 @@ mod tests {
 
         assert!(scope.shutdown(Duration::from_secs(1)).await);
     }
+
+    #[tokio::test]
+    async fn shutdown_waits_for_accepted_completion_without_cancelling_it() {
+        let scope = ProcessControlPlane::new();
+        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
+        let completed = Arc::new(AtomicBool::new(false));
+        let completed_task = completed.clone();
+        assert!(
+            scope
+                .spawn_completion(async move {
+                    let _ = release_rx.await;
+                    completed_task.store(true, Ordering::Release);
+                })
+                .is_ok()
+        );
+
+        let shutdown_scope = scope.clone();
+        let shutdown =
+            tokio::spawn(async move { shutdown_scope.shutdown(Duration::from_secs(1)).await });
+        tokio::task::yield_now().await;
+        assert!(!shutdown.is_finished());
+        assert!(!completed.load(Ordering::Acquire));
+
+        release_tx.send(()).unwrap();
+        assert!(shutdown.await.unwrap());
+        assert!(completed.load(Ordering::Acquire));
+    }
 }

src/maestro/orchestrator.rs

@@ -12,6 +12,7 @@ use crate::network::probe::{decide_network_capabilities, log_probe_result, run_p
 use crate::proxy::direct_buffer_budget::{DirectBufferBudget, resolve_direct_buffer_hard_limit};
 use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
 use crate::proxy::shared_state::ProxySharedState;
+use crate::proxy::user_admission::UserAdmissionAuthority;
 use crate::startup::{COMPONENT_API_BOOTSTRAP, COMPONENT_NETWORK_PROBE};
 use crate::stats::telemetry::TelemetryPolicy;
 use crate::stats::{QuotaStore, Stats};
@@ -106,9 +107,17 @@ pub(super) async fn run_telemt_core(
         configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes,
         "Direct relay buffer budget initialized"
     );
-    let shared_state =
-        ProxySharedState::new_with_direct_buffer_budget(direct_buffer_budget.clone());
-    shared_state.apply_user_config(&config.access.users, &config.access.user_enabled);
+    let user_admission = UserAdmissionAuthority::new_with_quota_store(quota_store.clone());
+    let shared_state = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
+        direct_buffer_budget.clone(),
+        user_admission,
+    );
+    let _ = shared_state.activate_user_config_source(
+        1,
+        None,
+        &config.access.users,
+        &config.access.user_enabled,
+    );
     shared_state.traffic_limiter.apply_policy(
         config.access.user_rate_limits.clone(),
         config.access.cidr_rate_limits.clone(),

src/maestro/reload_supervisor.rs

@@ -303,8 +303,9 @@ impl ReloadSupervisor {
         let replaced = {
             let listener_manager = self.listener_manager.lock().await;
             let config = new_runtime.config();
-            let _ = new_runtime.proxy_shared.apply_user_config_if_epoch(
-                user_admission_epoch,
+            let _ = new_runtime.proxy_shared.activate_user_config_source(
+                new_runtime.id,
+                Some(user_admission_epoch),
                 &config.access.users,
                 &config.access.user_enabled,
             );

src/maestro/runtime_build.rs

@@ -192,6 +192,7 @@ pub(crate) async fn prepare_runtime(
     let max_connections = Arc::new(Semaphore::new(max_connections_limit));
     let (config_watcher_activation, config_watcher_activation_rx) = watch::channel(false);
     let watches = runtime_tasks::spawn_runtime_tasks(
+        generation_id,
         &config,
         config_path,
         &probe,

src/maestro/runtime_startup.rs

@@ -229,6 +229,7 @@ pub(super) async fn prepare_runtime(
     }
 
     let runtime_watches = runtime_tasks::spawn_runtime_tasks(
+        1,
         &config,
         config_path,
         probe,

src/maestro/runtime_tasks.rs

@@ -90,6 +90,7 @@ impl RuntimeLogFilter {
 
 #[allow(clippy::too_many_arguments)]
 pub(crate) async fn spawn_runtime_tasks(
+    generation_id: u64,
     config: &Arc<ProxyConfig>,
     config_path: &Path,
     probe: &NetworkProbe,
@@ -288,9 +289,14 @@ pub(crate) async fn spawn_runtime_tasks(
                 break;
             }
             let cfg = config_rx_user_enabled.borrow_and_update().clone();
-            for (user, cancelled) in shared_user_enabled
-                .apply_user_config(&cfg.access.users, &cfg.access.user_enabled)
-            {
+            let Some(cancelled_users) = shared_user_enabled.apply_user_config_from_source(
+                generation_id,
+                &cfg.access.users,
+                &cfg.access.user_enabled,
+            ) else {
+                continue;
+            };
+            for (user, cancelled) in cancelled_users {
                 if cancelled > 0 {
                     info!(
                         user = %user,