use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use arc_swap::ArcSwap; use tokio::sync::Barrier; use super::*; use crate::config::ProxyConfig; use crate::maestro::generation::test_runtime_generation; fn test_runtime() -> ( Arc, Arc, ) { let generation = test_runtime_generation(1, ProxyConfig::default()); let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(Arc::clone(&generation)))); (runtime, generation) } async fn stop_runtime( runtime: Arc, generation: Arc, ) { runtime.shutdown().await; generation.stop_sessions().await; generation.stop_background_tasks().await; } #[tokio::test] async fn pause_is_idempotent_and_resume_advances_only_real_transitions() { let (runtime, generation) = test_runtime(); let initial = runtime.operator_lifecycle_status(); assert_eq!(initial.state, OperatorLifecycleState::Running); assert_eq!(initial.epoch, 0); let paused = runtime.pause_operator().await.unwrap(); assert_eq!(paused.state, OperatorLifecycleState::Paused); assert!(!paused.admission_open); let repeated_pause = runtime.pause_operator().await.unwrap(); assert_eq!(repeated_pause.epoch, paused.epoch); let resumed = runtime.resume_operator().await.unwrap(); assert_eq!(resumed.state, OperatorLifecycleState::Running); assert!(resumed.admission_open); let repeated_resume = runtime.resume_operator().await.unwrap(); assert_eq!(repeated_resume.epoch, resumed.epoch); stop_runtime(runtime, generation).await; } #[tokio::test] async fn pause_waits_for_pre_cutover_admission_and_rejects_late_registration() { let (runtime, generation) = test_runtime(); let registration = runtime.try_operator_admission().unwrap(); let pause_runtime = Arc::clone(&runtime); let pause = tokio::spawn(async move { pause_runtime.pause_operator().await.unwrap() }); tokio::task::yield_now().await; assert!(!pause.is_finished()); assert_eq!( runtime.try_operator_admission().err(), Some(super::super::ManagerError::AdmissionPaused) ); drop(registration); let paused = pause.await.unwrap(); assert_eq!(paused.state, OperatorLifecycleState::Paused); stop_runtime(runtime, generation).await; } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn pause_fence_leaves_no_late_admission_commits_under_scheduler_pressure() { const ATTEMPTS: usize = 10_000; let (runtime, generation) = test_runtime(); let start = Arc::new(Barrier::new(ATTEMPTS + 1)); let committed = Arc::new(AtomicUsize::new(0)); let mut attempts = tokio::task::JoinSet::new(); for _ in 0..ATTEMPTS { let runtime = Arc::clone(&runtime); let start = Arc::clone(&start); let committed = Arc::clone(&committed); attempts.spawn(async move { start.wait().await; if let Ok(_registration) = runtime.try_operator_admission() { tokio::task::yield_now().await; committed.fetch_add(1, Ordering::Release); } }); } start.wait().await; runtime.pause_operator().await.unwrap(); let committed_at_pause = committed.load(Ordering::Acquire); while let Some(result) = attempts.join_next().await { result.unwrap(); } assert_eq!(committed.load(Ordering::Acquire), committed_at_pause); assert!(matches!( runtime.try_operator_admission(), Err(super::super::ManagerError::AdmissionPaused) )); stop_runtime(runtime, generation).await; } #[tokio::test] async fn empty_drain_completes_gracefully_and_stays_closed_until_resume() { let (runtime, generation) = test_runtime(); let accepted = runtime.drain_operator(Duration::from_secs(30)).await.unwrap(); assert_eq!(accepted.state, OperatorLifecycleState::Draining); let completed = tokio::time::timeout(Duration::from_secs(1), async { loop { let status = runtime.operator_lifecycle_status(); if status.state == OperatorLifecycleState::Drained { break status; } tokio::task::yield_now().await; } }) .await .unwrap(); let drain = completed.drain.unwrap(); assert_eq!(drain.state, OperatorDrainState::Completed); assert_eq!(drain.outcome, Some(OperatorDrainOutcome::Graceful)); assert!(!completed.admission_open); let resumed = runtime.resume_operator().await.unwrap(); assert_eq!(resumed.state, OperatorLifecycleState::Running); stop_runtime(runtime, generation).await; } #[tokio::test] async fn resume_after_force_commit_cancels_wait_but_preserves_force_evidence() { let (runtime, generation) = test_runtime(); let sequence = 7; let cancellation = CancellationToken::new(); { let mut inner = runtime.operator_lifecycle.inner.lock(); runtime.operator_lifecycle.admission.close(); inner.active = Some(ActiveDrain { sequence, cancellation, }); inner.drain = Some(OperatorDrainStatus { operation_id: format!("wd1.{}.{sequence:016x}", runtime.runtime_instance()), state: OperatorDrainState::Draining, outcome: None, timeout_secs: 30, started_epoch_millis: 1, deadline_epoch_millis: 30_001, completed_epoch_millis: None, remaining_sessions: 1, remaining_streams: 0, remaining_websockets: 0, force_close_signalled: false, }); runtime .operator_lifecycle .transition_locked(&mut inner, OperatorLifecycleState::Draining); runtime.operator_lifecycle.publish_locked(&inner); } assert!(runtime.operator_lifecycle.commit_force( sequence, WorkCounts { sessions: 1, streams: 0, websockets: 0, } )); let resumed = runtime.resume_operator().await.unwrap(); assert_eq!(resumed.state, OperatorLifecycleState::Running); let drain = resumed.drain.unwrap(); assert_eq!(drain.state, OperatorDrainState::Cancelled); assert_eq!(drain.outcome, Some(OperatorDrainOutcome::Cancelled)); assert!(drain.force_close_signalled); stop_runtime(runtime, generation).await; } #[tokio::test] async fn reload_cannot_reopen_operator_pause_and_resume_cannot_override_disabled_config() { let mut initial_config = ProxyConfig::default(); initial_config.web.enabled = true; let initial = test_runtime_generation(1, initial_config); let active = Arc::new(ArcSwap::from(Arc::clone(&initial))); let runtime = WebProcessRuntime::start(Arc::clone(&active)); let paused = runtime.pause_operator().await.unwrap(); assert!(!paused.admission_open); let disabled = test_runtime_generation(2, ProxyConfig::default()); active.store(Arc::clone(&disabled)); let after_reload = runtime.operator_lifecycle_status(); assert_eq!(after_reload.state, OperatorLifecycleState::Paused); let resumed = runtime.resume_operator().await.unwrap(); assert!(resumed.admission_open); assert!(!resumed.effective_new_work_admission); runtime.shutdown().await; initial.stop_sessions().await; initial.stop_background_tasks().await; disabled.stop_sessions().await; disabled.stop_background_tasks().await; } #[tokio::test] async fn shutdown_terminally_prevents_resume() { let (runtime, generation) = test_runtime(); runtime.pause_operator().await.unwrap(); let drain = runtime.begin_shutdown(); assert!(matches!( runtime.resume_operator().await, Err(OperatorLifecycleError::Closed) )); let _ = drain .wait_until(tokio::time::Instant::now() + Duration::from_secs(1)) .await; generation.stop_sessions().await; generation.stop_background_tasks().await; }