Atomic Maestro Sessions + Shutdown gate

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

src/maestro/reload_supervisor_tests.rs

@@ -0,0 +1,320 @@
+use super::*;
+use crate::config::ProxyConfig;
+use crate::maestro::generation::test_runtime_generation;
+use crate::maestro::reload::{ReloadRequest, ReloadSubmitError};
+use crate::stats::QuotaStore;
+use tokio::sync::Notify;
+use tracing_subscriber::{EnvFilter, Registry};
+
+struct ReloadFixture {
+    supervisor: Arc<ReloadSupervisor>,
+    control: ReloadControl,
+    command: ReloadCommand,
+    old_runtime: Arc<RuntimeGeneration>,
+    new_runtime: Arc<RuntimeGeneration>,
+    runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
+}
+
+fn runtime_log_filter() -> RuntimeLogFilter {
+    let (_layer, handle) =
+        tracing_subscriber::reload::Layer::<EnvFilter, Registry>::new(EnvFilter::new("info"));
+    RuntimeLogFilter::new(handle)
+}
+
+async fn fixture(request: ReloadRequest) -> ReloadFixture {
+    let old_runtime = test_runtime_generation(1, ProxyConfig::default());
+    let new_config = Arc::new(ProxyConfig::default());
+    let new_runtime = test_runtime_generation(2, new_config.as_ref().clone());
+    let active_runtime = Arc::new(ArcSwap::from(old_runtime.clone()));
+    let (control, commands) = ReloadControl::channel(old_runtime.id);
+    let accepted = control
+        .submit(new_config.clone(), "revision".to_string(), request.clone())
+        .await
+        .unwrap();
+    let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None));
+    let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(old_runtime.watch_state()));
+    let supervisor = Arc::new(ReloadSupervisor {
+        active_runtime,
+        control: control.clone(),
+        commands,
+        config_path: PathBuf::new(),
+        quota_store: Arc::new(QuotaStore::default()),
+        detected_ips_tx,
+        runtime_log_filter: runtime_log_filter(),
+        runtime_watch_tx,
+    });
+    let command = ReloadCommand {
+        reload_id: accepted.reload_id,
+        target_generation: accepted.target_generation,
+        config: new_config,
+        config_revision: accepted.config_revision,
+        request,
+    };
+    ReloadFixture {
+        supervisor,
+        control,
+        command,
+        old_runtime,
+        new_runtime,
+        runtime_watch_rx,
+    }
+}
+
+struct DropSignal(Arc<Notify>);
+
+impl Drop for DropSignal {
+    fn drop(&mut self) {
+        self.0.notify_one();
+    }
+}
+
+#[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(_)
+        ));
+    }
+}
+#[tokio::test]
+async fn revision_rollback_keeps_old_generation_and_cleans_candidate() {
+    let fixture = fixture(ReloadRequest {
+        failure_policy: ReloadFailurePolicy::Rollback,
+        ..ReloadRequest::default()
+    })
+    .await;
+    let candidate_dropped = Arc::new(Notify::new());
+    let candidate_drop = candidate_dropped.clone();
+    assert!(fixture.new_runtime.spawn_session(async move {
+        let _drop_signal = DropSignal(candidate_drop);
+        std::future::pending::<()>().await;
+    }));
+    tokio::task::yield_now().await;
+
+    fixture
+        .supervisor
+        .activate_prepared(
+            fixture.command,
+            fixture.old_runtime.clone(),
+            PreparedRuntime {
+                generation: fixture.new_runtime,
+                detected_ips: (None, None),
+            },
+            RevisionGateAction::Rollback("revision changed".to_string()),
+            |_| -> Result<(), String> { panic!("DNS activation must not run on rollback") },
+        )
+        .await;
+
+    tokio::time::timeout(Duration::from_secs(1), candidate_dropped.notified())
+        .await
+        .unwrap();
+    assert_eq!(fixture.supervisor.active_runtime.load().id, 1);
+    assert_eq!(
+        fixture
+            .runtime_watch_rx
+            .borrow()
+            .as_ref()
+            .unwrap()
+            .generation_id,
+        1
+    );
+    assert!(fixture.old_runtime.spawn_session(async {}));
+    let status = fixture.control.status(1).await.unwrap();
+    assert_eq!(status.state, ReloadPhase::RolledBack);
+    fixture.old_runtime.stop_sessions().await;
+}
+
+#[tokio::test]
+async fn dns_failure_policy_controls_rollback_or_keep_new() {
+    for policy in [ReloadFailurePolicy::Rollback, ReloadFailurePolicy::KeepNew] {
+        let fixture = fixture(ReloadRequest {
+            failure_policy: policy,
+            ..ReloadRequest::default()
+        })
+        .await;
+        fixture
+            .supervisor
+            .activate_prepared(
+                fixture.command,
+                fixture.old_runtime.clone(),
+                PreparedRuntime {
+                    generation: fixture.new_runtime.clone(),
+                    detected_ips: (None, None),
+                },
+                RevisionGateAction::Proceed,
+                |_| Err("invalid DNS entry".to_string()),
+            )
+            .await;
+
+        let status = fixture.control.status(1).await.unwrap();
+        match policy {
+            ReloadFailurePolicy::Rollback => {
+                assert_eq!(fixture.supervisor.active_runtime.load().id, 1);
+                assert_eq!(status.state, ReloadPhase::RolledBack);
+                assert!(fixture.old_runtime.spawn_session(async {}));
+                fixture.old_runtime.stop_sessions().await;
+            }
+            ReloadFailurePolicy::KeepNew => {
+                assert_eq!(fixture.supervisor.active_runtime.load().id, 2);
+                assert_eq!(status.state, ReloadPhase::Succeeded);
+                assert_eq!(status.warnings.len(), 1);
+                assert!(!fixture.old_runtime.spawn_session(async {}));
+                fixture.new_runtime.stop_sessions().await;
+            }
+        }
+    }
+}
+
+#[tokio::test]
+async fn drain_publishes_new_generation_before_old_sessions_finish() {
+    let mut fixture = fixture(ReloadRequest {
+        mode: ReloadMode::Drain,
+        timeout_secs: Some(30),
+        ..ReloadRequest::default()
+    })
+    .await;
+    let old_started = Arc::new(Notify::new());
+    let old_release = Arc::new(Notify::new());
+    let started = old_started.clone();
+    let release = old_release.clone();
+    assert!(fixture.old_runtime.spawn_session(async move {
+        started.notify_one();
+        release.notified().await;
+    }));
+    old_started.notified().await;
+
+    let supervisor = fixture.supervisor.clone();
+    let old_runtime = fixture.old_runtime.clone();
+    let new_runtime = fixture.new_runtime.clone();
+    let activation = tokio::spawn(async move {
+        supervisor
+            .activate_prepared(
+                fixture.command,
+                old_runtime,
+                PreparedRuntime {
+                    generation: new_runtime,
+                    detected_ips: (None, None),
+                },
+                RevisionGateAction::Proceed,
+                |_| Ok(()),
+            )
+            .await;
+    });
+
+    fixture.runtime_watch_rx.changed().await.unwrap();
+    assert_eq!(
+        fixture
+            .runtime_watch_rx
+            .borrow()
+            .as_ref()
+            .unwrap()
+            .generation_id,
+        2
+    );
+    assert!(!activation.is_finished());
+    assert!(!fixture.old_runtime.spawn_session(async {}));
+
+    old_release.notify_one();
+    activation.await.unwrap();
+    assert_eq!(
+        fixture.control.status(1).await.unwrap().state,
+        ReloadPhase::Succeeded
+    );
+    fixture.new_runtime.stop_sessions().await;
+}
+
+#[tokio::test(start_paused = true)]
+async fn drain_timeout_cancels_old_sessions_and_records_one_warning() {
+    let mut fixture = fixture(ReloadRequest {
+        mode: ReloadMode::Drain,
+        timeout_secs: Some(1),
+        ..ReloadRequest::default()
+    })
+    .await;
+    let dropped = Arc::new(Notify::new());
+    let drop_signal = dropped.clone();
+    assert!(fixture.old_runtime.spawn_session(async move {
+        let _drop_signal = DropSignal(drop_signal);
+        std::future::pending::<()>().await;
+    }));
+    tokio::task::yield_now().await;
+
+    let supervisor = fixture.supervisor.clone();
+    let old_runtime = fixture.old_runtime.clone();
+    let new_runtime = fixture.new_runtime.clone();
+    let activation = tokio::spawn(async move {
+        supervisor
+            .activate_prepared(
+                fixture.command,
+                old_runtime,
+                PreparedRuntime {
+                    generation: new_runtime,
+                    detected_ips: (None, None),
+                },
+                RevisionGateAction::Proceed,
+                |_| Ok(()),
+            )
+            .await;
+    });
+    fixture.runtime_watch_rx.changed().await.unwrap();
+    tokio::task::yield_now().await;
+    tokio::time::advance(Duration::from_secs(1)).await;
+    activation.await.unwrap();
+
+    dropped.notified().await;
+    let status = fixture.control.status(1).await.unwrap();
+    assert_eq!(status.state, ReloadPhase::Succeeded);
+    assert_eq!(status.warnings.len(), 1);
+    assert!(status.warnings[0].contains("exceeded drain timeout"));
+    fixture.new_runtime.stop_sessions().await;
+}
+
+#[tokio::test]
+async fn quiesce_joins_idle_supervisor_and_rejects_later_submissions() {
+    let runtime = test_runtime_generation(1, ProxyConfig::default());
+    let active_runtime = Arc::new(ArcSwap::from(runtime.clone()));
+    let (control, commands) = ReloadControl::channel(runtime.id);
+    let (detected_ips_tx, _detected_ips_rx) = watch::channel((None, None));
+    let (runtime_watch_tx, _runtime_watch_rx) = watch::channel(Some(runtime.watch_state()));
+    let handle = ReloadSupervisor::spawn(
+        active_runtime,
+        control.clone(),
+        commands,
+        PathBuf::new(),
+        Arc::new(QuotaStore::default()),
+        detected_ips_tx,
+        runtime_log_filter(),
+        runtime_watch_tx,
+    );
+
+    tokio::time::timeout(Duration::from_secs(1), handle.quiesce())
+        .await
+        .unwrap();
+    let result = control
+        .submit(
+            Arc::new(ProxyConfig::default()),
+            "revision".to_string(),
+            ReloadRequest::default(),
+        )
+        .await;
+
+    assert_eq!(result, Err(ReloadSubmitError::MaestroUnavailable));
+    runtime.stop_sessions().await;
+}