Harden Maestro reload lifecycle and readiness barriers

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

src/api/mod.rs

@@ -20,9 +20,9 @@ use tokio::sync::{Mutex, RwLock, Semaphore, watch};
 use tokio::time::timeout;
 use tracing::{debug, info, warn};
 
-use crate::config::{ApiGrayAction, ProxyConfig};
+use crate::config::ApiGrayAction;
 use crate::ip_tracker::UserIpTracker;
-use crate::maestro::generation::RuntimeGeneration;
+use crate::maestro::generation::{RuntimeGeneration, RuntimeWatchState};
 use crate::maestro::reload::{ReloadControl, ReloadRequest, ReloadSubmitError};
 use crate::proxy::route_mode::RouteRuntimeController;
 use crate::proxy::shared_state::ProxySharedState;
@@ -240,8 +240,6 @@ pub async fn serve(
     route_runtime: Arc<RouteRuntimeController>,
     proxy_shared: Arc<ProxySharedState>,
     upstream_manager: Arc<UpstreamManager>,
-    config_rx: watch::Receiver<Arc<ProxyConfig>>,
-    admission_rx: watch::Receiver<bool>,
     config_path: PathBuf,
     quota_state_path: PathBuf,
     detected_ips_rx: watch::Receiver<(Option<IpAddr>, Option<IpAddr>)>,
@@ -249,6 +247,7 @@ pub async fn serve(
     startup_tracker: Arc<StartupTracker>,
     reload_control: ReloadControl,
     mut active_runtime_rx: watch::Receiver<Option<Arc<ArcSwap<RuntimeGeneration>>>>,
+    mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
 ) {
     let active_runtime = loop {
         if let Some(active_runtime) = active_runtime_rx.borrow().clone() {
@@ -259,6 +258,17 @@ pub async fn serve(
             return;
         }
     };
+    let initial_watch_state = loop {
+        if let Some(watch_state) = runtime_watch_rx.borrow().clone() {
+            break watch_state;
+        }
+        if runtime_watch_rx.changed().await.is_err() {
+            warn!("Runtime watch channel closed before API bootstrap");
+            return;
+        }
+    };
+    let config_rx = initial_watch_state.config_rx.clone();
+    let admission_rx = initial_watch_state.admission_rx.clone();
     let listener = match TcpListener::bind(listen).await {
         Ok(listener) => listener,
         Err(error) => {
@@ -306,8 +316,7 @@ pub async fn serve(
     });
 
     spawn_runtime_watchers(
-        config_rx.clone(),
-        admission_rx.clone(),
+        runtime_watch_rx,
         runtime_state.clone(),
         shared.runtime_events.clone(),
     );

src/api/runtime_watch.rs

@@ -4,63 +4,283 @@ use std::time::{SystemTime, UNIX_EPOCH};
 
 use tokio::sync::watch;
 
-use crate::config::ProxyConfig;
+use crate::maestro::generation::RuntimeWatchState;
 
 use super::ApiRuntimeState;
 use super::events::ApiEventStore;
 
 pub(super) fn spawn_runtime_watchers(
-    config_rx: watch::Receiver<Arc<ProxyConfig>>,
-    admission_rx: watch::Receiver<bool>,
+    runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
     runtime_state: Arc<ApiRuntimeState>,
     runtime_events: Arc<ApiEventStore>,
 ) {
-    let mut config_rx_reload = config_rx;
-    let runtime_state_reload = runtime_state.clone();
-    let runtime_events_reload = runtime_events.clone();
-    tokio::spawn(async move {
-        loop {
-            if config_rx_reload.changed().await.is_err() {
-                break;
-            }
-            runtime_state_reload
-                .config_reload_count
-                .fetch_add(1, Ordering::Relaxed);
-            runtime_state_reload
-                .last_config_reload_epoch_secs
-                .store(now_epoch_secs(), Ordering::Relaxed);
-            runtime_events_reload.record("config.reload.applied", "config receiver updated");
-        }
-    });
+    spawn_config_watcher(
+        runtime_watch_rx.clone(),
+        runtime_state.clone(),
+        runtime_events.clone(),
+    );
+    spawn_admission_watcher(runtime_watch_rx, runtime_state, runtime_events);
+}
 
-    let mut admission_rx_watch = admission_rx;
+fn spawn_config_watcher(
+    mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
+    runtime_state: Arc<ApiRuntimeState>,
+    runtime_events: Arc<ApiEventStore>,
+) {
     tokio::spawn(async move {
-        runtime_state
-            .admission_open
-            .store(*admission_rx_watch.borrow(), Ordering::Relaxed);
-        runtime_events.record(
-            "admission.state",
-            format!("accepting_new_connections={}", *admission_rx_watch.borrow()),
-        );
+        let Some(mut current) = runtime_watch_rx.borrow().clone() else {
+            return;
+        };
         loop {
-            if admission_rx_watch.changed().await.is_err() {
-                break;
+            tokio::select! {
+                biased;
+                changed = runtime_watch_rx.changed() => {
+                    if changed.is_err() {
+                        break;
+                    }
+                    let Some(next) = runtime_watch_rx.borrow().clone() else {
+                        continue;
+                    };
+                    if next.generation_id != current.generation_id {
+                        current = next;
+                        record_config_reload(
+                            &runtime_state,
+                            &runtime_events,
+                            format!("runtime generation {} activated", current.generation_id),
+                        );
+                    }
+                }
+                changed = current.config_rx.changed() => {
+                    if changed.is_err() {
+                        let Some(next) = wait_for_new_generation(
+                            &mut runtime_watch_rx,
+                            current.generation_id,
+                        ).await else {
+                            break;
+                        };
+                        current = next;
+                        record_config_reload(
+                            &runtime_state,
+                            &runtime_events,
+                            format!("runtime generation {} activated", current.generation_id),
+                        );
+                        continue;
+                    }
+                    if active_generation_id(&runtime_watch_rx) != Some(current.generation_id) {
+                        continue;
+                    }
+                    record_config_reload(
+                        &runtime_state,
+                        &runtime_events,
+                        format!("generation {} config receiver updated", current.generation_id),
+                    );
+                }
             }
-            let admission_open = *admission_rx_watch.borrow();
-            runtime_state
-                .admission_open
-                .store(admission_open, Ordering::Relaxed);
-            runtime_events.record(
-                "admission.state",
-                format!("accepting_new_connections={}", admission_open),
-            );
         }
     });
 }
 
+fn spawn_admission_watcher(
+    mut runtime_watch_rx: watch::Receiver<Option<RuntimeWatchState>>,
+    runtime_state: Arc<ApiRuntimeState>,
+    runtime_events: Arc<ApiEventStore>,
+) {
+    tokio::spawn(async move {
+        let Some(mut current) = runtime_watch_rx.borrow().clone() else {
+            return;
+        };
+        record_admission_state(&runtime_state, &runtime_events, &current);
+        loop {
+            tokio::select! {
+                biased;
+                changed = runtime_watch_rx.changed() => {
+                    if changed.is_err() {
+                        break;
+                    }
+                    let Some(next) = runtime_watch_rx.borrow().clone() else {
+                        continue;
+                    };
+                    if next.generation_id != current.generation_id {
+                        current = next;
+                        record_admission_state(&runtime_state, &runtime_events, &current);
+                    }
+                }
+                changed = current.admission_rx.changed() => {
+                    if changed.is_err() {
+                        let Some(next) = wait_for_new_generation(
+                            &mut runtime_watch_rx,
+                            current.generation_id,
+                        ).await else {
+                            break;
+                        };
+                        current = next;
+                        record_admission_state(&runtime_state, &runtime_events, &current);
+                        continue;
+                    }
+                    if active_generation_id(&runtime_watch_rx) == Some(current.generation_id) {
+                        record_admission_state(&runtime_state, &runtime_events, &current);
+                    }
+                }
+            }
+        }
+    });
+}
+
+fn active_generation_id(
+    runtime_watch_rx: &watch::Receiver<Option<RuntimeWatchState>>,
+) -> Option<u64> {
+    runtime_watch_rx
+        .borrow()
+        .as_ref()
+        .map(|state| state.generation_id)
+}
+
+async fn wait_for_new_generation(
+    runtime_watch_rx: &mut watch::Receiver<Option<RuntimeWatchState>>,
+    previous_generation_id: u64,
+) -> Option<RuntimeWatchState> {
+    loop {
+        if let Some(state) = runtime_watch_rx.borrow().clone()
+            && state.generation_id != previous_generation_id
+        {
+            return Some(state);
+        }
+        if runtime_watch_rx.changed().await.is_err() {
+            return None;
+        }
+    }
+}
+
+fn record_config_reload(
+    runtime_state: &ApiRuntimeState,
+    runtime_events: &ApiEventStore,
+    context: String,
+) {
+    runtime_state
+        .config_reload_count
+        .fetch_add(1, Ordering::Relaxed);
+    runtime_state
+        .last_config_reload_epoch_secs
+        .store(now_epoch_secs(), Ordering::Relaxed);
+    runtime_events.record("config.reload.applied", context);
+}
+
+fn record_admission_state(
+    runtime_state: &ApiRuntimeState,
+    runtime_events: &ApiEventStore,
+    current: &RuntimeWatchState,
+) {
+    let admission_open = *current.admission_rx.borrow();
+    runtime_state
+        .admission_open
+        .store(admission_open, Ordering::Relaxed);
+    runtime_events.record(
+        "admission.state",
+        format!(
+            "generation={} accepting_new_connections={}",
+            current.generation_id, admission_open
+        ),
+    );
+}
+
 fn now_epoch_secs() -> u64 {
     SystemTime::now()
         .duration_since(UNIX_EPOCH)
         .unwrap_or_default()
         .as_secs()
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::config::ProxyConfig;
+    use std::sync::atomic::{AtomicBool, AtomicU64};
+    use std::time::Duration;
+
+    fn state(
+        generation_id: u64,
+    ) -> (
+        RuntimeWatchState,
+        watch::Sender<Arc<ProxyConfig>>,
+        watch::Sender<bool>,
+    ) {
+        let (config_tx, config_rx) = watch::channel(Arc::new(ProxyConfig::default()));
+        let (admission_tx, admission_rx) = watch::channel(true);
+        (
+            RuntimeWatchState {
+                generation_id,
+                config_rx,
+                admission_rx,
+            },
+            config_tx,
+            admission_tx,
+        )
+    }
+
+    fn runtime_state() -> Arc<ApiRuntimeState> {
+        Arc::new(ApiRuntimeState {
+            process_started_at_epoch_secs: 1,
+            config_reload_count: AtomicU64::new(0),
+            last_config_reload_epoch_secs: AtomicU64::new(0),
+            admission_open: AtomicBool::new(false),
+        })
+    }
+
+    async fn wait_for_count(runtime_state: &ApiRuntimeState, expected: u64) {
+        tokio::time::timeout(Duration::from_secs(1), async {
+            loop {
+                if runtime_state.config_reload_count.load(Ordering::Relaxed) == expected {
+                    break;
+                }
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+    }
+
+    #[tokio::test]
+    async fn watchers_follow_only_the_active_generation() {
+        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));
+        spawn_runtime_watchers(runtime_watch_rx, runtime_state.clone(), events.clone());
+        tokio::task::yield_now().await;
+
+        assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 0);
+        initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
+        wait_for_count(&runtime_state, 1).await;
+
+        let (next, next_config_tx, next_admission_tx) = state(2);
+        runtime_watch_tx.send_replace(Some(next));
+        wait_for_count(&runtime_state, 2).await;
+
+        initial_config_tx.send_replace(Arc::new(ProxyConfig::default()));
+        initial_admission_tx.send_replace(false);
+        tokio::task::yield_now().await;
+        assert_eq!(runtime_state.config_reload_count.load(Ordering::Relaxed), 2);
+        assert!(runtime_state.admission_open.load(Ordering::Relaxed));
+
+        next_config_tx.send_replace(Arc::new(ProxyConfig::default()));
+        next_admission_tx.send_replace(false);
+        wait_for_count(&runtime_state, 3).await;
+        tokio::time::timeout(Duration::from_secs(1), async {
+            while runtime_state.admission_open.load(Ordering::Relaxed) {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .unwrap();
+
+        let snapshot = events.snapshot(16);
+        assert_eq!(
+            snapshot
+                .events
+                .iter()
+                .filter(|event| event.event_type == "config.reload.applied")
+                .count(),
+            3
+        );
+    }
+}

src/api/runtime_zero.rs

@@ -1,4 +1,5 @@
 use std::sync::atomic::Ordering;
+use std::time::{SystemTime, UNIX_EPOCH};
 
 use serde::Serialize;
 
@@ -162,7 +163,7 @@ pub(super) fn build_system_info_data(
         build_time_utc,
         rustc_version,
         process_started_at_epoch_secs: shared.runtime_state.process_started_at_epoch_secs,
-        uptime_seconds: shared.stats.uptime_secs(),
+        uptime_seconds: process_uptime_seconds(shared.runtime_state.process_started_at_epoch_secs),
         config_path: shared.config_path.display().to_string(),
         config_hash: revision.to_string(),
         config_reload_count: shared
@@ -173,6 +174,18 @@ pub(super) fn build_system_info_data(
     }
 }
 
+fn process_uptime_seconds(process_started_at_epoch_secs: u64) -> f64 {
+    let now_epoch_secs = SystemTime::now()
+        .duration_since(UNIX_EPOCH)
+        .unwrap_or_default()
+        .as_secs();
+    process_uptime_seconds_at(process_started_at_epoch_secs, now_epoch_secs)
+}
+
+fn process_uptime_seconds_at(process_started_at_epoch_secs: u64, now_epoch_secs: u64) -> f64 {
+    now_epoch_secs.saturating_sub(process_started_at_epoch_secs) as f64
+}
+
 pub(super) async fn build_runtime_gates_data(
     shared: &ApiShared,
     cfg: &ProxyConfig,
@@ -339,3 +352,14 @@ fn me_writer_pick_mode_label(mode: MeWriterPickMode) -> &'static str {
         MeWriterPickMode::P2c => "p2c",
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::process_uptime_seconds_at;
+
+    #[test]
+    fn process_uptime_is_monotonic_and_saturating() {
+        assert_eq!(process_uptime_seconds_at(100, 135), 35.0);
+        assert_eq!(process_uptime_seconds_at(135, 100), 0.0);
+    }
+}