Atomic Maestro Sessions + Shutdown gate

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

src/api/mod.rs

@@ -23,7 +23,7 @@ use tracing::{debug, info, warn};
 use crate::config::ApiGrayAction;
 use crate::ip_tracker::UserIpTracker;
 use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
-use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError};
+use crate::maestro::reload::{ReloadAccepted, ReloadControl, ReloadRequest, ReloadSubmitError};
 use crate::proxy::route_mode::RouteRuntimeController;
 use crate::proxy::shared_state::ProxySharedState;
 use crate::startup::StartupTracker;
@@ -37,6 +37,8 @@ mod events;
 mod http_utils;
 mod model;
 mod patch;
+#[cfg(test)]
+mod reload_tests;
 mod runtime_edge;
 mod runtime_init;
 mod runtime_min;
@@ -182,6 +184,35 @@ fn reload_status_route_id(path: &str) -> Option<u64> {
         .and_then(|id| id.parse().ok())
 }
 
+async fn submit_reload_from_disk(
+    config_path: &std::path::Path,
+    mutation_lock: &Mutex<()>,
+    reload_control: &ReloadControl,
+    expected_revision: Option<&str>,
+    request: ReloadRequest,
+) -> Result<(ReloadAccepted, String), ApiFailure> {
+    let _guard = mutation_lock.lock().await;
+    ensure_expected_revision(config_path, expected_revision).await?;
+    let revision = current_revision(config_path).await?;
+    let config = Arc::new(load_config_for_reload(config_path).await?);
+    let accepted = reload_control
+        .submit(config, revision.clone(), request)
+        .await
+        .map_err(|error| match error {
+            ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
+                StatusCode::CONFLICT,
+                "reload_in_progress",
+                format!("Reload {} is already in progress", reload_id),
+            ),
+            ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
+                StatusCode::SERVICE_UNAVAILABLE,
+                "maestro_unavailable",
+                "Maestro reload coordinator is unavailable",
+            ),
+        })?;
+    Ok((accepted, revision))
+}
+
 fn allowed_methods_for_path(path: &str) -> Option<&'static str> {
     match path {
         "/v1/health"
@@ -740,26 +771,14 @@ async fn handle(
                     .unwrap_or_default();
                 request.validate().map_err(ApiFailure::bad_request)?;
 
-                let _guard = shared.mutation_lock.lock().await;
-                ensure_expected_revision(&shared.config_path, expected_revision.as_deref()).await?;
-                let revision = current_revision(&shared.config_path).await?;
-                let config = Arc::new(load_config_for_reload(&shared.config_path).await?);
-                let accepted = shared
-                    .reload_control
-                    .submit(config, revision.clone(), request)
-                    .await
-                    .map_err(|error| match error {
-                        ReloadSubmitError::InProgress(reload_id) => ApiFailure::new(
-                            StatusCode::CONFLICT,
-                            "reload_in_progress",
-                            format!("Reload {} is already in progress", reload_id),
-                        ),
-                        ReloadSubmitError::MaestroUnavailable => ApiFailure::new(
-                            StatusCode::SERVICE_UNAVAILABLE,
-                            "maestro_unavailable",
-                            "Maestro reload coordinator is unavailable",
-                        ),
-                    })?;
+                let (accepted, revision) = submit_reload_from_disk(
+                    &shared.config_path,
+                    shared.mutation_lock.as_ref(),
+                    &shared.reload_control,
+                    expected_revision.as_deref(),
+                    request,
+                )
+                .await?;
                 Ok(success_response(StatusCode::ACCEPTED, accepted, revision))
             }
             ("PATCH", "/v1/config") => {

src/api/reload_tests.rs

@@ -0,0 +1,119 @@
+use super::*;
+use crate::config::ProxyConfig;
+
+async fn config_file() -> (tempfile::TempDir, PathBuf, String) {
+    let directory = tempfile::tempdir().unwrap();
+    let path = directory.path().join("config.toml");
+    let mut config = ProxyConfig::default();
+    config.server.max_connections = 4_242;
+    let body = toml::to_string_pretty(&config).unwrap();
+    tokio::fs::write(&path, &body).await.unwrap();
+    let revision = config_store::compute_revision(&body);
+    (directory, path, revision)
+}
+
+#[tokio::test]
+async fn reload_submission_uses_matching_disk_revision_and_snapshot() {
+    let (_directory, path, revision) = config_file().await;
+    let mutation_lock = Mutex::new(());
+    let (control, mut commands) = ReloadControl::channel(1);
+    let request = ReloadRequest::default();
+
+    let (accepted, response_revision) = submit_reload_from_disk(
+        &path,
+        &mutation_lock,
+        &control,
+        Some(&revision),
+        request.clone(),
+    )
+    .await
+    .unwrap();
+    let command = commands.recv().await.unwrap();
+
+    assert_eq!(response_revision, revision);
+    assert_eq!(accepted.config_revision, revision);
+    assert_eq!(command.config_revision, revision);
+    assert_eq!(command.request, request);
+    assert_eq!(command.config.server.max_connections, 4_242);
+}
+
+#[tokio::test]
+async fn revision_conflict_rejects_without_enqueuing_reload() {
+    let (_directory, path, _revision) = config_file().await;
+    let mutation_lock = Mutex::new(());
+    let (control, _commands) = ReloadControl::channel(1);
+
+    let error = submit_reload_from_disk(
+        &path,
+        &mutation_lock,
+        &control,
+        Some("stale-revision"),
+        ReloadRequest::default(),
+    )
+    .await
+    .unwrap_err();
+
+    assert_eq!(error.status, StatusCode::CONFLICT);
+    assert_eq!(error.code, "revision_conflict");
+    assert_eq!(control.in_progress().await, None);
+}
+
+#[tokio::test]
+async fn reload_conflict_and_closed_coordinator_map_to_http_contract() {
+    let (_directory, path, _revision) = config_file().await;
+    let mutation_lock = Mutex::new(());
+    let (control, mut commands) = ReloadControl::channel(1);
+    let _accepted = submit_reload_from_disk(
+        &path,
+        &mutation_lock,
+        &control,
+        None,
+        ReloadRequest::default(),
+    )
+    .await
+    .unwrap();
+    let _command = commands.recv().await.unwrap();
+
+    let conflict = submit_reload_from_disk(
+        &path,
+        &mutation_lock,
+        &control,
+        None,
+        ReloadRequest::default(),
+    )
+    .await
+    .unwrap_err();
+    assert_eq!(conflict.status, StatusCode::CONFLICT);
+    assert_eq!(conflict.code, "reload_in_progress");
+
+    control.fail(1, "test cleanup").await;
+    drop(commands);
+    let unavailable = submit_reload_from_disk(
+        &path,
+        &mutation_lock,
+        &control,
+        None,
+        ReloadRequest::default(),
+    )
+    .await
+    .unwrap_err();
+    assert_eq!(unavailable.status, StatusCode::SERVICE_UNAVAILABLE);
+    assert_eq!(unavailable.code, "maestro_unavailable");
+}
+
+#[test]
+fn reload_routes_expose_only_documented_methods_and_ids() {
+    assert_eq!(
+        allowed_methods_for_path("/v1/system/reload"),
+        Some(ALLOW_POST)
+    );
+    assert_eq!(
+        allowed_methods_for_path("/v1/system/reload/42"),
+        Some(ALLOW_GET)
+    );
+    assert_eq!(reload_status_route_id("/v1/system/reload/42"), Some(42));
+    assert_eq!(
+        reload_status_route_id("/v1/system/reload/not-a-number"),
+        None
+    );
+}

src/api/runtime_watch.rs

@@ -14,19 +14,20 @@ pub(super) fn spawn_runtime_watchers(
     runtime_state: Arc<ApiRuntimeState>,
     runtime_events: Arc<ApiEventStore>,
 ) {
-    spawn_config_watcher(
+    let _config_watcher = spawn_config_watcher(
         runtime_watch_rx.clone(),
         runtime_state.clone(),
         runtime_events.clone(),
     );
-    spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
+    let _admission_watcher =
+        spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
 }
 
 fn spawn_config_watcher(
     mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
     runtime_state: Arc<ApiRuntimeState>,
     runtime_events: Arc<ApiEventStore>,
-) {
+) -> tokio::task::JoinHandle<()> {
     tokio::spawn(async move {
         let Some(mut current) = runtime_watch_rx.borrow().clone() else {
             return;
@@ -77,14 +78,14 @@ fn spawn_config_watcher(
                 }
             }
         }
-    });
+    })
 }
 
 fn spawn_admission_watcher(
     mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
     runtime_state: Arc<ApiRuntimeState>,
     runtime_events: Arc<ApiEventStore>,
-) {
+) -> tokio::task::JoinHandle<()> {
     tokio::spawn(async move {
         let Some(mut current) = runtime_watch_rx.borrow().clone() else {
             return;
@@ -123,7 +124,7 @@ fn spawn_admission_watcher(
                 }
             }
         }
-    });
+    })
 }
 
 fn active_generation_id(
@@ -283,4 +284,36 @@ mod tests {
             3
         );
     }
+
+    #[tokio::test]
+    async fn watcher_recovers_from_closed_generation_and_exits_with_process_channel() {
+        let (initial, initial_config_tx, _initial_admission_tx) = state(1);
+        let (runtime_watch_tx, runtime_watch_rx) = watch::channel(Some(initial));
+        let runtime_state = runtime_state();
+        let events = Arc::new(ApiEventStore::new(16));
+        let watcher = spawn_config_watcher(runtime_watch_rx, runtime_state.clone(), events.clone());
+        drop(initial_config_tx);
+        tokio::task::yield_now().await;
+
+        let (next, next_config_tx, _next_admission_tx) = state(2);
+        runtime_watch_tx.send_replace(Some(next));
+        wait_for_count(&runtime_state, 1).await;
+        next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
+        wait_for_count(&runtime_state, 2).await;
+
+        drop(runtime_watch_tx);
+        tokio::time::timeout(Duration::from_secs(1), watcher)
+            .await
+            .unwrap()
+            .unwrap();
+        assert_eq!(
+            events
+                .snapshot(16)
+                .events
+                .iter()
+                .filter(|event| event.event_type == "config.reload.applied")
+                .count(),
+            2
+        );
+    }
 }