Harden Maestro reload lifecycle and readiness barriers

Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com>

src/maestro/reload_supervisor.rs

@@ -8,7 +8,7 @@ use tracing::{info, warn};
 
 use crate::stats::QuotaStore;
 
-use super::generation::RuntimeGeneration;
+use super::generation::{RuntimeGeneration, RuntimeWatchState};
 use super::reload::{
     ReloadCommand, ReloadCommandReceiver, ReloadControl, ReloadFailurePolicy, ReloadMode,
     ReloadPhase,
@@ -24,6 +24,48 @@ pub(crate) struct ReloadSupervisor {
     quota_store: Arc<QuotaStore>,
     detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
     runtime_log_filter: RuntimeLogFilter,
+    runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
+}
+
+#[derive(Debug, PartialEq, Eq)]
+enum RevisionGateAction {
+    Proceed,
+    Warn(String),
+    Rollback(String),
+}
+
+fn revision_gate_action(
+    accepted_revision: &str,
+    current_revision: Result<String, String>,
+    failure_policy: ReloadFailurePolicy,
+) -> RevisionGateAction {
+    let warning = match current_revision {
+        Ok(current) if current == accepted_revision => return RevisionGateAction::Proceed,
+        Ok(current) => format!(
+            "config revision changed during preparation: accepted={} current={}",
+            accepted_revision, current
+        ),
+        Err(error) => format!("config revision verification failed: {}", error),
+    };
+    match failure_policy {
+        ReloadFailurePolicy::KeepNew => RevisionGateAction::Warn(warning),
+        ReloadFailurePolicy::Rollback => RevisionGateAction::Rollback(warning),
+    }
+}
+
+async fn stop_background_and_middle_end(generation: &RuntimeGeneration) -> bool {
+    generation.stop_background_tasks().await;
+    let Some(pool) = generation.current_me_pool().await else {
+        return false;
+    };
+    tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
+        .await
+        .is_err()
+}
+
+async fn cleanup_candidate(generation: &RuntimeGeneration) -> bool {
+    generation.stop_sessions().await;
+    stop_background_and_middle_end(generation).await
 }
 
 impl ReloadSupervisor {
@@ -36,6 +78,7 @@ impl ReloadSupervisor {
         quota_store: Arc<QuotaStore>,
         detected_ips_tx: watch::Sender<(Option<std::net::IpAddr>, Option<std::net::IpAddr>)>,
         runtime_log_filter: RuntimeLogFilter,
+        runtime_watch_tx: watch::Sender<Option<RuntimeWatchState>>,
     ) {
         let supervisor = Self {
             active_runtime,
@@ -45,6 +88,7 @@ impl ReloadSupervisor {
             quota_store,
             detected_ips_tx,
             runtime_log_filter,
+            runtime_watch_tx,
         };
         tokio::spawn(supervisor.run());
     }
@@ -81,30 +125,23 @@ impl ReloadSupervisor {
             }
         };
 
-        if let Ok(current_revision) =
-            crate::api::config_store::current_revision_for_maestro(&self.config_path).await
-            && current_revision != command.config_revision
-        {
-            let warning = format!(
-                "config revision changed during preparation: accepted={} current={}",
-                command.config_revision, current_revision
-            );
-            if command.request.failure_policy == ReloadFailurePolicy::Rollback {
+        let revision_action = revision_gate_action(
+            &command.config_revision,
+            crate::api::config_store::current_revision_for_maestro(&self.config_path).await,
+            command.request.failure_policy,
+        );
+        match revision_action {
+            RevisionGateAction::Proceed => {}
+            RevisionGateAction::Warn(warning) => {
+                self.control.add_warning(command.reload_id, warning).await;
+            }
+            RevisionGateAction::Rollback(warning) => {
+                let _ = cleanup_candidate(&prepared.generation).await;
                 self.runtime_log_filter
                     .apply_reload(&old_runtime.config().general.log_level);
-                prepared.generation.stop_sessions().await;
-                if let Some(pool) = prepared.generation.current_me_pool().await {
-                    let _ = tokio::time::timeout(
-                        Duration::from_secs(2),
-                        pool.shutdown_send_close_conn_all(),
-                    )
-                    .await;
-                }
-                prepared.generation.stop_background_tasks().await;
                 self.control.rolled_back(command.reload_id, warning).await;
                 return;
             }
-            self.control.add_warning(command.reload_id, warning).await;
         }
 
         self.control
@@ -112,37 +149,26 @@ impl ReloadSupervisor {
             .await;
         let new_runtime = prepared.generation;
         old_runtime.stop_accepting_sessions();
-        let replaced = self.active_runtime.swap(new_runtime.clone());
         if let Err(error) = crate::network::dns_overrides::install_entries(
             &new_runtime.config().network.dns_overrides,
         ) {
             let message = format!("runtime DNS activation failed: {}", error);
             if command.request.failure_policy == ReloadFailurePolicy::Rollback {
-                let candidate = self.active_runtime.swap(replaced.clone());
-                replaced.resume_accepting_sessions();
+                old_runtime.resume_accepting_sessions();
+                let _ = cleanup_candidate(&new_runtime).await;
                 self.runtime_log_filter
-                    .apply_reload(&replaced.config().general.log_level);
-                let _ = crate::network::dns_overrides::install_entries(
-                    &replaced.config().network.dns_overrides,
-                );
-                candidate.stop_sessions().await;
-                if let Some(pool) = candidate.current_me_pool().await {
-                    let _ = tokio::time::timeout(
-                        Duration::from_secs(2),
-                        pool.shutdown_send_close_conn_all(),
-                    )
-                    .await;
-                }
-                candidate.stop_background_tasks().await;
+                    .apply_reload(&old_runtime.config().general.log_level);
                 self.control.rolled_back(command.reload_id, message).await;
                 return;
             }
             self.control.add_warning(command.reload_id, message).await;
         }
+        let replaced = self.active_runtime.swap(new_runtime.clone());
         self.detected_ips_tx.send_replace(prepared.detected_ips);
         self.runtime_log_filter
             .apply_reload(&new_runtime.config().general.log_level);
-        crate::synlimit_control::reconcile_synlimit_rules(&new_runtime.config()).await;
+        self.runtime_watch_tx
+            .send_replace(Some(new_runtime.watch_state()));
 
         info!(
             reload_id = command.reload_id,
@@ -177,11 +203,7 @@ impl ReloadSupervisor {
             }
         }
 
-        if let Some(pool) = replaced.current_me_pool().await
-            && tokio::time::timeout(Duration::from_secs(2), pool.shutdown_send_close_conn_all())
-                .await
-                .is_err()
-        {
+        if stop_background_and_middle_end(&replaced).await {
             let warning = format!(
                 "generation {} Middle-End close broadcast timed out",
                 replaced.id
@@ -189,9 +211,46 @@ impl ReloadSupervisor {
             warn!(reload_id = command.reload_id, warning = %warning);
             self.control.add_warning(command.reload_id, warning).await;
         }
-        replaced.stop_background_tasks().await;
         self.control
             .succeed(command.reload_id, new_runtime.id)
             .await;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn revision_gate_proceeds_only_on_verified_match() {
+        assert_eq!(
+            revision_gate_action(
+                "accepted",
+                Ok("accepted".to_string()),
+                ReloadFailurePolicy::Rollback,
+            ),
+            RevisionGateAction::Proceed
+        );
+    }
+
+    #[test]
+    fn revision_gate_applies_failure_policy_to_mismatch_and_read_error() {
+        for result in [
+            Ok("changed".to_string()),
+            Err("read failed".to_string()),
+        ] {
+            assert!(matches!(
+                revision_gate_action(
+                    "accepted",
+                    result.clone(),
+                    ReloadFailurePolicy::KeepNew,
+                ),
+                RevisionGateAction::Warn(_)
+            ));
+            assert!(matches!(
+                revision_gate_action("accepted", result, ReloadFailurePolicy::Rollback),
+                RevisionGateAction::Rollback(_)
+            ));
+        }
+    }
+}