Process-wide concurrency + Cancellation ownership fixes

src/maestro/generation.rs

@@ -22,6 +22,11 @@ use crate::tls_front::TlsFrontCache;
 use crate::transport::UpstreamManager;
 use crate::transport::middle_proxy::MePool;
 
+// Cancellation guards preserve runtime ownership across preparation and drain futures.
+mod lifecycle;
+pub(crate) use lifecycle::RuntimeTaskScopePreparationGuard;
+use lifecycle::SessionDrainCancellationGuard;
+
 const SESSION_STOP_TIMEOUT: Duration = Duration::from_secs(5);
 const BACKGROUND_STOP_TIMEOUT: Duration = Duration::from_secs(5);
 const SESSION_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1);
@@ -145,12 +150,17 @@ impl RuntimeTaskScope {
         self.cancel.clone()
     }
 
-    /// Cancels the scope and waits within the bounded background-task budget.
-    pub(crate) async fn stop(&self) {
+    /// Synchronously closes task admission and signals every tracked task.
+    pub(crate) fn begin_stop(&self) {
         self.admission.close();
-        self.admission.wait_for_registrations().await;
         self.cancel.cancel();
         self.tracker.close();
+    }
+
+    /// Cancels the scope and waits within the bounded background-task budget.
+    pub(crate) async fn stop(&self) {
+        self.begin_stop();
+        self.admission.wait_for_registrations().await;
         let _ = tokio::time::timeout(BACKGROUND_STOP_TIMEOUT, self.tracker.wait()).await;
     }
 }
