From 012dc07a988139ca6e89701a74d2b365b34fd825 Mon Sep 17 00:00:00 2001 From: Alexey <247128645+axkurcom@users.noreply.github.com> Date: Fri, 28 Aug 2026 22:39:18 +0300 Subject: [PATCH] WEB Ephemeral Lifecycle Co-Authored-By: brekotis <93345790+brekotis@users.noreply.github.com> --- src/api/web_runtime.rs | 146 +++-- src/api/web_runtime/request.rs | 18 + src/api/web_runtime/tests.rs | 51 ++ src/web/http/operator_lifecycle_tests.rs | 345 ++++++++++++ src/web/http/session.rs | 5 +- src/web/http/tests.rs | 3 + src/web/manager.rs | 16 +- src/web/manager/admission.rs | 24 +- src/web/manager/credentials.rs | 1 + src/web/manager/lifecycle.rs | 3 + src/web/manager/operator_lifecycle.rs | 550 +++++++++++++++++++ src/web/manager/operator_lifecycle/status.rs | 90 +++ src/web/manager/operator_lifecycle/tests.rs | 220 ++++++++ src/web/manager/session_creation.rs | 2 + src/web/manager/websocket.rs | 1 + src/web/session/uplink.rs | 14 +- src/web/session/websocket.rs | 6 +- src/web/session/websocket/tests.rs | 4 +- 18 files changed, 1434 insertions(+), 65 deletions(-) create mode 100644 src/api/web_runtime/tests.rs create mode 100644 src/web/http/operator_lifecycle_tests.rs create mode 100644 src/web/manager/operator_lifecycle.rs create mode 100644 src/web/manager/operator_lifecycle/status.rs create mode 100644 src/web/manager/operator_lifecycle/tests.rs diff --git a/src/api/web_runtime.rs b/src/api/web_runtime.rs index b9a25be..b64ad84 100644 --- a/src/api/web_runtime.rs +++ b/src/api/web_runtime.rs @@ -1,5 +1,5 @@ use std::sync::Arc; -use std::time::Instant; +use std::time::{Duration, Instant}; use http_body_util::Full; use hyper::body::{Bytes, Incoming}; @@ -13,12 +13,14 @@ use super::model::ApiFailure; use super::{ALLOW_GET, ALLOW_POST, ApiShared}; use crate::config::ProxyConfig; use crate::web::control::{WebRuntimeLifecycle, WebRuntimePublication}; -use crate::web::manager::{ControlError, SessionDetail, WebProcessRuntime}; +use crate::web::manager::{ + ControlError, OperatorLifecycleError, SessionDetail, WebProcessRuntime, +}; // Exact JSON DTOs and strict query parsing stay independent from route dispatch. mod request; use request::{ - CloseRequest, RuntimeInstanceRequest, parse_session_query, parse_session_ref, + CloseRequest, DrainRequest, RuntimeInstanceRequest, parse_session_query, parse_session_ref, valid_runtime_instance, }; @@ -27,6 +29,9 @@ const SESSIONS_PATH: &str = "/v1/runtime/web/sessions"; const CLOSE_PATH: &str = "/v1/runtime/web/sessions/close"; const DEBUG_CLEAR_PATH: &str = "/v1/runtime/web/debug/clear"; const LEARNING_RESET_PATH: &str = "/v1/runtime/web/carrier-learning/reset"; +const LIFECYCLE_PAUSE_PATH: &str = "/v1/runtime/web/lifecycle/pause"; +const LIFECYCLE_DRAIN_PATH: &str = "/v1/runtime/web/lifecycle/drain"; +const LIFECYCLE_RESUME_PATH: &str = "/v1/runtime/web/lifecycle/resume"; const SESSION_DETAIL_PREFIX: &str = "/v1/runtime/web/sessions/"; const OPERATION_PREFIX: &str = "/v1/runtime/web/operations/"; const MAX_CONTROL_BODY_BYTES: usize = 64 * 1024; @@ -35,7 +40,12 @@ const MAX_CONTROL_BODY_BYTES: usize = 64 * 1024; pub(super) fn allowed_methods(path: &str) -> Option<&'static str> { match path { STATUS_PATH | SESSIONS_PATH => Some(ALLOW_GET), - CLOSE_PATH | DEBUG_CLEAR_PATH | LEARNING_RESET_PATH => Some(ALLOW_POST), + CLOSE_PATH + | DEBUG_CLEAR_PATH + | LEARNING_RESET_PATH + | LIFECYCLE_PAUSE_PATH + | LIFECYCLE_DRAIN_PATH + | LIFECYCLE_RESUME_PATH => Some(ALLOW_POST), _ if detail_ref(path).is_some() || operation_ref(path).is_some() => Some(ALLOW_GET), _ => None, } @@ -105,6 +115,67 @@ pub(super) async fn handle( .map_err(control_failure)?; Ok(success_response(StatusCode::OK, status, revision)) } + ("POST", LIFECYCLE_PAUSE_PATH) => { + require_mutable(config)?; + reject_query(query)?; + require_json_content_type(&request)?; + let request = read_json::( + request.into_body(), + body_limit.min(MAX_CONTROL_BODY_BYTES), + ) + .await?; + let runtime = control_runtime(shared)?; + require_runtime_instance(&runtime, &request.runtime_instance)?; + let status = runtime.pause_operator().await.map_err(lifecycle_failure)?; + shared.runtime_events.record( + "api.web.lifecycle.pause.ok", + format!("epoch={}", status.epoch), + ); + Ok(success_response(StatusCode::OK, status, revision)) + } + ("POST", LIFECYCLE_DRAIN_PATH) => { + require_mutable(config)?; + reject_query(query)?; + require_json_content_type(&request)?; + let request = read_json::( + request.into_body(), + body_limit.min(MAX_CONTROL_BODY_BYTES), + ) + .await?; + let timeout = drain_timeout(request.timeout_secs)?; + let runtime = control_runtime(shared)?; + require_runtime_instance(&runtime, &request.runtime_instance)?; + let status = runtime + .drain_operator(timeout) + .await + .map_err(lifecycle_failure)?; + shared.runtime_events.record( + "api.web.lifecycle.drain.accepted", + format!( + "epoch={} timeout_secs={}", + status.epoch, request.timeout_secs + ), + ); + Ok(success_response(StatusCode::ACCEPTED, status, revision)) + } + ("POST", LIFECYCLE_RESUME_PATH) => { + require_mutable(config)?; + reject_query(query)?; + require_json_content_type(&request)?; + let request = read_json::( + request.into_body(), + body_limit.min(MAX_CONTROL_BODY_BYTES), + ) + .await?; + let runtime = control_runtime(shared)?; + require_runtime_instance(&runtime, &request.runtime_instance)?; + let status = runtime.resume_operator().await.map_err(lifecycle_failure)?; + shared.runtime_events.record( + "api.web.lifecycle.resume.ok", + format!("epoch={}", status.epoch), + ); + Ok(success_response(StatusCode::OK, status, revision)) + } ("POST", CLOSE_PATH) => { require_mutable(config)?; reject_query(query)?; @@ -195,6 +266,8 @@ struct WebStatusData { listeners: Vec, effective_config_enabled: bool, #[serde(skip_serializing_if = "Option::is_none")] + operator_lifecycle: Option, + #[serde(skip_serializing_if = "Option::is_none")] runtime: Option, } @@ -221,6 +294,7 @@ impl WebStatusData { WebRuntimeLifecycle::DeadlineExceeded => "deadline_exceeded", }) }; + let operator_lifecycle = runtime.map(WebProcessRuntime::operator_lifecycle_status); Self { lifecycle: publication.lifecycle.as_str(), lifecycle_epoch: publication.epoch, @@ -233,6 +307,7 @@ impl WebStatusData { .map(ToString::to_string) .collect(), effective_config_enabled, + operator_lifecycle, runtime: runtime.map(WebProcessRuntime::try_status), } } @@ -369,6 +444,17 @@ fn control_failure(error: ControlError) -> ApiFailure { } } +fn lifecycle_failure(error: OperatorLifecycleError) -> ApiFailure { + match error { + OperatorLifecycleError::Closed => runtime_unavailable(WebRuntimeLifecycle::Draining), + OperatorLifecycleError::OperationInProgress => ApiFailure::new( + StatusCode::CONFLICT, + "web_lifecycle_in_progress", + "Another WEB drain operation is active", + ), + } +} + fn snapshot_busy() -> ApiFailure { ApiFailure::new( StatusCode::SERVICE_UNAVAILABLE, @@ -377,6 +463,15 @@ fn snapshot_busy() -> ApiFailure { ) } +fn drain_timeout(timeout_secs: u64) -> Result { + if !(1..=3600).contains(&timeout_secs) { + return Err(ApiFailure::bad_request( + "timeout_secs must be within 1..=3600", + )); + } + Ok(Duration::from_secs(timeout_secs)) +} + fn reject_query(query: Option<&str>) -> Result<(), ApiFailure> { if query.is_some_and(|query| !query.is_empty()) { return Err(ApiFailure::bad_request( @@ -401,44 +496,5 @@ fn millis(duration: std::time::Duration) -> u64 { } #[cfg(test)] -mod tests { - use super::*; - use hyper::header::HeaderValue; - - #[test] - fn route_table_keeps_status_read_only_and_controls_post_only() { - assert_eq!(allowed_methods(STATUS_PATH), Some(ALLOW_GET)); - assert_eq!(allowed_methods(SESSIONS_PATH), Some(ALLOW_GET)); - assert_eq!(allowed_methods(CLOSE_PATH), Some(ALLOW_POST)); - assert_eq!(allowed_methods(DEBUG_CLEAR_PATH), Some(ALLOW_POST)); - assert_eq!(allowed_methods(LEARNING_RESET_PATH), Some(ALLOW_POST)); - assert_eq!( - allowed_methods("/v1/runtime/web/sessions/ws1.instance.0000000000000001"), - Some(ALLOW_GET) - ); - } - - #[test] - fn control_content_type_is_exact_and_single() { - let exact = Request::builder() - .header(CONTENT_TYPE, "application/json") - .body(()) - .unwrap(); - assert!(require_json_content_type(&exact).is_ok()); - - let parameterized = Request::builder() - .header(CONTENT_TYPE, "application/json; charset=utf-8") - .body(()) - .unwrap(); - assert!(require_json_content_type(¶meterized).is_err()); - - let mut duplicated = Request::builder() - .header(CONTENT_TYPE, "application/json") - .body(()) - .unwrap(); - duplicated - .headers_mut() - .append(CONTENT_TYPE, HeaderValue::from_static("application/json")); - assert!(require_json_content_type(&duplicated).is_err()); - } -} +#[path = "web_runtime/tests.rs"] +mod tests; diff --git a/src/api/web_runtime/request.rs b/src/api/web_runtime/request.rs index 8a2f011..259b96a 100644 --- a/src/api/web_runtime/request.rs +++ b/src/api/web_runtime/request.rs @@ -22,6 +22,16 @@ pub(super) struct RuntimeInstanceRequest { pub(super) runtime_instance: String, } +/// Process-fenced graceful drain request with one bounded relative deadline. +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct DrainRequest { + /// Random process identifier copied from WEB runtime status. + pub(super) runtime_instance: String, + /// Relative drain deadline frozen into one monotonic server deadline. + pub(super) timeout_secs: u64, +} + /// One process-fenced asynchronous close request. #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -331,6 +341,14 @@ mod tests { })) .is_err() ); + assert!( + serde_json::from_value::(serde_json::json!({ + "runtime_instance": runtime_instance, + "timeout_secs": 30, + "extra": true, + })) + .is_err() + ); } #[test] diff --git a/src/api/web_runtime/tests.rs b/src/api/web_runtime/tests.rs new file mode 100644 index 0000000..3fd1d69 --- /dev/null +++ b/src/api/web_runtime/tests.rs @@ -0,0 +1,51 @@ +use super::*; + +use hyper::header::HeaderValue; + +#[test] +fn route_table_keeps_status_read_only_and_controls_post_only() { + assert_eq!(allowed_methods(STATUS_PATH), Some(ALLOW_GET)); + assert_eq!(allowed_methods(SESSIONS_PATH), Some(ALLOW_GET)); + assert_eq!(allowed_methods(CLOSE_PATH), Some(ALLOW_POST)); + assert_eq!(allowed_methods(DEBUG_CLEAR_PATH), Some(ALLOW_POST)); + assert_eq!(allowed_methods(LEARNING_RESET_PATH), Some(ALLOW_POST)); + assert_eq!(allowed_methods(LIFECYCLE_PAUSE_PATH), Some(ALLOW_POST)); + assert_eq!(allowed_methods(LIFECYCLE_DRAIN_PATH), Some(ALLOW_POST)); + assert_eq!(allowed_methods(LIFECYCLE_RESUME_PATH), Some(ALLOW_POST)); + assert_eq!( + allowed_methods("/v1/runtime/web/sessions/ws1.instance.0000000000000001"), + Some(ALLOW_GET) + ); +} + +#[test] +fn control_content_type_is_exact_and_single() { + let exact = Request::builder() + .header(CONTENT_TYPE, "application/json") + .body(()) + .unwrap(); + assert!(require_json_content_type(&exact).is_ok()); + + let parameterized = Request::builder() + .header(CONTENT_TYPE, "application/json; charset=utf-8") + .body(()) + .unwrap(); + assert!(require_json_content_type(¶meterized).is_err()); + + let mut duplicated = Request::builder() + .header(CONTENT_TYPE, "application/json") + .body(()) + .unwrap(); + duplicated + .headers_mut() + .append(CONTENT_TYPE, HeaderValue::from_static("application/json")); + assert!(require_json_content_type(&duplicated).is_err()); +} + +#[test] +fn drain_timeout_is_bounded_to_the_public_contract() { + assert_eq!(drain_timeout(1).unwrap(), Duration::from_secs(1)); + assert_eq!(drain_timeout(3600).unwrap(), Duration::from_secs(3600)); + assert!(drain_timeout(0).is_err()); + assert!(drain_timeout(3601).is_err()); +} diff --git a/src/web/http/operator_lifecycle_tests.rs b/src/web/http/operator_lifecycle_tests.rs new file mode 100644 index 0000000..9cd3275 --- /dev/null +++ b/src/web/http/operator_lifecycle_tests.rs @@ -0,0 +1,345 @@ +use super::*; + +use std::time::Duration; + +use sha2::{Digest, Sha256}; + +use crate::web::manager::{ManagerError, OperatorLifecycleError}; + +fn create_request(bootstrap: &str, hello: &[u8]) -> Vec { + let mut request = format!( + "POST /api/v1/session HTTP/1.1\r\nHost: proxy.example.com\r\nX-Forwarded-For: 192.0.2.10\r\nAuthorization: Bearer {bootstrap}\r\nContent-Type: application/octet-stream\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + hello.len() + ) + .into_bytes(); + request.extend_from_slice(hello); + request +} + +fn negotiated_request( + bootstrap: &str, + hello: &[u8], + attempt: u8, + failure: Option<&str>, +) -> Vec { + let failure = failure + .map(|failure| format!("X-Carrier-Failure: {failure}\r\n")) + .unwrap_or_default(); + let mut request = format!( + "POST /api/v1/session HTTP/1.1\r\nHost: proxy.example.com\r\nX-Forwarded-For: 192.0.2.10\r\nAuthorization: Bearer {bootstrap}\r\nContent-Type: application/octet-stream\r\nX-Carrier-Capabilities: https,https-lanes,websocket,websocket-lanes\r\nX-Carrier-Attempt: {attempt}\r\n{failure}Content-Length: {}\r\nConnection: close\r\n\r\n", + hello.len() + ) + .into_bytes(); + request.extend_from_slice(hello); + request +} + +fn token_hash(token: &str) -> crate::web::manager::TokenHash { + let raw = base64::engine::general_purpose::URL_SAFE_NO_PAD + .decode(token) + .unwrap(); + Sha256::digest(raw).into() +} + +async fn live_runtime() -> ( + Arc, + Arc, + TcpListener, +) { + let generation = test_runtime_generation( + 1, + runtime_config([71; 32], WebCarrier::Https), + ); + let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(Arc::clone(&generation)))); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + (runtime, generation, listener) +} + +fn issue_bootstrap(runtime: &Arc) -> String { + let profile = runtime + .active_generation() + .config() + .web + .runtime + .as_ref() + .unwrap() + .profiles[0] + .clone(); + runtime + .issue_bootstrap(profile, "192.0.2.10".parse().unwrap()) + .unwrap() + .token +} + +async fn create_session( + listener: &TcpListener, + runtime: &Arc, + bootstrap: &str, +) -> (Vec, String) { + let hello = frame::encode(FrameType::Hello, 0, &[1]); + let response = request(listener, runtime, create_request(bootstrap, &hello)).await; + let (headers, _) = split_response(&response); + assert!(headers.starts_with(b"HTTP/1.1 200")); + let token = response_header(headers, "x-session-token").to_string(); + (hello.to_vec(), token) +} + +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_preserves_decoy_retry_and_exact_session_replay() { + let (runtime, generation, listener) = live_runtime().await; + let bootstrap = issue_bootstrap(&runtime); + let hello = frame::encode(FrameType::Hello, 0, &[1]); + + runtime.pause_operator().await.unwrap(); + let paused_create = request( + &listener, + &runtime, + create_request(&bootstrap, &hello), + ) + .await; + let (paused_headers, _) = split_response(&paused_create); + assert!(paused_headers.starts_with(b"HTTP/1.1 503")); + assert_eq!(response_header(paused_headers, "retry-after"), "1"); + + let profile = runtime + .active_generation() + .config() + .web + .runtime + .as_ref() + .unwrap() + .profiles[0] + .clone(); + assert!(matches!( + runtime.issue_bootstrap(profile, "192.0.2.11".parse().unwrap()), + Err(ManagerError::AdmissionPaused) + )); + let bridge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode([71; 32]); + let decoy = format!( + "GET /?bridge={bridge} HTTP/1.1\r\nHost: proxy.example.com\r\nX-Forwarded-For: 192.0.2.11\r\nConnection: close\r\n\r\n" + ) + .into_bytes(); + let decoy = request(&listener, &runtime, decoy).await; + let (decoy_headers, decoy_body) = split_response(&decoy); + assert!(decoy_headers.starts_with(b"HTTP/1.1 200")); + assert!(!decoy_body.windows(11).any(|window| window == b"bootstrap=\"")); + + runtime.resume_operator().await.unwrap(); + let created = request( + &listener, + &runtime, + create_request(&bootstrap, &hello), + ) + .await; + let (created_headers, _) = split_response(&created); + assert!(created_headers.starts_with(b"HTTP/1.1 200")); + let token = response_header(created_headers, "x-session-token").to_string(); + + runtime.pause_operator().await.unwrap(); + let replay = request( + &listener, + &runtime, + create_request(&bootstrap, &hello), + ) + .await; + let (replay_headers, _) = split_response(&replay); + assert!(replay_headers.starts_with(b"HTTP/1.1 200")); + assert_eq!(response_header(replay_headers, "x-session-token"), token); + + stop_runtime(runtime, generation).await; +} + +#[tokio::test] +async fn paused_open_is_stream_local_and_does_not_charge_limit_hits() { + let (runtime, generation, listener) = live_runtime().await; + let bootstrap = issue_bootstrap(&runtime); + let (_, token) = create_session(&listener, &runtime, &bootstrap).await; + let before = serde_json::to_value(runtime.try_status()).unwrap(); + let limit_hits = before["limit_hits"].as_u64().unwrap(); + let streams_rejected = before["streams_rejected"].as_u64().unwrap(); + + runtime.pause_operator().await.unwrap(); + let open = frame::encode(FrameType::Open, 7, &[]); + let mut uplink = format!( + "POST /api/v1/up HTTP/1.1\r\nHost: proxy.example.com\r\nX-Forwarded-For: 192.0.2.10\r\nAuthorization: Bearer {token}\r\nContent-Type: application/octet-stream\r\nX-Up-Seq: 1\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + open.len() + ) + .into_bytes(); + uplink.extend_from_slice(&open); + let rejected = request(&listener, &runtime, uplink).await; + assert!(rejected.starts_with(b"HTTP/1.1 204")); + + let after = serde_json::to_value(runtime.try_status()).unwrap(); + assert_eq!(after["limit_hits"].as_u64(), Some(limit_hits)); + assert_eq!( + after["streams_rejected"].as_u64(), + Some(streams_rejected + 1) + ); + assert_eq!(after["streams"]["live"].as_u64(), Some(0)); + assert!( + runtime + .get_session(token_hash(&token), "proxy.example.com") + .is_ok() + ); + + stop_runtime(runtime, generation).await; +} + +#[tokio::test] +async fn paused_replacement_preserves_old_session_and_attempt_replay() { + let generation = test_runtime_generation( + 1, + negotiation_runtime_config( + [72; 32], + WebCarrier::Https, + false, + Arc::from([WebCarrier::Https, WebCarrier::HttpsLanes]), + ), + ); + let runtime = WebProcessRuntime::start(Arc::new(ArcSwap::from(Arc::clone(&generation)))); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let bootstrap = issue_bootstrap(&runtime); + let hello = frame::encode(FrameType::Hello, 0, &[1]); + let first_request = negotiated_request(&bootstrap, &hello, 1, None); + let first = request(&listener, &runtime, first_request.clone()).await; + let (first_headers, _) = split_response(&first); + assert!(first_headers.starts_with(b"HTTP/1.1 200")); + let first_token = response_header(first_headers, "x-session-token").to_string(); + + runtime.pause_operator().await.unwrap(); + let replacement_request = negotiated_request(&bootstrap, &hello, 2, Some("timeout")); + let rejected = request(&listener, &runtime, replacement_request.clone()).await; + let (rejected_headers, _) = split_response(&rejected); + assert!(rejected_headers.starts_with(b"HTTP/1.1 503")); + assert_eq!(response_header(rejected_headers, "retry-after"), "1"); + assert!( + runtime + .get_session(token_hash(&first_token), "proxy.example.com") + .is_ok() + ); + + let replay = request(&listener, &runtime, first_request).await; + let (replay_headers, _) = split_response(&replay); + assert!(replay_headers.starts_with(b"HTTP/1.1 200")); + assert_eq!(response_header(replay_headers, "x-session-token"), first_token); + + runtime.resume_operator().await.unwrap(); + let replacement = request(&listener, &runtime, replacement_request).await; + let (replacement_headers, _) = split_response(&replacement); + assert!(replacement_headers.starts_with(b"HTTP/1.1 200")); + assert_ne!( + response_header(replacement_headers, "x-session-token"), + first_token + ); + + stop_runtime(runtime, generation).await; +} + +#[tokio::test] +async fn client_delete_during_drain_completes_naturally() { + let (runtime, generation, listener) = live_runtime().await; + let bootstrap = issue_bootstrap(&runtime); + let (_, token) = create_session(&listener, &runtime, &bootstrap).await; + runtime.drain_operator(Duration::from_secs(30)).await.unwrap(); + + let delete = format!( + "DELETE /api/v1/session HTTP/1.1\r\nHost: proxy.example.com\r\nX-Forwarded-For: 192.0.2.10\r\nAuthorization: Bearer {token}\r\nContent-Length: 0\r\nConnection: close\r\n\r\n" + ) + .into_bytes(); + let deleted = request(&listener, &runtime, delete).await; + assert!(deleted.starts_with(b"HTTP/1.1 204")); + + let completed = tokio::time::timeout(Duration::from_secs(1), async { + loop { + let status = runtime.operator_lifecycle_status(); + if serde_json::to_value(&status).unwrap()["state"] == "drained" { + break status; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let completed = serde_json::to_value(&completed).unwrap(); + assert_eq!(completed["drain"]["outcome"], "graceful"); + assert_eq!(completed["drain"]["force_close_signalled"], false); + stop_runtime(runtime, generation).await; +} + +#[tokio::test(start_paused = true)] +async fn deadline_force_closes_all_sessions_and_requires_explicit_resume() { + let (runtime, generation, listener) = live_runtime().await; + let bootstrap = issue_bootstrap(&runtime); + let (_, token) = create_session(&listener, &runtime, &bootstrap).await; + + let accepted = runtime.drain_operator(Duration::from_secs(30)).await.unwrap(); + assert_eq!(serde_json::to_value(&accepted).unwrap()["state"], "draining"); + assert!(matches!( + runtime.drain_operator(Duration::from_secs(30)).await, + Err(OperatorLifecycleError::OperationInProgress) + )); + tokio::time::advance(Duration::from_secs(30)).await; + let completed = tokio::time::timeout(Duration::from_secs(1), async { + loop { + let status = runtime.operator_lifecycle_status(); + if serde_json::to_value(&status).unwrap()["state"] == "drained" { + break status; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let completed = serde_json::to_value(&completed).unwrap(); + assert_eq!(completed["drain"]["state"], "completed"); + assert_eq!(completed["drain"]["outcome"], "forced"); + assert_eq!(completed["drain"]["force_close_signalled"], true); + assert_eq!(completed["admission_open"], false); + assert!( + runtime + .get_session(token_hash(&token), "proxy.example.com") + .is_err() + ); + + let resumed = runtime.resume_operator().await.unwrap(); + let resumed = serde_json::to_value(&resumed).unwrap(); + assert_eq!(resumed["state"], "running"); + assert_eq!(resumed["admission_open"], true); + stop_runtime(runtime, generation).await; +} + +#[tokio::test] +async fn resume_before_deadline_cancels_drain_without_closing_existing_session() { + let (runtime, generation, listener) = live_runtime().await; + let bootstrap = issue_bootstrap(&runtime); + let (_, token) = create_session(&listener, &runtime, &bootstrap).await; + + runtime.drain_operator(Duration::from_secs(30)).await.unwrap(); + let still_draining = runtime.pause_operator().await.unwrap(); + assert_eq!( + serde_json::to_value(&still_draining).unwrap()["state"], + "draining" + ); + let resumed = runtime.resume_operator().await.unwrap(); + let resumed = serde_json::to_value(&resumed).unwrap(); + assert_eq!(resumed["state"], "running"); + assert_eq!(resumed["drain"]["state"], "cancelled"); + assert_eq!(resumed["drain"]["outcome"], "cancelled"); + assert_eq!(resumed["drain"]["force_close_signalled"], false); + assert!( + runtime + .get_session(token_hash(&token), "proxy.example.com") + .is_ok() + ); + + stop_runtime(runtime, generation).await; +} diff --git a/src/web/http/session.rs b/src/web/http/session.rs index 41fcbb6..1834d0a 100644 --- a/src/web/http/session.rs +++ b/src/web/http/session.rs @@ -189,7 +189,10 @@ pub(super) async fn handle_session( response } Err( - error @ (ManagerError::Limit | ManagerError::Backpressure | ManagerError::Concurrent), + error @ (ManagerError::Limit + | ManagerError::Backpressure + | ManagerError::Concurrent + | ManagerError::AdmissionPaused), ) => { runtime.trace().record_profile_lifecycle( client_ip, diff --git a/src/web/http/tests.rs b/src/web/http/tests.rs index f8414c2..01e0330 100644 --- a/src/web/http/tests.rs +++ b/src/web/http/tests.rs @@ -30,6 +30,9 @@ mod session_policy_tests; // Runtime control integration stays separate from carrier protocol scenarios. #[path = "control_tests.rs"] mod control_tests; +// Reversible operator lifecycle coverage stays separate from terminal shutdown tests. +#[path = "operator_lifecycle_tests.rs"] +mod operator_lifecycle_tests; const TEST_CARRIER_DEADLINES_SECS: [u64; 4] = [3, 5, 8, 12]; diff --git a/src/web/manager.rs b/src/web/manager.rs index e249858..e66b189 100644 --- a/src/web/manager.rs +++ b/src/web/manager.rs @@ -34,6 +34,11 @@ mod admission; // Shutdown and expiry work remain outside request-path coordination. mod lifecycle; pub(crate) use lifecycle::WebShutdownOutcome; +// Reversible operator admission stays independent from terminal process shutdown. +mod operator_lifecycle; +pub(crate) use operator_lifecycle::{ + OperatorLifecycleError, OperatorLifecycleStatus, +}; // Queue and WebSocket allocations share one process-owned data-plane budget. mod budget; // WebSocket admission, replacement, and liveness are process-scoped. @@ -79,6 +84,8 @@ pub(crate) enum ManagerError { Committed, /// The process or session has stopped accepting work. Closed, + /// Operator pause or drain temporarily rejects new WEB work. + AdmissionPaused, } impl ManagerError { @@ -92,6 +99,7 @@ impl ManagerError { Self::Concurrent => "concurrent", Self::Committed => "committed", Self::Closed => "closed", + Self::AdmissionPaused => "admission_paused", } } } @@ -155,6 +163,7 @@ pub(crate) struct WebProcessRuntime { websocket_next_id: AtomicU64, websocket_clock: std::time::Instant, websocket_notify: Arc, + operator_lifecycle: operator_lifecycle::OperatorLifecycle, data_budget: Arc, control_operations: Mutex, next_control_operation_id: AtomicU64, @@ -200,8 +209,13 @@ impl WebProcessRuntime { .saturating_sub(limits.websocket_http_connection_reserve); let lane_poll_limit = limits.max_http_handlers / 2; let lane_aux_poll_limit = (lane_poll_limit / 2).max(1); + let runtime_instance: Arc = + Arc::from(format!("{:032x}", rand::random::())); let runtime = Arc::new(Self { - runtime_instance: Arc::from(format!("{:032x}", rand::random::())), + operator_lifecycle: operator_lifecycle::OperatorLifecycle::new(Arc::clone( + &runtime_instance, + )), + runtime_instance, active_runtime, trace, http_connections: Arc::new(Semaphore::new(limits.max_http_connections)), diff --git a/src/web/manager/admission.rs b/src/web/manager/admission.rs index 3ad3988..18ae6a6 100644 --- a/src/web/manager/admission.rs +++ b/src/web/manager/admission.rs @@ -13,11 +13,21 @@ impl WebProcessRuntime { max_streams: usize, client_ip: IpAddr, public_addr: SocketAddr, - ) -> Option { + ) -> Result { + let _operator_admission = match self.try_operator_admission() { + Ok(admission) => admission, + Err(error) => { + self.streams_rejected.fetch_add(1, Ordering::Relaxed); + return Err(error); + } + }; let now = Instant::now(); let mut state = self.stream_admission.lock(); - if state.closed - || state.streams_live >= self.limits.max_streams_global + if state.closed { + self.streams_rejected.fetch_add(1, Ordering::Relaxed); + return Err(super::ManagerError::Closed); + } + if state.streams_live >= self.limits.max_streams_global || state .streams_per_profile .get(&profile_key) @@ -33,17 +43,17 @@ impl WebProcessRuntime { { self.streams_rejected.fetch_add(1, Ordering::Relaxed); self.limit_hits.fetch_add(1, Ordering::Relaxed); - return None; + return Err(super::ManagerError::Limit); } let Some(peer_port) = allocate_stream_port(&mut state, client_ip, public_addr) else { self.streams_rejected.fetch_add(1, Ordering::Relaxed); self.limit_hits.fetch_add(1, Ordering::Relaxed); - return None; + return Err(super::ManagerError::Limit); }; state.streams_live += 1; *state.streams_per_profile.entry(profile_key).or_insert(0) += 1; self.streams_opened.fetch_add(1, Ordering::Relaxed); - Some(peer_port) + Ok(peer_port) } /// Releases one live logical-stream slot after its relay task exits. @@ -60,6 +70,8 @@ impl WebProcessRuntime { } state.streams_live = state.streams_live.saturating_sub(1); decrement_map(&mut state.streams_per_profile, &profile_key); + drop(state); + self.notify_operator_work_changed(); } /// Records a logical stream rejected outside manager quota acquisition. diff --git a/src/web/manager/credentials.rs b/src/web/manager/credentials.rs index ecc3482..5e081df 100644 --- a/src/web/manager/credentials.rs +++ b/src/web/manager/credentials.rs @@ -66,6 +66,7 @@ impl WebProcessRuntime { if !config.web.enabled || !generation.proxy_shared.is_user_enabled(&profile.user) { return Err(ManagerError::Closed); } + let _operator_admission = self.try_operator_admission()?; let now = Instant::now(); let mut state = self.state.lock(); remove_expired_locked(&mut state, now); diff --git a/src/web/manager/lifecycle.rs b/src/web/manager/lifecycle.rs index 4985495..3853393 100644 --- a/src/web/manager/lifecycle.rs +++ b/src/web/manager/lifecycle.rs @@ -80,11 +80,14 @@ impl WebProcessRuntime { remove_bootstrap_locked(&mut state, bootstrap_hash); } self.sessions_closed.fetch_add(1, Ordering::Relaxed); + drop(state); + self.notify_operator_work_changed(); } /// Closes every WEB authority gate before any graceful wait begins. pub(crate) fn begin_shutdown(self: &std::sync::Arc) -> WebShutdownDrain { let started = TokioInstant::now(); + self.close_operator_lifecycle(); self.shutdown.cancel(); self.close_control_submission_gate(); self.close_websockets(); diff --git a/src/web/manager/operator_lifecycle.rs b/src/web/manager/operator_lifecycle.rs new file mode 100644 index 0000000..672bd73 --- /dev/null +++ b/src/web/manager/operator_lifecycle.rs @@ -0,0 +1,550 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use arc_swap::ArcSwap; +use parking_lot::Mutex; +use tokio::sync::{Mutex as AsyncMutex, Notify}; +use tokio::time::Instant as TokioInstant; +use tokio_util::sync::CancellationToken; + +use super::WebProcessRuntime; + +// Serialized state-machine values stay separate from synchronization mechanics. +mod status; +pub(crate) use status::{ + OperatorDrainOutcome, OperatorDrainState, OperatorDrainStatus, OperatorLifecycleState, + OperatorLifecycleStatus, +}; + +const OPERATOR_ADMISSION_CLOSED: usize = 1 << (usize::BITS - 1); +const OPERATOR_REGISTRATION_COUNT: usize = OPERATOR_ADMISSION_CLOSED - 1; +const DRAIN_REF_VERSION: &str = "wd1"; + +/// Stable operator-control rejection category. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum OperatorLifecycleError { + /// Process shutdown terminally closed the command gate. + Closed, + /// One drain already owns the single active operation slot. + OperationInProgress, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +struct WorkCounts { + sessions: usize, + streams: usize, + websockets: usize, +} + +impl WorkCounts { + fn is_zero(self) -> bool { + self.sessions == 0 && self.streams == 0 && self.websockets == 0 + } +} + +struct OperatorAdmission { + state: AtomicUsize, + registrations_drained: Notify, +} + +pub(super) struct OperatorRegistration<'a> { + admission: &'a OperatorAdmission, +} + +impl OperatorAdmission { + fn new() -> Self { + Self { + state: AtomicUsize::new(0), + registrations_drained: Notify::new(), + } + } + + fn try_register(&self) -> Option> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state & OPERATOR_ADMISSION_CLOSED != 0 + || state & OPERATOR_REGISTRATION_COUNT == OPERATOR_REGISTRATION_COUNT + { + return None; + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return Some(OperatorRegistration { admission: self }), + Err(observed) => state = observed, + } + } + } + + fn close(&self) { + self.state + .fetch_or(OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel); + } + + fn reopen(&self) { + self.state + .fetch_and(!OPERATOR_ADMISSION_CLOSED, Ordering::AcqRel); + } + + fn is_closed(&self) -> bool { + self.state.load(Ordering::Acquire) & OPERATOR_ADMISSION_CLOSED != 0 + } + + async fn wait_for_registrations(&self) { + loop { + let notified = self.registrations_drained.notified(); + if self.state.load(Ordering::Acquire) & OPERATOR_REGISTRATION_COUNT == 0 { + return; + } + notified.await; + } + } +} + +impl Drop for OperatorRegistration<'_> { + fn drop(&mut self) { + let previous = self.admission.state.fetch_sub(1, Ordering::AcqRel); + if previous & OPERATOR_REGISTRATION_COUNT == 1 { + self.admission.registrations_drained.notify_waiters(); + } + } +} + +struct ActiveDrain { + sequence: u64, + cancellation: CancellationToken, +} + +struct OperatorLifecycleInner { + state: OperatorLifecycleState, + epoch: u64, + since: Instant, + terminal: bool, + active: Option, + drain: Option, +} + +#[derive(Clone)] +struct OperatorSnapshot { + state: OperatorLifecycleState, + epoch: u64, + since: Instant, + terminal: bool, + drain: Option, +} + +pub(super) struct OperatorLifecycle { + runtime_instance: Arc, + admission: OperatorAdmission, + commands: AsyncMutex<()>, + inner: Mutex, + published: ArcSwap, + work_changed: Notify, + next_operation_id: AtomicU64, +} + +impl OperatorLifecycle { + pub(super) fn new(runtime_instance: Arc) -> Self { + let since = Instant::now(); + let snapshot = OperatorSnapshot { + state: OperatorLifecycleState::Running, + epoch: 0, + since, + terminal: false, + drain: None, + }; + Self { + runtime_instance, + admission: OperatorAdmission::new(), + commands: AsyncMutex::new(()), + inner: Mutex::new(OperatorLifecycleInner { + state: snapshot.state, + epoch: snapshot.epoch, + since, + terminal: false, + active: None, + drain: None, + }), + published: ArcSwap::from_pointee(snapshot), + work_changed: Notify::new(), + next_operation_id: AtomicU64::new(1), + } + } + + pub(super) fn try_register(&self) -> Option> { + self.admission.try_register() + } + + pub(super) fn notify_work_changed(&self) { + if self.admission.is_closed() { + self.work_changed.notify_waiters(); + } + } + + pub(super) fn status(&self, config_enabled: bool) -> OperatorLifecycleStatus { + let snapshot = self.published.load(); + let admission_open = !snapshot.terminal + && snapshot.state == OperatorLifecycleState::Running; + OperatorLifecycleStatus { + state: snapshot.state, + epoch: snapshot.epoch, + age_ms: millis(Instant::now().saturating_duration_since(snapshot.since)), + admission_open, + effective_new_work_admission: admission_open && config_enabled, + drain: snapshot.drain.clone(), + } + } + + fn is_terminal(&self) -> bool { + self.published.load().terminal + } + + fn publish_locked(&self, inner: &OperatorLifecycleInner) { + self.published.store(Arc::new(OperatorSnapshot { + state: inner.state, + epoch: inner.epoch, + since: inner.since, + terminal: inner.terminal, + drain: inner.drain.clone(), + })); + } + + fn transition_locked( + &self, + inner: &mut OperatorLifecycleInner, + state: OperatorLifecycleState, + ) { + if inner.state != state { + inner.state = state; + inner.epoch = inner.epoch.saturating_add(1); + inner.since = Instant::now(); + } + } + + fn update_counts(&self, sequence: u64, counts: WorkCounts) -> bool { + let mut inner = self.inner.lock(); + if inner + .active + .as_ref() + .is_none_or(|active| active.sequence != sequence) + { + return false; + } + let Some(drain) = inner.drain.as_mut() else { + return false; + }; + if drain.remaining_sessions != counts.sessions + || drain.remaining_streams != counts.streams + || drain.remaining_websockets != counts.websockets + { + drain.remaining_sessions = counts.sessions; + drain.remaining_streams = counts.streams; + drain.remaining_websockets = counts.websockets; + self.publish_locked(&inner); + } + true + } + + fn commit_force(&self, sequence: u64, counts: WorkCounts) -> bool { + let mut inner = self.inner.lock(); + if inner + .active + .as_ref() + .is_none_or(|active| active.sequence != sequence) + { + return false; + } + self.transition_locked(&mut inner, OperatorLifecycleState::ForceClosing); + let Some(drain) = inner.drain.as_mut() else { + return false; + }; + drain.state = OperatorDrainState::ForceClosing; + drain.remaining_sessions = counts.sessions; + drain.remaining_streams = counts.streams; + drain.remaining_websockets = counts.websockets; + drain.force_close_signalled = true; + self.publish_locked(&inner); + true + } + + fn complete(&self, sequence: u64, forced: bool) { + let mut inner = self.inner.lock(); + if inner + .active + .as_ref() + .is_none_or(|active| active.sequence != sequence) + { + return; + } + inner.active = None; + self.transition_locked(&mut inner, OperatorLifecycleState::Drained); + if let Some(drain) = inner.drain.as_mut() { + drain.state = OperatorDrainState::Completed; + drain.outcome = Some(if forced { + OperatorDrainOutcome::Forced + } else { + OperatorDrainOutcome::Graceful + }); + drain.completed_epoch_millis = Some(crate::web::trace::store_epoch_millis()); + drain.remaining_sessions = 0; + drain.remaining_streams = 0; + drain.remaining_websockets = 0; + } + self.publish_locked(&inner); + } + + fn close_terminal(&self) { + let mut inner = self.inner.lock(); + self.admission.close(); + if inner.terminal { + return; + } + inner.terminal = true; + if let Some(active) = inner.active.take() { + active.cancellation.cancel(); + if let Some(drain) = inner.drain.as_mut() { + drain.state = OperatorDrainState::Cancelled; + drain.outcome = Some(OperatorDrainOutcome::Cancelled); + drain.completed_epoch_millis = Some(crate::web::trace::store_epoch_millis()); + } + } + self.publish_locked(&inner); + self.work_changed.notify_waiters(); + } +} + +impl WebProcessRuntime { + /// Returns the current lock-free operator lifecycle snapshot. + pub(crate) fn operator_lifecycle_status(&self) -> OperatorLifecycleStatus { + let config_enabled = self.active_generation().config().web.enabled; + self.operator_lifecycle.status(config_enabled) + } + + /// Pauses new WEB work after every pre-cutover admission section completes. + pub(crate) async fn pause_operator( + self: &Arc, + ) -> Result { + let _command = self.operator_lifecycle.commands.lock().await; + { + let mut inner = self.operator_lifecycle.inner.lock(); + if inner.terminal || self.shutdown.is_cancelled() { + return Err(OperatorLifecycleError::Closed); + } + self.operator_lifecycle.admission.close(); + if matches!( + inner.state, + OperatorLifecycleState::Running | OperatorLifecycleState::Drained + ) { + self.operator_lifecycle + .transition_locked(&mut inner, OperatorLifecycleState::Paused); + } + self.operator_lifecycle.publish_locked(&inner); + } + self.operator_lifecycle + .admission + .wait_for_registrations() + .await; + if self.shutdown.is_cancelled() || self.operator_lifecycle.is_terminal() { + return Err(OperatorLifecycleError::Closed); + } + Ok(self.operator_lifecycle_status()) + } + + /// Starts one asynchronous graceful WEB drain under an absolute deadline. + pub(crate) async fn drain_operator( + self: &Arc, + timeout: Duration, + ) -> Result { + let _command = self.operator_lifecycle.commands.lock().await; + let sequence = self + .operator_lifecycle + .next_operation_id + .fetch_add(1, Ordering::Relaxed); + let cancellation = CancellationToken::new(); + let started = TokioInstant::now(); + let deadline = started + timeout; + let started_epoch_millis = crate::web::trace::store_epoch_millis(); + { + let mut inner = self.operator_lifecycle.inner.lock(); + if inner.terminal || self.shutdown.is_cancelled() { + return Err(OperatorLifecycleError::Closed); + } + if inner.active.is_some() + || matches!( + inner.state, + OperatorLifecycleState::Draining | OperatorLifecycleState::ForceClosing + ) + { + return Err(OperatorLifecycleError::OperationInProgress); + } + self.operator_lifecycle.admission.close(); + inner.active = Some(ActiveDrain { + sequence, + cancellation: cancellation.clone(), + }); + inner.drain = Some(OperatorDrainStatus { + operation_id: format!( + "{DRAIN_REF_VERSION}.{}.{sequence:016x}", + self.operator_lifecycle.runtime_instance + ), + state: OperatorDrainState::Draining, + outcome: None, + timeout_secs: timeout.as_secs(), + started_epoch_millis, + deadline_epoch_millis: started_epoch_millis + .saturating_add(timeout.as_millis().min(u128::from(u64::MAX)) as u64), + completed_epoch_millis: None, + remaining_sessions: 0, + remaining_streams: 0, + remaining_websockets: 0, + force_close_signalled: false, + }); + self.operator_lifecycle + .transition_locked(&mut inner, OperatorLifecycleState::Draining); + self.operator_lifecycle.publish_locked(&inner); + } + self.operator_lifecycle + .admission + .wait_for_registrations() + .await; + if self.shutdown.is_cancelled() || self.operator_lifecycle.is_terminal() { + return Err(OperatorLifecycleError::Closed); + } + let counts = self.operator_work_counts(); + if !self.operator_lifecycle.update_counts(sequence, counts) { + return Err(OperatorLifecycleError::Closed); + } + let runtime = Arc::clone(self); + self.spawn_auxiliary(async move { + runtime + .run_operator_drain(sequence, deadline, cancellation) + .await; + }); + Ok(self.operator_lifecycle_status()) + } + + /// Resumes operator admission and invalidates any active drain waiter. + pub(crate) async fn resume_operator( + self: &Arc, + ) -> Result { + let _command = self.operator_lifecycle.commands.lock().await; + let counts = self.operator_work_counts(); + let mut inner = self.operator_lifecycle.inner.lock(); + if inner.terminal || self.shutdown.is_cancelled() { + return Err(OperatorLifecycleError::Closed); + } + if let Some(active) = inner.active.take() { + active.cancellation.cancel(); + if let Some(drain) = inner.drain.as_mut() { + drain.state = OperatorDrainState::Cancelled; + drain.outcome = Some(OperatorDrainOutcome::Cancelled); + drain.completed_epoch_millis = Some(crate::web::trace::store_epoch_millis()); + drain.remaining_sessions = counts.sessions; + drain.remaining_streams = counts.streams; + drain.remaining_websockets = counts.websockets; + } + } + self.operator_lifecycle + .transition_locked(&mut inner, OperatorLifecycleState::Running); + self.operator_lifecycle.admission.reopen(); + self.operator_lifecycle.publish_locked(&inner); + drop(inner); + self.operator_lifecycle.work_changed.notify_waiters(); + Ok(self.operator_lifecycle_status()) + } + + pub(super) fn try_operator_admission( + &self, + ) -> Result, super::ManagerError> { + self.operator_lifecycle + .try_register() + .ok_or(super::ManagerError::AdmissionPaused) + } + + pub(super) fn notify_operator_work_changed(&self) { + self.operator_lifecycle.notify_work_changed(); + } + + pub(super) fn close_operator_lifecycle(&self) { + self.operator_lifecycle.close_terminal(); + } + + fn operator_work_counts(&self) -> WorkCounts { + let sessions = self.state.lock().sessions.len(); + let streams = self.stream_admission.lock().streams_live; + let websockets = self.websockets.lock().status().entries; + WorkCounts { + sessions, + streams, + websockets, + } + } + + async fn run_operator_drain( + self: Arc, + sequence: u64, + deadline: TokioInstant, + cancellation: CancellationToken, + ) { + let mut forced = false; + loop { + let notified = self.operator_lifecycle.work_changed.notified(); + let counts = self.operator_work_counts(); + if !self.operator_lifecycle.update_counts(sequence, counts) { + return; + } + if counts.is_zero() { + self.operator_lifecycle.complete(sequence, forced); + return; + } + if forced { + tokio::select! { + biased; + _ = cancellation.cancelled() => return, + _ = notified => {} + } + continue; + } + tokio::select! { + biased; + _ = cancellation.cancelled() => return, + _ = notified => {}, + _ = tokio::time::sleep_until(deadline) => { + let counts = self.operator_work_counts(); + if counts.is_zero() { + self.operator_lifecycle.complete(sequence, false); + return; + } + // Freeze the close set while admission is still closed. If resume + // wins the following epoch check, this snapshot is discarded. + let sessions = self + .state + .lock() + .sessions + .values() + .cloned() + .collect::>(); + if !self.operator_lifecycle.commit_force(sequence, counts) { + return; + } + forced = true; + for session in sessions { + session.close(); + } + } + } + } + } +} + +fn millis(duration: Duration) -> u64 { + duration.as_millis().min(u128::from(u64::MAX)) as u64 +} + +#[cfg(test)] +#[path = "operator_lifecycle/tests.rs"] +mod tests; diff --git a/src/web/manager/operator_lifecycle/status.rs b/src/web/manager/operator_lifecycle/status.rs new file mode 100644 index 0000000..dbc018c --- /dev/null +++ b/src/web/manager/operator_lifecycle/status.rs @@ -0,0 +1,90 @@ +use serde::Serialize; + +/// Reversible operator lifecycle state independent from terminal WEB shutdown. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum OperatorLifecycleState { + /// Operator admission is open subject to config and generation policy. + Running, + /// New WEB work is paused without closing existing work. + Paused, + /// Existing WEB work is completing before the absolute deadline. + Draining, + /// The deadline fired and close signals were sent to remaining sessions. + ForceClosing, + /// Every tracked WEB session, stream, and WebSocket completed. + Drained, +} + +/// Current or retained drain-operation phase. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum OperatorDrainState { + /// Existing work is inside its graceful completion interval. + Draining, + /// Forced session closure was signalled and zero is not confirmed yet. + ForceClosing, + /// The operation reached confirmed zero. + Completed, + /// Explicit resume or terminal process shutdown cancelled the waiter. + Cancelled, +} + +/// Terminal result retained for the latest drain operation. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum OperatorDrainOutcome { + /// All tracked work completed before forced closure was committed. + Graceful, + /// Tracked work reached zero after the deadline forced session closure. + Forced, + /// The drain waiter was cancelled before confirmed zero. + Cancelled, +} + +/// API-visible status for the active or latest process-local drain. +#[derive(Clone, Debug, Serialize)] +pub(crate) struct OperatorDrainStatus { + /// Opaque process-fenced drain identifier. + pub(crate) operation_id: String, + /// Current operation phase. + pub(crate) state: OperatorDrainState, + /// Terminal outcome when the operation no longer waits for zero. + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) outcome: Option, + /// Frozen relative deadline accepted from the API. + pub(crate) timeout_secs: u64, + /// Wall-clock projection retained only for operator correlation. + pub(crate) started_epoch_millis: u64, + /// Wall-clock projection of the monotonic deadline. + pub(crate) deadline_epoch_millis: u64, + /// Wall-clock completion or cancellation time. + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) completed_epoch_millis: Option, + /// Most recent exact live-session sample. + pub(crate) remaining_sessions: usize, + /// Most recent exact logical-stream ownership sample. + pub(crate) remaining_streams: usize, + /// Most recent exact session-owned WebSocket sample. + pub(crate) remaining_websockets: usize, + /// Whether the deadline won and committed the forced-close snapshot. + pub(crate) force_close_signalled: bool, +} + +/// API-visible snapshot of reversible WEB operator lifecycle state. +#[derive(Clone, Debug, Serialize)] +pub(crate) struct OperatorLifecycleStatus { + /// Stable process-local state machine value. + pub(crate) state: OperatorLifecycleState, + /// Monotonic state-transition epoch independent from config revision. + pub(crate) epoch: u64, + /// Monotonic age of the current state. + pub(crate) age_ms: u64, + /// Whether the operator-owned admission fence is open. + pub(crate) admission_open: bool, + /// Operator and effective config admission conjunction. + pub(crate) effective_new_work_admission: bool, + /// Active or latest drain retained until replacement or process restart. + #[serde(skip_serializing_if = "Option::is_none")] + pub(crate) drain: Option, +} diff --git a/src/web/manager/operator_lifecycle/tests.rs b/src/web/manager/operator_lifecycle/tests.rs new file mode 100644 index 0000000..6ce2ce8 --- /dev/null +++ b/src/web/manager/operator_lifecycle/tests.rs @@ -0,0 +1,220 @@ +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; +} diff --git a/src/web/manager/session_creation.rs b/src/web/manager/session_creation.rs index 4e068f9..e0474ab 100644 --- a/src/web/manager/session_creation.rs +++ b/src/web/manager/session_creation.rs @@ -189,6 +189,7 @@ impl WebProcessRuntime { ip_learning_eligible, carrier_deadline_at: entry.carrier_deadline_at.ok_or(ManagerError::Protocol)?, }; + let _operator_admission = self.try_operator_admission()?; state .bootstraps .get_mut(&bootstrap_hash) @@ -272,6 +273,7 @@ impl WebProcessRuntime { let Some(carrier) = candidates.first().copied() else { return Err(ManagerError::Protocol); }; + let _operator_admission = self.try_operator_admission()?; if !admit_initial(self, &mut state, now, client_ip, profile_key, &profile) { return Err(ManagerError::Limit); } diff --git a/src/web/manager/websocket.rs b/src/web/manager/websocket.rs index a8abed1..7a36073 100644 --- a/src/web/manager/websocket.rs +++ b/src/web/manager/websocket.rs @@ -186,6 +186,7 @@ impl Drop for WebSocketConnection { drop(self.slot.take()); self.entry.released.cancel(); runtime.websocket_notify.notify_waiters(); + runtime.notify_operator_work_changed(); } } } diff --git a/src/web/session/uplink.rs b/src/web/session/uplink.rs index 0e8cdc9..49da439 100644 --- a/src/web/session/uplink.rs +++ b/src/web/session/uplink.rs @@ -288,12 +288,14 @@ impl WebSession { return None; } let manager = self.manager.upgrade()?; - let peer_port = manager.try_acquire_stream( - self.profile_key, - self.profile.max_streams, - self.client_ip, - self.profile.public_addr, - )?; + let peer_port = manager + .try_acquire_stream( + self.profile_key, + self.profile.max_streams, + self.client_ip, + self.profile.public_addr, + ) + .ok()?; if state.active_peer_ports.insert(peer_port) { return Some(peer_port); } diff --git a/src/web/session/websocket.rs b/src/web/session/websocket.rs index 9f09adb..400d389 100644 --- a/src/web/session/websocket.rs +++ b/src/web/session/websocket.rs @@ -229,14 +229,12 @@ impl WebSession { let Some(manager) = self.manager.upgrade() else { return Err(ManagerError::Closed); }; - let Some(peer_port) = manager.try_acquire_stream( + let peer_port = manager.try_acquire_stream( self.profile_key, self.profile.max_streams, self.client_ip, self.profile.public_addr, - ) else { - return Err(ManagerError::Limit); - }; + )?; if !state.active_peer_ports.insert(peer_port) { manager.release_stream( self.profile_key, diff --git a/src/web/session/websocket/tests.rs b/src/web/session/websocket/tests.rs index c1d8b37..45f964b 100644 --- a/src/web/session/websocket/tests.rs +++ b/src/web/session/websocket/tests.rs @@ -153,7 +153,7 @@ async fn rejected_open_retains_stream_quota_until_lane_socket_teardown() { runtime.session.client_ip, runtime.session.profile.public_addr, ) - .is_none() + .is_err() ); runtime.session.close_websocket_lane(reservation); @@ -271,7 +271,7 @@ async fn closed_session_keeps_stream_owned_quota_until_task_completion() { runtime.session.client_ip, runtime.session.profile.public_addr, ) - .is_none() + .is_err() ); runtime.session.wait().await; let peer_port = runtime