@@ -298,24 +308,31 @@ impl RuntimeGeneration {
     /// Waits for registered sessions and cancels them when the deadline expires.
     pub(crate) async fn drain_sessions(&self, timeout: Duration) -> bool {
         self.stop_accepting_sessions();
+        let mut cancellation_guard = SessionDrainCancellationGuard::new(self);
         self.session_admission.wait_for_registrations().await;
         self.sessions.close();
         if tokio::time::timeout(timeout, self.sessions.wait())
             .await
             .is_ok()
         {
+            cancellation_guard.disarm();
             return true;
         }
+        cancellation_guard.disarm();
         self.stop_sessions().await;
         false
     }
 
-    /// Cancels all sessions and waits within the bounded session-stop budget.
-    pub(crate) async fn stop_sessions(&self) {
+    fn begin_stop_sessions(&self) {
         self.stop_accepting_sessions();
-        self.session_admission.wait_for_registrations().await;
         self.session_cancel.cancel();
         self.sessions.close();
+    }
+
+    /// Cancels all sessions and waits within the bounded session-stop budget.
+    pub(crate) async fn stop_sessions(&self) {
+        self.begin_stop_sessions();
+        self.session_admission.wait_for_registrations().await;
         let _ = tokio::time::timeout(SESSION_STOP_TIMEOUT, self.sessions.wait()).await;
     }
 
@@ -335,6 +352,8 @@ impl RuntimeGeneration {
 
 impl Drop for RuntimeGeneration {
     fn drop(&mut self) {
+        self.background_tasks.begin_stop();
+        self.begin_stop_sessions();
         if let Some(pool) = self.me_pool.as_ref() {
             pool.begin_shutdown();
         }

src/maestro/generation/lifecycle.rs

@@ -0,0 +1,157 @@
+use super::*;
+
+/// Cancels tasks if runtime preparation exits before ownership reaches a generation.
+#[must_use = "runtime preparation guards must be disarmed after ownership transfer"]
+pub(crate) struct RuntimeTaskScopePreparationGuard {
+    scope: RuntimeTaskScope,
+    armed: bool,
+}
+
+impl RuntimeTaskScopePreparationGuard {
+    /// Arms cancellation for a newly created, not-yet-published task scope.
+    pub(crate) fn new(scope: RuntimeTaskScope) -> Self {
+        Self { scope, armed: true }
+    }
+
+    /// Confirms that a runtime generation now owns the task scope.
+    pub(crate) fn disarm(mut self) {
+        self.armed = false;
+    }
+}
+
+impl Drop for RuntimeTaskScopePreparationGuard {
+    fn drop(&mut self) {
+        if self.armed {
+            self.scope.begin_stop();
+        }
+    }
+}
+
+pub(super) struct SessionDrainCancellationGuard<'a> {
+    generation: &'a RuntimeGeneration,
+    armed: bool,
+}
+
+impl<'a> SessionDrainCancellationGuard<'a> {
+    pub(super) fn new(generation: &'a RuntimeGeneration) -> Self {
+        Self {
+            generation,
+            armed: true,
+        }
+    }
+
+    pub(super) fn disarm(&mut self) {
+        self.armed = false;
+    }
+}
+
+impl Drop for SessionDrainCancellationGuard<'_> {
+    fn drop(&mut self) {
+        if self.armed {
+            self.generation.begin_stop_sessions();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    struct NotifyOnDrop(Arc<Notify>);
+
+    impl Drop for NotifyOnDrop {
+        fn drop(&mut self) {
+            self.0.notify_one();
+        }
+    }
+
+    #[tokio::test]
+    async fn preparation_guard_drop_cancels_scope_and_rejects_late_spawn() {
+        let scope = RuntimeTaskScope::new();
+        let guard = RuntimeTaskScopePreparationGuard::new(scope.clone());
+        drop(guard);
+
+        assert!(scope.cancellation_token().is_cancelled());
+        let ran = Arc::new(AtomicUsize::new(0));
+        let ran_task = Arc::clone(&ran);
+        scope.spawn(async move {
+            ran_task.fetch_add(1, Ordering::AcqRel);
+        });
+        tokio::task::yield_now().await;
+        assert_eq!(ran.load(Ordering::Acquire), 0);
+    }
+
+    #[tokio::test]
+    async fn preparation_guard_disarm_transfers_ownership() {
+        let scope = RuntimeTaskScope::new();
+        RuntimeTaskScopePreparationGuard::new(scope.clone()).disarm();
+
+        assert!(!scope.cancellation_token().is_cancelled());
+        scope.stop().await;
+    }
+
+    #[tokio::test]
+    async fn aborted_scope_stop_still_cancels_children() {
+        let scope = RuntimeTaskScope::new();
+        let registration = scope.admission.try_register().unwrap();
+        let started = Arc::new(Notify::new());
+        let dropped = Arc::new(Notify::new());
+        let started_task = Arc::clone(&started);
+        let dropped_task = Arc::clone(&dropped);
+        scope.spawn(async move {
+            let _drop_signal = NotifyOnDrop(dropped_task);
+            started_task.notify_one();
+            std::future::pending::<()>().await;
+        });
+        started.notified().await;
+
+        let stop_scope = scope.clone();
+        let stop = tokio::spawn(async move {
+            stop_scope.stop().await;
+        });
+        scope.cancellation_token().cancelled().await;
+        stop.abort();
+        assert!(stop.await.unwrap_err().is_cancelled());
+        drop(registration);
+
+        tokio::time::timeout(Duration::from_secs(1), dropped.notified())
+            .await
+            .expect("scope cancellation must drop the tracked child");
+        assert!(scope.admission.try_register().is_none());
+    }
+
+    #[tokio::test]
+    async fn aborted_graceful_drain_forces_session_cancellation() {
+        let generation = test_runtime_generation(1, ProxyConfig::default());
+        let registration = generation.session_admission.try_register().unwrap();
+        let started = Arc::new(Notify::new());
+        let dropped = Arc::new(Notify::new());
+        let started_task = Arc::clone(&started);
+        let dropped_task = Arc::clone(&dropped);
+        assert!(generation.spawn_session(async move {
+            let _drop_signal = NotifyOnDrop(dropped_task);
+            started_task.notify_one();
+            std::future::pending::<()>().await;
+        }));
+        started.notified().await;
+
+        let drain_generation = Arc::clone(&generation);
+        let drain = tokio::spawn(async move {
+            drain_generation.drain_sessions(Duration::from_secs(60)).await
+        });
+        while generation.session_admission.state.load(Ordering::Acquire)
+            & SESSION_ADMISSION_CLOSED
+            == 0
+        {
+            tokio::task::yield_now().await;
+        }
+        drain.abort();
+        assert!(drain.await.unwrap_err().is_cancelled());
+        assert!(generation.session_cancel.is_cancelled());
+        drop(registration);
+
+        tokio::time::timeout(Duration::from_secs(1), dropped.notified())
+            .await
+            .expect("aborted graceful drain must cancel existing sessions");
+    }
+}

src/maestro/orchestrator.rs

@@ -3,7 +3,7 @@ use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 
 use arc_swap::ArcSwap;
-use tokio::sync::{RwLock, watch};
+use tokio::sync::{RwLock, Semaphore, watch};
 use tracing::{error, info};
 
 use crate::api;
@@ -12,7 +12,9 @@ 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::traffic_limiter::TrafficLimiter;
 use crate::proxy::user_admission::UserAdmissionAuthority;
+use crate::proxy::user_connection_authority::UserConnectionAuthority;
 use crate::startup::{COMPONENT_API_BOOTSTRAP, COMPONENT_NETWORK_PROBE};
 use crate::stats::telemetry::TelemetryPolicy;
 use crate::stats::{QuotaStore, Stats};
@@ -47,10 +49,16 @@ pub(super) async fn run_telemt_core(
     } = bootstrap::bootstrap(privilege_drop_requested).await?;
 
     let quota_store = Arc::new(QuotaStore::default());
-    let stats = Arc::new(Stats::with_quota_store(quota_store.clone()));
+    let connection_authority = Arc::new(UserConnectionAuthority::default());
+    let stats = Arc::new(Stats::with_process_authorities(
+        quota_store.clone(),
+        connection_authority,
+    ));
     let tls_full_cert_budget = Arc::new(TlsFullCertBudget::new());
     let process_control_plane = control_plane::ProcessControlPlane::new();
     let runtime_task_scope = generation::RuntimeTaskScope::new();
+    let runtime_task_scope_guard =
+        generation::RuntimeTaskScopePreparationGuard::new(runtime_task_scope.clone());
     stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
     let quota_state_path = config.general.quota_state_path.clone();
     let quota_state =
@@ -72,14 +80,11 @@ pub(super) async fn run_telemt_core(
         .with_dns_overrides(&config.network.dns_overrides)?,
     );
     let ip_tracker = Arc::new(UserIpTracker::new());
-    ip_tracker
-        .load_limits(
+    let _ = ip_tracker
+        .apply_policy_from_source(
+            1,
             config.access.user_max_unique_ips_global_each,
             &config.access.user_max_unique_ips,
-        )
-        .await;
-    ip_tracker
-        .set_limit_policy(
             config.access.user_max_unique_ips_mode,
             config.access.user_max_unique_ips_window_secs,
         )
@@ -102,14 +107,22 @@ pub(super) async fn run_telemt_core(
     let direct_buffer_hard_limit =
         resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
     let direct_buffer_budget = DirectBufferBudget::new(direct_buffer_hard_limit);
+    direct_buffer_budget.activate_controller(1);
     info!(
         hard_limit_bytes = direct_buffer_hard_limit,
         configured_override_bytes = config.general.direct_relay_buffer_budget_max_bytes,
         "Direct relay buffer budget initialized"
     );
     let user_admission = UserAdmissionAuthority::new_with_quota_store(quota_store.clone());
-    let shared_state = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
+    let traffic_limiter = TrafficLimiter::new();
+    let _ = traffic_limiter.apply_policy_from_source(
+        1,
+        config.access.user_rate_limits.clone(),
+        config.access.cidr_rate_limits.clone(),
+    );
+    let shared_state = ProxySharedState::new_with_process_authorities(
         direct_buffer_budget.clone(),
+        traffic_limiter,
         user_admission,
     );
     let _ = shared_state.activate_user_config_source(
@@ -118,10 +131,12 @@ pub(super) async fn run_telemt_core(
         &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(),
-    );
+    let max_connections_limit = if config.server.max_connections == 0 {
+        Semaphore::MAX_PERMITS
+    } else {
+        config.server.max_connections as usize
+    };
+    let max_connections = Arc::new(Semaphore::new(max_connections_limit));
     let web_trace = WebTraceStore::new(config.web.debug.clone(), &config.web.limits);
     let web_runtime_control = WebRuntimeControl::new();
 
@@ -303,6 +318,7 @@ pub(super) async fn run_telemt_core(
         ip_tracker.clone(),
         shared_state.clone(),
         direct_buffer_budget,
+        max_connections,
         route_runtime.clone(),
         api_me_pool.clone(),
         runtime_task_scope.clone(),
@@ -333,6 +349,7 @@ pub(super) async fn run_telemt_core(
         runtime.max_connections,
         runtime_task_scope,
     );
+    runtime_task_scope_guard.disarm();
     let active_runtime = Arc::new(ArcSwap::from(runtime_generation));
     let bound = listeners::bind_listeners(
         &runtime.config,

src/maestro/reload_supervisor.rs

@@ -175,9 +175,14 @@ impl ReloadSupervisor {
             resolved.effective,
             &self.config_path,
             self.quota_store.clone(),
+            old_runtime.stats.connection_authority(),
             self.runtime_log_filter.clone(),
             self.tls_full_cert_budget.clone(),
             old_runtime.proxy_shared.user_admission(),
+            old_runtime.ip_tracker.clone(),
+            old_runtime.proxy_shared.traffic_limiter.clone(),
+            old_runtime.proxy_shared.direct_buffer_budget.clone(),
+            old_runtime.max_connections.clone(),
         )
         .await
         {
@@ -300,15 +305,37 @@ impl ReloadSupervisor {
         } else {
             None
         };
+        let config = new_runtime.config();
+        let _ = new_runtime.proxy_shared.activate_user_config_source(
+            new_runtime.id,
+            Some(user_admission_epoch),
+            &config.access.users,
+            &config.access.user_enabled,
+        );
+        let _ = new_runtime
+            .ip_tracker
+            .apply_policy_from_source(
+                new_runtime.id,
+                config.access.user_max_unique_ips_global_each,
+                &config.access.user_max_unique_ips,
+                config.access.user_max_unique_ips_mode,
+                config.access.user_max_unique_ips_window_secs,
+            )
+            .await;
+        let _ = new_runtime
+            .proxy_shared
+            .traffic_limiter
+            .apply_policy_from_source(
+                new_runtime.id,
+                config.access.user_rate_limits.clone(),
+                config.access.cidr_rate_limits.clone(),
+            );
+        new_runtime
+            .proxy_shared
+            .direct_buffer_budget
+            .activate_controller(new_runtime.id);
         let replaced = {
             let listener_manager = self.listener_manager.lock().await;
-            let config = new_runtime.config();
-            let _ = new_runtime.proxy_shared.activate_user_config_source(
-                new_runtime.id,
-                Some(user_admission_epoch),
-                &config.access.users,
-                &config.access.user_enabled,
-            );
             old_runtime.stop_accepting_sessions();
             listener_manager.activate_runtime_generation(new_runtime.clone())
         };

src/maestro/runtime_build.rs

@@ -12,11 +12,13 @@ use crate::crypto::SecureRandom;
 use crate::ip_tracker::UserIpTracker;
 use crate::network::probe::{decide_network_capabilities, run_probe};
 use crate::proxy::direct_buffer_budget::{
-    DirectBufferBudget, resolve_direct_buffer_hard_limit, run_direct_buffer_budget_controller,
+    DirectBufferBudget, run_direct_buffer_budget_controller,
 };
 use crate::proxy::route_mode::{RelayRouteMode, RouteRuntimeController};
 use crate::proxy::shared_state::ProxySharedState;
+use crate::proxy::traffic_limiter::TrafficLimiter;
 use crate::proxy::user_admission::UserAdmissionAuthority;
+use crate::proxy::user_connection_authority::UserConnectionAuthority;
 use crate::startup::StartupTracker;
 use crate::stats::beobachten::BeobachtenStore;
 use crate::stats::telemetry::TelemetryPolicy;
@@ -27,7 +29,9 @@ use crate::transport::UpstreamManager;
 use crate::transport::middle_proxy::MePool;
 
 use super::admission;
-use super::generation::{RuntimeGeneration, RuntimeTaskScope};
+use super::generation::{
+    RuntimeGeneration, RuntimeTaskScope, RuntimeTaskScopePreparationGuard,
+};
 use super::listeners::listener_rebind_supported;
 use super::runtime_tasks::RuntimeLogFilter;
 use super::{me_startup, runtime_tasks, tls_bootstrap};
@@ -49,9 +53,14 @@ pub(crate) async fn prepare_runtime(
     config: ProxyConfig,
     config_path: &Path,
     quota_store: Arc<QuotaStore>,
+    connection_authority: Arc<UserConnectionAuthority>,
     runtime_log_filter: RuntimeLogFilter,
     tls_full_cert_budget: Arc<TlsFullCertBudget>,
     user_admission: Arc<UserAdmissionAuthority>,
+    ip_tracker: Arc<UserIpTracker>,
+    traffic_limiter: Arc<TrafficLimiter>,
+    direct_buffer_budget: Arc<DirectBufferBudget>,
+    max_connections: Arc<Semaphore>,
 ) -> Result<PreparedRuntime, String> {
     let user_admission_epoch = user_admission.epoch();
     config
@@ -63,7 +72,11 @@ pub(crate) async fn prepare_runtime(
         .as_secs();
     let startup_tracker = Arc::new(StartupTracker::new(started_at_epoch_secs));
     let task_scope = RuntimeTaskScope::new();
-    let stats = Arc::new(Stats::with_quota_store(quota_store));
+    let task_scope_guard = RuntimeTaskScopePreparationGuard::new(task_scope.clone());
+    let stats = Arc::new(Stats::with_process_authorities(
+        quota_store,
+        connection_authority,
+    ));
     stats.apply_telemetry_policy(TelemetryPolicy::from_config(&config.general.telemetry));
 
     let upstream_manager = Arc::new(
@@ -80,31 +93,11 @@ pub(crate) async fn prepare_runtime(
         .with_dns_overrides(&config.network.dns_overrides)
         .map_err(|error| format!("DNS override preparation failed: {}", error))?,
     );
-    let ip_tracker = Arc::new(UserIpTracker::new());
-    ip_tracker
-        .load_limits(
-            config.access.user_max_unique_ips_global_each,
-            &config.access.user_max_unique_ips,
-        )
-        .await;
-    ip_tracker
-        .set_limit_policy(
-            config.access.user_max_unique_ips_mode,
-            config.access.user_max_unique_ips_window_secs,
-        )
-        .await;
-
-    let hard_limit =
-        resolve_direct_buffer_hard_limit(config.general.direct_relay_buffer_budget_max_bytes).await;
-    let direct_buffer_budget = DirectBufferBudget::new(hard_limit);
-    let proxy_shared = ProxySharedState::new_with_direct_buffer_budget_and_user_admission(
+    let proxy_shared = ProxySharedState::new_with_process_authorities(
         direct_buffer_budget.clone(),
+        traffic_limiter,
         user_admission,
     );
-    proxy_shared.traffic_limiter.apply_policy(
-        config.access.user_rate_limits.clone(),
-        config.access.cidr_rate_limits.clone(),
-    );
 
     let probe = run_probe(
         &config.network,
@@ -184,12 +177,6 @@ pub(crate) async fn prepare_runtime(
         Duration::from_secs(config.access.replay_window_secs),
     ));
     let buffer_pool = Arc::new(BufferPool::with_config(64 * 1024, 4096));
-    let max_connections_limit = if config.server.max_connections == 0 {
-        Semaphore::MAX_PERMITS
-    } else {
-        config.server.max_connections as usize
-    };
-    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,
@@ -288,10 +275,12 @@ pub(crate) async fn prepare_runtime(
         conntrack_scope.cancellation_token(),
     ));
     task_scope.spawn(run_direct_buffer_budget_controller(
+        generation_id,
         direct_buffer_budget,
         buffer_pool.clone(),
         stats.clone(),
         proxy_shared.clone(),
+        max_connections.clone(),
         config.server.max_connections,
     ));
     let generation = RuntimeGeneration::new(
@@ -313,6 +302,7 @@ pub(crate) async fn prepare_runtime(
         max_connections,
         task_scope,
     );
+    task_scope_guard.disarm();
     drop(admission_tx);
 
     Ok(PreparedRuntime {
@@ -414,6 +404,17 @@ pub(crate) fn resolve_reload_config(
         effective.server.metrics_listen = old.server.metrics_listen.clone();
         effective.server.metrics_port = old.server.metrics_port;
     }
+    if old.server.max_connections != desired.server.max_connections {
+        fields.push("server.max_connections".to_string());
+        effective.server.max_connections = old.server.max_connections;
+    }
+    if old.general.direct_relay_buffer_budget_max_bytes
+        != desired.general.direct_relay_buffer_budget_max_bytes
+    {
+        fields.push("general.direct_relay_buffer_budget_max_bytes".to_string());
+        effective.general.direct_relay_buffer_budget_max_bytes =
+            old.general.direct_relay_buffer_budget_max_bytes;
+    }
     if old.general.quota_state_path != desired.general.quota_state_path {
         fields.push("general.quota_state_path".to_string());
         effective.general.quota_state_path = old.general.quota_state_path.clone();

src/maestro/runtime_build_tests.rs

@@ -94,6 +94,39 @@ fn global_mss_profiles_are_deferred_with_the_listener_socket_group() {
     assert!(!resolved.runtime_changed);
 }
 
+#[test]
+fn process_wide_connection_and_direct_buffer_envelopes_are_restart_only() {
+    let old = ProxyConfig::default();
+    let mut desired = old.clone();
+    desired.server.max_connections = old.server.max_connections.saturating_add(1);
+    desired.general.direct_relay_buffer_budget_max_bytes = old
+        .general
+        .direct_relay_buffer_budget_max_bytes
+        .saturating_add(4 * 1024);
+
+    let resolved = resolve_reload_config(&old, &desired).unwrap();
+
+    assert_eq!(
+        resolved.deferred_process_fields,
+        vec![
+            "server.max_connections".to_string(),
+            "general.direct_relay_buffer_budget_max_bytes".to_string(),
+        ]
+    );
+    assert_eq!(
+        resolved.effective.server.max_connections,
+        old.server.max_connections
+    );
+    assert_eq!(
+        resolved
+            .effective
+            .general
+            .direct_relay_buffer_budget_max_bytes,
+        old.general.direct_relay_buffer_budget_max_bytes
+    );
+    assert!(!resolved.runtime_changed);
+}
+
 #[test]
 fn mixed_reload_retains_process_state_and_applies_runtime_state() {
     let old = ProxyConfig::default();

src/maestro/runtime_startup.rs

@@ -56,6 +56,7 @@ pub(super) async fn prepare_runtime(
     ip_tracker: Arc<UserIpTracker>,
     shared_state: Arc<ProxySharedState>,
     direct_buffer_budget: Arc<DirectBufferBudget>,
+    max_connections: Arc<Semaphore>,
     route_runtime: Arc<RouteRuntimeController>,
     api_me_pool: Arc<RwLock<Option<Arc<MePool>>>>,
     runtime_task_scope: RuntimeTaskScope,
@@ -69,13 +70,6 @@ pub(super) async fn prepare_runtime(
     let beobachten = Arc::new(BeobachtenStore::new());
     let rng = Arc::new(SecureRandom::new());
 
-    let max_connections_limit = if config.server.max_connections == 0 {
-        Semaphore::MAX_PERMITS
-    } else {
-        config.server.max_connections as usize
-    };
-    let max_connections = Arc::new(Semaphore::new(max_connections_limit));
-
     let me2dc_fallback = config.general.me2dc_fallback;
     let me_init_retry_attempts = config.general.me_init_retry_attempts;
     if use_middle_proxy && !decision.ipv4_me && !decision.ipv6_me {
@@ -346,10 +340,12 @@ pub(super) async fn prepare_runtime(
         conntrack_scope.cancellation_token(),
     ));
     runtime_task_scope.spawn(run_direct_buffer_budget_controller(
+        1,
         direct_buffer_budget,
         buffer_pool.clone(),
         stats,
         shared_state,
+        max_connections.clone(),
         config.server.max_connections,
     ));
 

src/maestro/runtime_tasks.rs

@@ -138,7 +138,9 @@ pub(crate) async fn spawn_runtime_tasks(
 
     let ip_tracker_maintenance = ip_tracker.clone();
     task_scope.spawn(async move {
-        ip_tracker_maintenance.run_periodic_maintenance().await;
+        ip_tracker_maintenance
+            .run_periodic_maintenance(generation_id)
+            .await;
     });
 
     let detected_ip_v4: Option<IpAddr> = probe.detected_ipv4.map(IpAddr::V4);
@@ -197,20 +199,7 @@ pub(crate) async fn spawn_runtime_tasks(
     let ip_tracker_policy = ip_tracker.clone();
     let mut config_rx_ip_limits = config_rx.clone();
     task_scope.spawn(async move {
-        let mut prev_limits = config_rx_ip_limits
-            .borrow()
-            .access
-            .user_max_unique_ips
-            .clone();
-        let mut prev_global_each = config_rx_ip_limits
-            .borrow()
-            .access
-            .user_max_unique_ips_global_each;
-        let mut prev_mode = config_rx_ip_limits.borrow().access.user_max_unique_ips_mode;
-        let mut prev_window = config_rx_ip_limits
-            .borrow()
-            .access
-            .user_max_unique_ips_window_secs;
+        let mut previous = config_rx_ip_limits.borrow().access.clone();
 
         loop {
             if config_rx_ip_limits.changed().await.is_err() {
@@ -218,39 +207,28 @@ pub(crate) async fn spawn_runtime_tasks(
             }
             let cfg = config_rx_ip_limits.borrow_and_update().clone();
 
-            if prev_limits != cfg.access.user_max_unique_ips
-                || prev_global_each != cfg.access.user_max_unique_ips_global_each
+            if previous.user_max_unique_ips != cfg.access.user_max_unique_ips
+                || previous.user_max_unique_ips_global_each
+                    != cfg.access.user_max_unique_ips_global_each
+                || previous.user_max_unique_ips_mode != cfg.access.user_max_unique_ips_mode
+                || previous.user_max_unique_ips_window_secs
+                    != cfg.access.user_max_unique_ips_window_secs
             {
-                ip_tracker_policy
-                    .load_limits(
+                let _ = ip_tracker_policy
+                    .apply_policy_from_source(
+                        generation_id,
                         cfg.access.user_max_unique_ips_global_each,
                         &cfg.access.user_max_unique_ips,
-                    )
-                    .await;
-                prev_limits = cfg.access.user_max_unique_ips.clone();
-                prev_global_each = cfg.access.user_max_unique_ips_global_each;
-            }
-
-            if prev_mode != cfg.access.user_max_unique_ips_mode
-                || prev_window != cfg.access.user_max_unique_ips_window_secs
-            {
-                ip_tracker_policy
-                    .set_limit_policy(
                         cfg.access.user_max_unique_ips_mode,
                         cfg.access.user_max_unique_ips_window_secs,
                     )
                     .await;
-                prev_mode = cfg.access.user_max_unique_ips_mode;
-                prev_window = cfg.access.user_max_unique_ips_window_secs;
+                previous = cfg.access.clone();
             }
         }
     });
 
     let limiter = shared_state.traffic_limiter.clone();
-    limiter.apply_policy(
-        config.access.user_rate_limits.clone(),
-        config.access.cidr_rate_limits.clone(),
-    );
     let mut config_rx_rate_limits = config_rx.clone();
     task_scope.spawn(async move {
         let mut prev_user_limits = config_rx_rate_limits
@@ -271,7 +249,8 @@ pub(crate) async fn spawn_runtime_tasks(
             if prev_user_limits != cfg.access.user_rate_limits
                 || prev_cidr_limits != cfg.access.cidr_rate_limits
             {
-                limiter.apply_policy(
+                let _ = limiter.apply_policy_from_source(
+                    generation_id,
                     cfg.access.user_rate_limits.clone(),
                     cfg.access.cidr_rate_limits.clone(),
                 );