From aa80391b45f206e7a7092a72e8c0e0447bb609c6 Mon Sep 17 00:00:00 2001 From: Nathan Flurry Date: Wed, 5 Aug 2026 18:37:05 -0700 Subject: [PATCH] refactor(native-sidecar): simplify VM lifecycle coordination --- crates/agentos-sidecar/src/acp/mod.rs | 167 +- .../src/execution/javascript/rpc.rs | 1 + crates/native-sidecar/src/extension.rs | 36 - .../native-sidecar/src/extension_services.rs | 152 +- crates/native-sidecar/src/lib.rs | 3 +- .../src/ownership_coordinator.rs | 2035 ++++++++--------- .../native-sidecar/src/request_operations.rs | 1066 +++++---- crates/native-sidecar/src/service.rs | 35 + crates/native-sidecar/src/state.rs | 8 + crates/native-sidecar/src/stdio.rs | 510 +++-- .../src/stdio/request_concurrency_tests.rs | 690 +++++- .../tests/architecture_guards.rs | 77 +- crates/runtime/src/lib.rs | 9 + request-concurrency-fix-prompt.md | 179 +- vm-lifecycle-gate-simplification-spec.md | 639 ++++++ 15 files changed, 3489 insertions(+), 2118 deletions(-) create mode 100644 vm-lifecycle-gate-simplification-spec.md diff --git a/crates/agentos-sidecar/src/acp/mod.rs b/crates/agentos-sidecar/src/acp/mod.rs index 3d3c343ecc..a31589b92d 100644 --- a/crates/agentos-sidecar/src/acp/mod.rs +++ b/crates/agentos-sidecar/src/acp/mod.rs @@ -16,8 +16,8 @@ use agentos_native_sidecar::wire::{ StreamChannel, WriteStdinRequest, }; use agentos_native_sidecar::{ - Extension, ExtensionContext, ExtensionFuture, ExtensionOrderingPolicy, ExtensionRequestClass, - ExtensionResponse, SidecarError, + Extension, ExtensionContext, ExtensionFuture, ExtensionRequestClass, ExtensionResponse, + SidecarError, }; use agentos_protocol::generated::v1::*; use agentos_protocol::ACP_EXTENSION_NAMESPACE; @@ -1351,22 +1351,6 @@ impl Extension for AcpExtension { ACP_EXTENSION_NAMESPACE } - fn request_ordering_key(&self, ownership: &OwnershipScope, payload: &[u8]) -> Option> { - let request = decode_request(payload).ok()?; - let session_id = durable_request_session_id(&request)?; - Some(durable_route_key(ownership, session_id).into_bytes()) - } - - fn request_ordering_policy( - &self, - _ownership: &OwnershipScope, - _payload: &[u8], - ) -> ExtensionOrderingPolicy { - // AcpRouteEntry performs the same bounded exclusion and must return - // ACP's established typed `session_busy` response. - ExtensionOrderingPolicy::ExtensionManaged - } - fn handle_request<'a>( &'a self, ctx: ExtensionContext, @@ -1503,57 +1487,6 @@ fn durable_route_key(ownership: &OwnershipScope, session_id: &str) -> String { key } -/// Return the public durable session identity targeted by a native ACP -/// request. Global queries and legacy browser-only messages have no durable -/// session conflict domain. Validation remains in the request handler; this -/// ingress classification hook is deliberately total and side-effect free. -fn durable_request_session_id(request: &AcpRequest) -> Option<&str> { - match request { - AcpRequest::AcpOpenSessionRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpGetDurableSessionRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpDeleteSessionRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpUnloadSessionRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpPromptRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpCancelPromptRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpRespondPermissionRequest(request) => Some(&request.session_id), - AcpRequest::AcpReadHistoryRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpGetSessionConfigRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpSetSessionConfigOptionRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpGetSessionCapabilitiesRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpGetSessionAgentInfoRequest(request) => { - Some(request.session_id.as_deref().unwrap_or("main")) - } - AcpRequest::AcpListDurableSessionsRequest(_) - | AcpRequest::AcpCreateSessionRequest(_) - | AcpRequest::AcpSessionRequest(_) - | AcpRequest::AcpGetSessionStateRequest(_) - | AcpRequest::AcpCloseSessionRequest(_) - | AcpRequest::AcpResumeSessionRequest(_) - | AcpRequest::AcpDeliverAgentOutputRequest(_) - | AcpRequest::AcpListAgentsRequest(_) => None, - } -} - fn session_store_error(error: agentos_native_sidecar::vm_sqlite::VmSqliteError) -> SidecarError { match error { error @ (agentos_native_sidecar::vm_sqlite::VmSqliteError::ResultTooLarge { .. } @@ -1770,6 +1703,7 @@ fn error_response(error: SidecarError) -> AcpResponse { fn error_code(error: &SidecarError) -> String { let code = match error { SidecarError::ResourceLimit(_) => "resource_limit", + SidecarError::RequestAdmission { code, .. } => return (*code).to_owned(), SidecarError::InvalidState(message) => message .split_once(':') .map(|(prefix, _)| prefix) @@ -1837,91 +1771,6 @@ mod tests { ); } - #[test] - fn durable_acp_requests_expose_an_opaque_owned_session_ordering_key() { - use agentos_native_sidecar::wire::VmOwnership; - - let extension = AcpExtension::new(); - let ownership = OwnershipScope::VmOwnership(VmOwnership { - connection_id: String::from("conn-1"), - session_id: String::from("owner-session"), - vm_id: String::from("vm-1"), - }); - let other_vm = OwnershipScope::VmOwnership(VmOwnership { - connection_id: String::from("conn-1"), - session_id: String::from("owner-session"), - vm_id: String::from("vm-2"), - }); - let default_prompt = serde_bare::to_vec(&AcpRequest::AcpPromptRequest(AcpPromptRequest { - session_id: None, - idempotency_key: None, - content: String::from("[]"), - })) - .expect("encode default prompt"); - let explicit_main_cancel = serde_bare::to_vec(&AcpRequest::AcpCancelPromptRequest( - AcpCancelPromptRequest { - session_id: Some(String::from("main")), - }, - )) - .expect("encode explicit main cancellation"); - let other_prompt = serde_bare::to_vec(&AcpRequest::AcpPromptRequest(AcpPromptRequest { - session_id: Some(String::from("other")), - idempotency_key: None, - content: String::from("[]"), - })) - .expect("encode other prompt"); - - let main_key = durable_route_key(&ownership, "main").into_bytes(); - assert_eq!( - extension.request_ordering_key(&ownership, &default_prompt), - Some(main_key.clone()), - "omitted durable session IDs must order as main" - ); - assert_eq!( - extension.request_ordering_key(&ownership, &explicit_main_cancel), - Some(main_key), - "progress and ordinary messages for one durable session must expose the same target key" - ); - assert_ne!( - extension.request_ordering_key(&ownership, &default_prompt), - extension.request_ordering_key(&ownership, &other_prompt), - "different durable ACP sessions must remain independent" - ); - assert_ne!( - extension.request_ordering_key(&ownership, &default_prompt), - extension.request_ordering_key(&other_vm, &default_prompt), - "identically named durable sessions in different VMs must remain independent" - ); - assert_eq!( - extension.request_ordering_policy(&ownership, &default_prompt), - ExtensionOrderingPolicy::ExtensionManaged, - "ACP must retain its protocol-specific session_busy rejection" - ); - } - - #[test] - fn global_and_invalid_acp_requests_have_no_session_ordering_key() { - use agentos_native_sidecar::wire::ConnectionOwnership; - - let extension = AcpExtension::new(); - let ownership = OwnershipScope::ConnectionOwnership(ConnectionOwnership { - connection_id: String::from("conn-1"), - }); - let list = serde_bare::to_vec(&AcpRequest::AcpListDurableSessionsRequest( - AcpListDurableSessionsRequest { - cursor: None, - limit: None, - }, - )) - .expect("encode list request"); - - assert_eq!(extension.request_ordering_key(&ownership, &list), None); - assert_eq!( - extension.request_ordering_key(&ownership, b"not an ACP request"), - None - ); - } - #[test] fn route_start_is_single_flight_and_releases_waiters_without_sleeping() { let route = Arc::new(AcpRouteEntry::new(String::from("conn-1"))); @@ -3058,6 +2907,16 @@ mod tests { )), "sqlite_result_limit" ); + assert_eq!( + error_code(&SidecarError::RequestAdmission { + code: "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT", + message: String::from("VM lifecycle work is pending"), + configuration_path: None, + retryable: true, + errno: "EBUSY", + }), + "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT" + ); } #[test] diff --git a/crates/native-sidecar/src/execution/javascript/rpc.rs b/crates/native-sidecar/src/execution/javascript/rpc.rs index e0aaac0ff3..b3397126db 100644 --- a/crates/native-sidecar/src/execution/javascript/rpc.rs +++ b/crates/native-sidecar/src/execution/javascript/rpc.rs @@ -4979,6 +4979,7 @@ fn format_unix_socket_resource( pub(crate) fn error_code(error: &SidecarError) -> &'static str { match error { SidecarError::ResourceLimit(_) => "ERR_AGENTOS_RESOURCE_LIMIT", + SidecarError::RequestAdmission { code, .. } => code, SidecarError::InvalidState(_) => "invalid_state", SidecarError::ProtocolVersionMismatch(_) => "protocol_version_mismatch", SidecarError::BridgeVersionMismatch(_) => "bridge_version_mismatch", diff --git a/crates/native-sidecar/src/extension.rs b/crates/native-sidecar/src/extension.rs index eff93b8ec7..f726745079 100644 --- a/crates/native-sidecar/src/extension.rs +++ b/crates/native-sidecar/src/extension.rs @@ -837,45 +837,9 @@ pub enum ExtensionRequestClass { Progress, } -/// Who enforces exclusivity for an extension-provided opaque ordering key. -/// -/// Core-exclusive keys receive a bounded, typed conflict rejection before the -/// second request starts. An extension may retain enforcement when it needs to -/// preserve a richer protocol-specific rejection (ACP's `session_busy` is the -/// motivating case); core still scopes and validates the opaque key. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum ExtensionOrderingPolicy { - CoreExclusive, - ExtensionManaged, -} - pub trait Extension: Send + Sync { fn namespace(&self) -> &str; - /// Return the extension-owned conflict key for an ordinary request. - /// - /// Core treats this value as opaque and combines it with the extension - /// namespace and connection ownership when it builds request-operation - /// metadata. Extensions whose identities are scoped more narrowly than a - /// connection (for example, to one VM) must include that ownership in the - /// returned key. `None` means the request has no extension-specific - /// ordering boundary. - fn request_ordering_key( - &self, - _ownership: &OwnershipScope, - _payload: &[u8], - ) -> Option> { - None - } - - fn request_ordering_policy( - &self, - _ownership: &OwnershipScope, - _payload: &[u8], - ) -> ExtensionOrderingPolicy { - ExtensionOrderingPolicy::CoreExclusive - } - fn request_class(&self, _payload: &[u8]) -> ExtensionRequestClass { ExtensionRequestClass::Ordinary } diff --git a/crates/native-sidecar/src/extension_services.rs b/crates/native-sidecar/src/extension_services.rs index d9742b02df..fe8b8fefe9 100644 --- a/crates/native-sidecar/src/extension_services.rs +++ b/crates/native-sidecar/src/extension_services.rs @@ -29,7 +29,7 @@ use crate::protocol::{ }; use crate::request_operations::{ OperationCancellation, OperationCancellationReason, RequestOperationMetadata, - RequestOrderingKey, + VmConcurrencyClass, }; use crate::service::NativeSidecar; use crate::state::SidecarError; @@ -325,7 +325,7 @@ fn with_vm_admission_class( ownership: &OwnershipScope, class: ExtensionServiceAdmissionClass, ) -> PreparedExtensionServiceCommand { - let OwnershipScope::VmOwnership(scope) = ownership else { + let OwnershipScope::VmOwnership(_) = ownership else { return prepared; }; prepared.admission = Some(ExtensionServiceAdmission { @@ -333,11 +333,7 @@ fn with_vm_admission_class( metadata: RequestOperationMetadata::new( ownership.clone(), prepared.operation, - RequestOrderingKey::VmOperation { - connection_id: scope.connection_id.clone(), - session_id: scope.session_id.clone(), - vm_id: scope.vm_id.clone(), - }, + VmConcurrencyClass::SharedVm, ), cancellation: OperationCancellation::new(), pre_admitted: None, @@ -846,6 +842,7 @@ pub(crate) enum ExtensionServiceCommand { #[derive(Clone)] pub(crate) struct RoutedExtensionServices { commands: mpsc::Sender, + progress_commands: mpsc::Sender, /// Wakes extension-side probes only after the protocol coordinator has /// drained the runtime producer queues. Runtime producers use a separate /// notification owned exclusively by the central process-event pump, so @@ -925,11 +922,13 @@ impl CompletedProcessEventStore { } impl RoutedExtensionServices { + #[cfg(test)] pub(crate) fn new( commands: mpsc::Sender, routed_process_event_notify: Arc, ) -> Self { Self { + progress_commands: commands.clone(), commands, routed_process_event_notify, process_event_broker: None, @@ -937,8 +936,9 @@ impl RoutedExtensionServices { } } - pub(crate) fn new_with_process_event_broker( + pub(crate) fn new_with_process_event_broker_and_progress( commands: mpsc::Sender, + progress_commands: mpsc::Sender, routed_process_event_notify: Arc, process_event_broker: ProcessEventBroker, ) -> Self { @@ -946,6 +946,7 @@ impl RoutedExtensionServices { CompletedProcessEventStore::new(process_event_broker.max_events()); Self { commands, + progress_commands, routed_process_event_notify, process_event_broker: Some(process_event_broker), completed_process_events, @@ -957,7 +958,42 @@ impl RoutedExtensionServices { T: Send + 'static, F: FnOnce(Reply) -> ExtensionServiceCommand + Send + 'static, { - let commands = self.commands.clone(); + self.call_on(self.commands.clone(), build) + } + + fn call_progress(&self, build: F) -> ExtensionFuture<'static, T> + where + T: Send + 'static, + F: FnOnce(Reply) -> ExtensionServiceCommand + Send + 'static, + { + let commands = self.progress_commands.clone(); + Box::pin(async move { + let (reply, response) = oneshot::channel(); + commands.try_send(build(reply)).map_err(|error| match error { + mpsc::error::TrySendError::Full(_) => SidecarError::InvalidState(String::from( + "ERR_AGENTOS_PROGRESS_SERVICE_LIMIT: reserved progress service admission is full; raise runtime.protocol.maxProgressFrames", + )), + mpsc::error::TrySendError::Closed(_) => SidecarError::Io(String::from( + "ERR_AGENTOS_EXTENSION_SERVICE_CLOSED: sidecar coordinator stopped; active extension operation was cancelled", + )), + })?; + response.await.map_err(|_| { + SidecarError::Io(String::from( + "ERR_AGENTOS_EXTENSION_SERVICE_REPLY_CLOSED: sidecar coordinator dropped an extension operation without a result", + )) + })? + }) + } + + fn call_on( + &self, + commands: mpsc::Sender, + build: F, + ) -> ExtensionFuture<'static, T> + where + T: Send + 'static, + F: FnOnce(Reply) -> ExtensionServiceCommand + Send + 'static, + { Box::pin(async move { let (reply, response) = oneshot::channel(); commands.send(build(reply)).await.map_err(|_| { @@ -1017,7 +1053,7 @@ impl ExtensionServices for RoutedExtensionServices { ownership: OwnershipScope, request: WriteStdinRequest, ) -> ExtensionFuture<'static, StdinWrittenResponse> { - self.call(move |reply| ExtensionServiceCommand::WriteStdin { + self.call_progress(move |reply| ExtensionServiceCommand::WriteStdin { ownership, request, reply, @@ -1029,7 +1065,7 @@ impl ExtensionServices for RoutedExtensionServices { ownership: OwnershipScope, request: CloseStdinRequest, ) -> ExtensionFuture<'static, StdinClosedResponse> { - self.call(move |reply| ExtensionServiceCommand::CloseStdin { + self.call_progress(move |reply| ExtensionServiceCommand::CloseStdin { ownership, request, reply, @@ -1041,7 +1077,7 @@ impl ExtensionServices for RoutedExtensionServices { ownership: OwnershipScope, request: KillProcessRequest, ) -> ExtensionFuture<'static, ProcessKilledResponse> { - self.call(move |reply| ExtensionServiceCommand::KillProcess { + self.call_progress(move |reply| ExtensionServiceCommand::KillProcess { ownership, request, reply, @@ -1329,6 +1365,98 @@ mod internal_event_lifecycle_tests { } } + #[tokio::test(flavor = "current_thread")] + async fn write_stdin_uses_the_reserved_progress_service_lane() { + let (ordinary_tx, mut ordinary_rx) = mpsc::channel(1); + let (progress_tx, mut progress_rx) = mpsc::channel(1); + let services = RoutedExtensionServices { + commands: ordinary_tx, + progress_commands: progress_tx, + routed_process_event_notify: Arc::new(Notify::new()), + process_event_broker: None, + completed_process_events: CompletedProcessEventStore::new(1), + }; + let ownership = OwnershipScope::vm("connection-a", "session-a", "vm-a"); + let write = services.write_stdin( + ownership.clone(), + WriteStdinRequest { + process_id: String::from("process-a"), + chunk: b"cancel\n".to_vec(), + }, + ); + let service = async { + assert!(matches!( + ordinary_rx.try_recv(), + Err(mpsc::error::TryRecvError::Empty) + )); + let command = progress_rx + .recv() + .await + .expect("WriteStdin reaches the progress service lane"); + let ExtensionServiceCommand::WriteStdin { + ownership: actual_ownership, + request, + reply, + } = command + else { + panic!("progress lane received a different service command"); + }; + assert_eq!(actual_ownership, ownership); + assert_eq!(request.process_id, "process-a"); + assert_eq!(request.chunk, b"cancel\n"); + reply + .send(Ok(StdinWrittenResponse { + process_id: request.process_id, + accepted_bytes: request.chunk.len() as u64, + })) + .expect("WriteStdin caller retains its reply waiter"); + }; + let (response, ()) = tokio::join!(write, service); + let response = response.expect("WriteStdin response"); + assert_eq!(response.accepted_bytes, 7); + } + + #[tokio::test(flavor = "current_thread")] + async fn saturated_progress_service_lane_returns_a_typed_limit() { + let (ordinary_tx, _ordinary_rx) = mpsc::channel(1); + let (progress_tx, _progress_rx) = mpsc::channel(1); + let services = RoutedExtensionServices { + commands: ordinary_tx, + progress_commands: progress_tx, + routed_process_event_notify: Arc::new(Notify::new()), + process_event_broker: None, + completed_process_events: CompletedProcessEventStore::new(1), + }; + let ownership = OwnershipScope::vm("connection-a", "session-a", "vm-a"); + let mut first = services.write_stdin( + ownership.clone(), + WriteStdinRequest { + process_id: String::from("process-a"), + chunk: vec![1], + }, + ); + let waker = std::task::Waker::noop(); + let mut context = Context::from_waker(waker); + assert!(matches!(first.as_mut().poll(&mut context), Poll::Pending)); + + let error = services + .write_stdin( + ownership, + WriteStdinRequest { + process_id: String::from("process-a"), + chunk: vec![2], + }, + ) + .await + .expect_err("full progress lane rejects without waiting"); + assert!(error + .to_string() + .contains("ERR_AGENTOS_PROGRESS_SERVICE_LIMIT")); + assert!(error + .to_string() + .contains("runtime.protocol.maxProgressFrames")); + } + fn prepared_internal_event( coordinator: &OwnershipCoordinator, ownership: &OwnershipScope, diff --git a/crates/native-sidecar/src/lib.rs b/crates/native-sidecar/src/lib.rs index 86a570c507..2f57297e1c 100644 --- a/crates/native-sidecar/src/lib.rs +++ b/crates/native-sidecar/src/lib.rs @@ -29,8 +29,7 @@ pub mod vm_sqlite; pub use agentos_sidecar_protocol::{generated_protocol, protocol, wire}; pub use extension::{ - Extension, ExtensionContext, ExtensionFuture, ExtensionOrderingPolicy, ExtensionRequestClass, - ExtensionResponse, + Extension, ExtensionContext, ExtensionFuture, ExtensionRequestClass, ExtensionResponse, }; pub use service::{DispatchResult, NativeSidecar, NativeSidecarConfig, SidecarError}; pub use state::EventSinkTransport; diff --git a/crates/native-sidecar/src/ownership_coordinator.rs b/crates/native-sidecar/src/ownership_coordinator.rs index 55797bcb73..3d21a7b043 100644 --- a/crates/native-sidecar/src/ownership_coordinator.rs +++ b/crates/native-sidecar/src/ownership_coordinator.rs @@ -6,10 +6,10 @@ //! mutex. A request keeps the returned permit while it runs; every actual //! access to mutable VM state remains a separate short service command. -use crate::extension::ExtensionOrderingPolicy; use crate::request_operations::{ - OperationCancellation, OperationCancellationReason, RequestOperationMetadata, - RequestOrderingKey, IN_FLIGHT_REQUEST_COUNT_PATH, + OperationCancellation, OperationCancellationReason, OperationTable, RequestOperationKey, + RequestOperationMetadata, ScopeOperationDrain, VmConcurrencyClass, + IN_FLIGHT_REQUEST_COUNT_PATH, }; use crate::wire::OwnershipScope; use agentos_runtime::RuntimeConfig; @@ -19,7 +19,7 @@ use std::sync::{Arc, Mutex, MutexGuard, Weak}; use tokio::sync::Notify; const CONNECTION_LIMIT_PATH: &str = "runtime.resources.maxConnections"; -const SESSION_LIMIT_PATH: &str = "runtime.protocol.maxInFlightRequests"; +const SESSION_LIMIT_PATH: &str = "runtime.protocol.maxSessionsPerConnection"; const VM_LIMIT_PATH: &str = "runtime.fairness.maxVms"; pub(crate) const INTERNAL_EVENT_OPERATION_LIMIT_PATH: &str = "runtime.protocol.maxProcessEvents"; @@ -39,11 +39,7 @@ impl OwnershipCoordinatorLimits { pub(crate) fn from_runtime_config(config: &RuntimeConfig) -> Self { Self { max_connections: config.resources.max_connections, - // There is not a second request queue hidden in this coordinator. - // Retained live session handles reuse the protocol's bounded - // request-count ceiling until a dedicated session-membership limit - // is introduced. - max_sessions_per_connection: config.protocol.max_in_flight_requests, + max_sessions_per_connection: config.protocol.max_sessions_per_connection, max_vms_per_session: config.fairness.max_vms, max_operations_per_entity: config.protocol.max_in_flight_requests, max_internal_event_operations_per_entity: config.protocol.max_process_events, @@ -78,6 +74,7 @@ impl OwnershipCoordinatorLimits { pub(crate) enum CoordinatorPhase { Open, Closing, + #[cfg(test)] Closed, } @@ -88,6 +85,7 @@ pub(crate) enum VmLifecyclePhase { Active, } +#[cfg(test)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct EntityCoordinatorSnapshot { pub(crate) phase: CoordinatorPhase, @@ -95,6 +93,7 @@ pub(crate) struct EntityCoordinatorSnapshot { pub(crate) child_count: usize, } +#[cfg(test)] #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct VmCoordinatorSnapshot { pub(crate) phase: CoordinatorPhase, @@ -130,9 +129,6 @@ pub(crate) enum OwnershipCoordinatorError { vm: String, lifecycle: VmLifecyclePhase, }, - OrderingConflict { - scope: String, - }, Cancelled { reason: OperationCancellationReason, }, @@ -151,11 +147,35 @@ impl OwnershipCoordinatorError { Self::OwnershipMismatch { .. } => "ERR_AGENTOS_COORDINATOR_OWNERSHIP", Self::Closing { .. } => "ERR_AGENTOS_COORDINATOR_CLOSING", Self::LifecycleConflict { .. } => "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT", - Self::OrderingConflict { .. } => "ERR_AGENTOS_ORDERING_CONFLICT", Self::Cancelled { .. } => "ERR_AGENTOS_COORDINATOR_CANCELLED", Self::NotDrained { .. } => "ERR_AGENTOS_COORDINATOR_NOT_DRAINED", } } + + pub(crate) fn configuration_path(&self) -> Option<&'static str> { + match self { + Self::Limit { + configuration_path, .. + } => Some(configuration_path), + _ => None, + } + } + + pub(crate) fn retryable(&self) -> bool { + matches!(self, Self::Limit { .. } | Self::LifecycleConflict { .. }) + } + + pub(crate) fn errno(&self) -> &'static str { + match self { + Self::Limit { .. } | Self::LifecycleConflict { .. } => "EAGAIN", + Self::Duplicate { .. } => "EEXIST", + Self::NotFound { .. } => "ENOENT", + Self::OwnershipMismatch { .. } => "EACCES", + Self::Closing { .. } => "ESHUTDOWN", + Self::Cancelled { .. } => "ECANCELED", + Self::NotDrained { .. } => "EBUSY", + } + } } impl fmt::Display for OwnershipCoordinatorError { @@ -192,11 +212,6 @@ impl fmt::Display for OwnershipCoordinatorError { "{}: VM {vm} lifecycle is {lifecycle:?}", self.code() ), - Self::OrderingConflict { scope } => write!( - formatter, - "{}: {scope} already has an active operation", - self.code() - ), Self::Cancelled { reason } => write!( formatter, "{}: coordinator admission was cancelled ({reason:?})", @@ -227,76 +242,103 @@ struct OwnershipCoordinatorInner { state: Mutex, } -#[derive(Debug, Default)] +#[derive(Debug)] struct OwnershipCoordinatorState { - connections: BTreeMap, + next_generation: u64, + connections: BTreeMap, } -#[derive(Clone, Debug)] -pub(crate) struct ConnectionCoordinator { - root: Weak, - inner: Arc, +impl Default for OwnershipCoordinatorState { + fn default() -> Self { + Self { + next_generation: 1, + connections: BTreeMap::new(), + } + } } #[derive(Debug)] -struct ConnectionCoordinatorInner { - connection_id: String, - limits: OwnershipCoordinatorLimits, - state: Mutex, - drained: Notify, +struct ConnectionRecord { + generation: u64, + phase: CoordinatorPhase, + sessions: BTreeMap, } #[derive(Debug)] -struct ConnectionCoordinatorState { +struct SessionRecord { + generation: u64, phase: CoordinatorPhase, - sessions: BTreeMap, - operations: BTreeMap, - extension_ordering: BTreeMap<(String, Vec), ()>, - next_operation_id: u64, + vms: BTreeMap, } -#[derive(Clone, Debug)] -pub(crate) struct SessionCoordinator { - parent: Weak, - inner: Arc, +#[derive(Debug)] +struct VmRecord { + generation: u64, + phase: CoordinatorPhase, + gate: Arc, } -#[derive(Debug)] -struct SessionCoordinatorInner { +#[derive(Clone, Debug)] +pub(crate) struct ConnectionCoordinator { + root: Weak, connection_id: String, - session_id: String, - limits: OwnershipCoordinatorLimits, - state: Mutex, - drained: Notify, + generation: u64, } -#[derive(Debug)] -struct SessionCoordinatorState { - phase: CoordinatorPhase, - vms: BTreeMap, - operations: BTreeMap, - next_operation_id: u64, +#[derive(Clone, Debug)] +pub(crate) struct SessionCoordinator { + root: Weak, + connection_id: String, + connection_generation: u64, + session_id: String, + generation: u64, } #[derive(Clone, Debug)] pub(crate) struct VmCoordinator { - parent: Weak, - inner: Arc, -} - + root: Weak, + connection_id: String, + connection_generation: u64, + session_id: String, + session_generation: u64, + vm_id: String, + generation: u64, + gate: Arc, +} + +/// Explicit per-VM lifecycle admission gate. +/// +/// The state machine is `Idle -> Pending -> Active -> Idle`. The transition to +/// `Pending` is the linearization point that closes ordinary VM admission; +/// operations registered before it drain, while later ordinary requests get +/// `ERR_AGENTOS_VM_LIFECYCLE_CONFLICT`. `Pending -> Active` occurs only after +/// every ordinary and internal permit has drained. Progress-critical internal +/// settlement events have separate bounded admission during `Pending`; during +/// `Active` they remain durably claimed but cannot mutate the VM. +/// +/// This is not a standard-library or Tokio `RwLock`. A standard lock cannot +/// cross `.await` without blocking a runtime worker. A Tokio `RwLock` would put +/// new ordinary work into an implicit waiter queue after lifecycle admission, +/// while this protocol must reject that work immediately. Neither lock models +/// the distinct bounded internal-settlement class, cancellation generations, +/// disposal closure, or the active counts needed for a safe drain. The mutex +/// here protects only non-suspending state transitions and is released before +/// waiting. A short global membership mutex is acceptable for the same reason: +/// the forbidden design is a lock held across request execution, not a lock +/// around bounded map and counter transitions. #[derive(Debug)] -struct VmCoordinatorInner { +struct VmLifecycleGate { connection_id: String, session_id: String, vm_id: String, limits: OwnershipCoordinatorLimits, - state: Mutex, + state: Mutex, changed: Notify, } #[derive(Debug)] -struct VmCoordinatorState { - phase: CoordinatorPhase, +struct VmGateState { + closing: bool, operations: BTreeMap, next_operation_id: u64, next_lifecycle_id: u64, @@ -321,7 +363,6 @@ pub(crate) enum InternalVmEventAdmission { Admitted(CoordinatorOperationPermit), Deferred(CoordinatorOperationPermit), } - #[derive(Debug)] enum VmLifecycleState { Idle, @@ -374,7 +415,7 @@ impl OwnershipCoordinator { connection_id: impl Into, ) -> Result { let connection_id = connection_id.into(); - let mut state = lock(&self.inner.state, "connection registry"); + let mut state = lock(&self.inner.state, "ownership membership"); if state.connections.contains_key(&connection_id) { return Err(OwnershipCoordinatorError::Duplicate { scope: "connection", @@ -387,22 +428,20 @@ impl OwnershipCoordinator { self.inner.limits.max_connections, CONNECTION_LIMIT_PATH, )?; + let generation = take_id(&mut state.next_generation); let connection = ConnectionCoordinator { root: Arc::downgrade(&self.inner), - inner: Arc::new(ConnectionCoordinatorInner { - connection_id: connection_id.clone(), - limits: self.inner.limits, - state: Mutex::new(ConnectionCoordinatorState { - phase: CoordinatorPhase::Open, - sessions: BTreeMap::new(), - operations: BTreeMap::new(), - extension_ordering: BTreeMap::new(), - next_operation_id: 1, - }), - drained: Notify::new(), - }), + connection_id: connection_id.clone(), + generation, }; - state.connections.insert(connection_id, connection.clone()); + state.connections.insert( + connection_id, + ConnectionRecord { + generation, + phase: CoordinatorPhase::Open, + sessions: BTreeMap::new(), + }, + ); Ok(connection) } @@ -410,14 +449,18 @@ impl OwnershipCoordinator { &self, connection_id: &str, ) -> Result { - lock(&self.inner.state, "connection registry") - .connections - .get(connection_id) - .cloned() - .ok_or_else(|| OwnershipCoordinatorError::NotFound { + let state = lock(&self.inner.state, "ownership membership"); + let record = state.connections.get(connection_id).ok_or_else(|| { + OwnershipCoordinatorError::NotFound { scope: "connection", id: connection_id.to_owned(), - }) + } + })?; + Ok(ConnectionCoordinator { + root: Arc::downgrade(&self.inner), + connection_id: connection_id.to_owned(), + generation: record.generation, + }) } pub(crate) async fn admit( @@ -425,61 +468,39 @@ impl OwnershipCoordinator { metadata: &RequestOperationMetadata, cancellation: OperationCancellation, ) -> Result { - validate_ordering_ownership(&metadata.ownership, &metadata.ordering_key)?; + validate_vm_concurrency_ownership(&metadata.ownership, &metadata.vm_concurrency)?; if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } - let resolved = self.resolve(&metadata.ownership)?; - let connection = resolved - .connection - .register_operation(cancellation.clone())?; - let session = match resolved.session.as_ref() { - Some(session) => Some(session.register_operation(cancellation.clone())?), - None => None, - }; - let mut permit = CoordinatorOperationPermit { - _connection: connection, - _session: session, - vm_operation: None, + vm_gate: None, vm_lifecycle: None, - _extension_ordering: None, }; - match &metadata.ordering_key { - RequestOrderingKey::VmOperation { .. } => { - permit.vm_operation = Some( - resolved - .vm - .as_ref() - .expect("VM ordering key requires VM ownership") - .register_operation(cancellation.clone())?, - ); - } - RequestOrderingKey::VmLifecycle { .. } => { - let admission = resolved - .vm - .as_ref() - .expect("VM lifecycle key requires VM ownership") - .begin_lifecycle(cancellation.clone())?; - permit.vm_lifecycle = Some(admission.wait().await?); - } - RequestOrderingKey::Extension { - namespace, - key, - policy, - .. - } => { - if *policy == ExtensionOrderingPolicy::CoreExclusive { - permit._extension_ordering = Some( - resolved - .connection - .register_extension_ordering(namespace.clone(), key.clone())?, + let lifecycle_admission = { + // Membership validation and gate admission share one lock order: + // membership, then the selected VM gate. Disposal uses the same + // order, so an operation is either admitted against one coherent + // entity generation or observes Closing; it cannot slip between a + // path check and gate registration. + let state = lock(&self.inner.state, "ownership membership"); + let gate = validate_ownership_locked(&state, &metadata.ownership)?; + match &metadata.vm_concurrency { + VmConcurrencyClass::SharedVm => { + permit.vm_gate = Some( + gate.expect("shared VM concurrency requires VM ownership") + .register_operation(cancellation.clone())?, ); + None } + VmConcurrencyClass::ExclusiveVmLifecycle => Some( + gate.expect("exclusive VM lifecycle concurrency requires VM ownership") + .begin_lifecycle(cancellation.clone())?, + ), + VmConcurrencyClass::OwnershipOnly => None, } - RequestOrderingKey::Connection(_) - | RequestOrderingKey::Session { .. } - | RequestOrderingKey::Unordered => {} + }; + if let Some(admission) = lifecycle_admission { + permit.vm_lifecycle = Some(admission.wait().await?); } if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); @@ -499,136 +520,31 @@ impl OwnershipCoordinator { metadata: &RequestOperationMetadata, cancellation: OperationCancellation, ) -> Result { - validate_ordering_ownership(&metadata.ownership, &metadata.ordering_key)?; + validate_vm_concurrency_ownership(&metadata.ownership, &metadata.vm_concurrency)?; if !matches!(&metadata.ownership, OwnershipScope::VmOwnership(_)) - || !matches!( - &metadata.ordering_key, - RequestOrderingKey::VmOperation { .. } - ) + || !matches!(&metadata.vm_concurrency, VmConcurrencyClass::SharedVm) { return Err(OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("VM ownership with VM-operation ordering"), + expected: String::from("VM ownership with shared-VM concurrency"), actual: format!( "{}/{}", ownership_label(&metadata.ownership), - ordering_label(&metadata.ordering_key) + vm_concurrency_label(&metadata.vm_concurrency) ), }); } if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } - let resolved = self.resolve(&metadata.ownership)?; - let vm = resolved - .vm - .as_ref() - .expect("internal VM event requires VM ownership"); - let session = resolved - .session - .as_ref() - .expect("internal VM event requires session ownership"); - - // Lock narrowest-to-broadest and register all three scopes as one - // non-suspending ownership transition. If active service capacity is - // unavailable, the claimed event still receives a bounded deferred - // registration, so disposal can cancel and drain it before it ever - // obtains an execution slot. - let mut vm_state = lock(&vm.inner.state, "VM coordinator"); - ensure_open(vm_state.phase, vm.label())?; - let mut session_state = lock(&session.inner.state, "session coordinator"); - ensure_open(session_state.phase, session.label())?; - let mut connection_state = lock(&resolved.connection.inner.state, "connection coordinator"); - ensure_open( - connection_state.phase, - format!("connection {}", resolved.connection.inner.connection_id), - )?; - if let Some(reason) = cancellation.reason() { - return Err(OwnershipCoordinatorError::Cancelled { reason }); - } - - let active_capacity = [&vm_state.operations, &session_state.operations] - .into_iter() - .all(|operations| { - operation_count(operations, OperationAdmissionClass::InternalEvent) - < self.inner.limits.max_internal_event_operations_per_entity - }) - && operation_count( - &connection_state.operations, - OperationAdmissionClass::InternalEvent, - ) < self.inner.limits.max_internal_event_operations_per_entity; - let class = if active_capacity { - OperationAdmissionClass::InternalEvent - } else { - OperationAdmissionClass::DeferredInternalEvent - }; - check_operation_limit( - &vm_state.operations, - class, - self.inner.limits, - "VM-owned operations", - "VM internal-event operations", - )?; - check_operation_limit( - &session_state.operations, - class, - self.inner.limits, - "session-owned operations", - "session internal-event operations", - )?; - check_operation_limit( - &connection_state.operations, - class, - self.inner.limits, - "connection-owned operations", - "connection internal-event operations", - )?; - - let vm_operation_id = take_id(&mut vm_state.next_operation_id); - vm_state.operations.insert( - vm_operation_id, - RegisteredOperation { - cancellation: cancellation.clone(), - class, - }, - ); - let session_operation_id = take_id(&mut session_state.next_operation_id); - session_state.operations.insert( - session_operation_id, - RegisteredOperation { - cancellation: cancellation.clone(), - class, - }, - ); - let connection_operation_id = take_id(&mut connection_state.next_operation_id); - connection_state.operations.insert( - connection_operation_id, - RegisteredOperation { - cancellation: cancellation.clone(), - class, - }, - ); - drop(connection_state); - drop(session_state); - drop(vm_state); - - let vm_operation = VmOperationRegistration { - inner: Arc::clone(&vm.inner), - id: vm_operation_id, - }; - let session = SessionOperationRegistration { - inner: Arc::clone(&session.inner), - id: session_operation_id, - }; - let connection = ConnectionOperationRegistration { - inner: Arc::clone(&resolved.connection.inner), - id: connection_operation_id, + let (gate_permit, class) = { + let state = lock(&self.inner.state, "ownership membership"); + let gate = validate_ownership_locked(&state, &metadata.ownership)? + .expect("internal VM event requires VM ownership"); + gate.register_internal(cancellation.clone())? }; let permit = CoordinatorOperationPermit { - _connection: connection, - _session: Some(session), - vm_operation: Some(vm_operation), + vm_gate: Some(gate_permit), vm_lifecycle: None, - _extension_ordering: None, }; if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); @@ -644,35 +560,42 @@ impl OwnershipCoordinator { }) } - pub(crate) fn begin_session_disposal( + pub(crate) fn begin_vm_disposal( &self, ownership: &OwnershipScope, reason: OperationCancellationReason, - ) -> Result { - let resolved = self.resolve(ownership)?; - let session = - resolved - .session - .ok_or_else(|| OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("session or VM ownership"), - actual: ownership_label(ownership), - })?; - session.begin_disposal(reason) + ) -> Result { + let (connection_id, session_id, vm_id) = ownership_ids(ownership); + let (Some(session_id), Some(vm_id)) = (session_id, vm_id) else { + return Err(OwnershipCoordinatorError::OwnershipMismatch { + expected: String::from("VM ownership"), + actual: ownership_label(ownership), + }); + }; + self.connection(connection_id)? + .session(session_id)? + .vm(vm_id)? + .begin_disposal(reason, None) } - pub(crate) fn begin_vm_disposal( + /// Begin VM disposal while closing the same ownership subtree in the + /// authoritative operation table. The dispose request is excluded from its + /// own drain; every other ordinary or progress operation is cancelled and + /// must release before teardown proceeds. + pub(crate) fn begin_vm_disposal_with_operations( &self, ownership: &OwnershipScope, reason: OperationCancellationReason, + operations: &OperationTable, + excluded: &RequestOperationKey, ) -> Result { - let resolved = self.resolve(ownership)?; - let vm = resolved - .vm - .ok_or_else(|| OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("VM ownership"), - actual: ownership_label(ownership), - })?; - vm.begin_disposal(reason) + // Close membership/gate admission first. A request that reaches the + // operation table in the tiny interval before scope closure still + // fails the coherent membership check and cannot touch the VM. + let mut disposal = self.begin_vm_disposal(ownership, reason)?; + disposal.operation_drain = + Some(operations.close_scope(ownership.clone(), reason, Some(excluded))); + Ok(disposal) } pub(crate) fn begin_connection_disposal( @@ -682,78 +605,49 @@ impl OwnershipCoordinator { ) -> Result { self.connection(connection_id)?.begin_disposal(reason) } - - fn resolve( - &self, - ownership: &OwnershipScope, - ) -> Result { - let (connection_id, session_id, vm_id) = ownership_ids(ownership); - let connection = self.connection(connection_id)?; - let session = match session_id { - Some(session_id) => Some(connection.session(session_id)?), - None => None, - }; - let vm = match (session.as_ref(), vm_id) { - (Some(session), Some(vm_id)) => Some(session.vm(vm_id)?), - _ => None, - }; - Ok(ResolvedCoordinators { - connection, - session, - vm, - }) - } -} - -struct ResolvedCoordinators { - connection: ConnectionCoordinator, - session: Option, - vm: Option, } impl ConnectionCoordinator { - pub(crate) fn connection_id(&self) -> &str { - &self.inner.connection_id - } - pub(crate) fn open_session( &self, session_id: impl Into, ) -> Result { let session_id = session_id.into(); - let mut state = lock(&self.inner.state, "connection coordinator"); + let root = upgrade_root(&self.root, "connection", &self.connection_id)?; + let mut state = lock(&root.state, "ownership membership"); + let connection = matching_connection_mut(&mut state, self)?; ensure_open( - state.phase, - format!("connection {}", self.inner.connection_id), + connection.phase, + format!("connection {}", self.connection_id), )?; - if state.sessions.contains_key(&session_id) { + if connection.sessions.contains_key(&session_id) { return Err(OwnershipCoordinatorError::Duplicate { scope: "session", - id: format!("{}:{session_id}", self.inner.connection_id), + id: format!("{}:{session_id}", self.connection_id), }); } check_limit( "session coordinators per connection", - state.sessions.len(), - self.inner.limits.max_sessions_per_connection, + connection.sessions.len(), + root.limits.max_sessions_per_connection, SESSION_LIMIT_PATH, )?; + let generation = take_id(&mut state.next_generation); + matching_connection_mut(&mut state, self)?.sessions.insert( + session_id.clone(), + SessionRecord { + generation, + phase: CoordinatorPhase::Open, + vms: BTreeMap::new(), + }, + ); let session = SessionCoordinator { - parent: Arc::downgrade(&self.inner), - inner: Arc::new(SessionCoordinatorInner { - connection_id: self.inner.connection_id.clone(), - session_id: session_id.clone(), - limits: self.inner.limits, - state: Mutex::new(SessionCoordinatorState { - phase: CoordinatorPhase::Open, - vms: BTreeMap::new(), - operations: BTreeMap::new(), - next_operation_id: 1, - }), - drained: Notify::new(), - }), + root: Arc::downgrade(&root), + connection_id: self.connection_id.clone(), + connection_generation: self.generation, + session_id, + generation, }; - state.sessions.insert(session_id, session.clone()); Ok(session) } @@ -761,22 +655,42 @@ impl ConnectionCoordinator { &self, session_id: &str, ) -> Result { - lock(&self.inner.state, "connection coordinator") - .sessions - .get(session_id) - .cloned() - .ok_or_else(|| OwnershipCoordinatorError::NotFound { + let root = upgrade_root(&self.root, "connection", &self.connection_id)?; + let state = lock(&root.state, "ownership membership"); + let connection = matching_connection(&state, self)?; + let session = connection.sessions.get(session_id).ok_or_else(|| { + OwnershipCoordinatorError::NotFound { scope: "session", - id: format!("{}:{session_id}", self.inner.connection_id), - }) + id: format!("{}:{session_id}", self.connection_id), + } + })?; + Ok(SessionCoordinator { + root: Arc::downgrade(&root), + connection_id: self.connection_id.clone(), + connection_generation: self.generation, + session_id: session_id.to_owned(), + generation: session.generation, + }) } + #[cfg(test)] pub(crate) fn snapshot(&self) -> EntityCoordinatorSnapshot { - let state = lock(&self.inner.state, "connection coordinator"); + let Some(root) = self.root.upgrade() else { + return closed_entity_snapshot(); + }; + let state = lock(&root.state, "ownership membership"); + let Ok(connection) = matching_connection(&state, self) else { + return closed_entity_snapshot(); + }; EntityCoordinatorSnapshot { - phase: state.phase, - active_operations: state.operations.len(), - child_count: state.sessions.len(), + phase: connection.phase, + active_operations: connection + .sessions + .values() + .flat_map(|session| session.vms.values()) + .map(|vm| vm.gate.active_operations()) + .sum(), + child_count: connection.sessions.len(), } } @@ -784,223 +698,267 @@ impl ConnectionCoordinator { &self, reason: OperationCancellationReason, ) -> Result { - let sessions = { - let mut state = lock(&self.inner.state, "connection coordinator"); + let root = upgrade_root(&self.root, "connection", &self.connection_id)?; + let gates = { + let mut state = lock(&root.state, "ownership membership"); + let connection = matching_connection_mut(&mut state, self)?; ensure_open( - state.phase, - format!("connection {}", self.inner.connection_id), + connection.phase, + format!("connection {}", self.connection_id), )?; - state.phase = CoordinatorPhase::Closing; - signal_all(&state.operations, reason); - state.sessions.values().cloned().collect::>() + connection.phase = CoordinatorPhase::Closing; + let mut gates = Vec::new(); + for session in connection.sessions.values_mut() { + session.phase = CoordinatorPhase::Closing; + for vm in session.vms.values_mut() { + vm.phase = CoordinatorPhase::Closing; + vm.gate.begin_closing(reason); + gates.push(Arc::clone(&vm.gate)); + } + } + gates }; - for session in &sessions { - session.force_closing(reason); - } Ok(ConnectionDisposal { - connection: self.clone(), - sessions, + root, + connection_id: self.connection_id.clone(), + generation: self.generation, + gates, completed: false, }) } - - fn register_operation( - &self, - cancellation: OperationCancellation, - ) -> Result { - self.register_operation_with_class(cancellation, OperationAdmissionClass::Ordinary) - } - - fn register_operation_with_class( - &self, - cancellation: OperationCancellation, - class: OperationAdmissionClass, - ) -> Result { - if let Some(reason) = cancellation.reason() { - return Err(OwnershipCoordinatorError::Cancelled { reason }); - } - let mut state = lock(&self.inner.state, "connection coordinator"); - ensure_open( - state.phase, - format!("connection {}", self.inner.connection_id), - )?; - check_operation_limit( - &state.operations, - class, - self.inner.limits, - "connection-owned operations", - "connection internal-event operations", - )?; - let id = take_id(&mut state.next_operation_id); - state.operations.insert( - id, - RegisteredOperation { - cancellation, - class, - }, - ); - Ok(ConnectionOperationRegistration { - inner: Arc::clone(&self.inner), - id, - }) - } - - fn register_extension_ordering( - &self, - namespace: String, - key: Vec, - ) -> Result { - let mut state = lock(&self.inner.state, "connection coordinator"); - ensure_open( - state.phase, - format!("connection {}", self.inner.connection_id), - )?; - let ordering_key = (namespace, key); - if state.extension_ordering.contains_key(&ordering_key) { - return Err(OwnershipCoordinatorError::OrderingConflict { - scope: format!( - "connection {}/extension/{}/{}-bytes", - self.inner.connection_id, - ordering_key.0, - ordering_key.1.len() - ), - }); - } - state.extension_ordering.insert(ordering_key.clone(), ()); - Ok(ExtensionOrderingRegistration { - inner: Arc::clone(&self.inner), - key: ordering_key, - }) - } } impl SessionCoordinator { - pub(crate) fn connection_id(&self) -> &str { - &self.inner.connection_id - } - - pub(crate) fn session_id(&self) -> &str { - &self.inner.session_id - } - pub(crate) fn open_vm( &self, vm_id: impl Into, ) -> Result { let vm_id = vm_id.into(); - let mut state = lock(&self.inner.state, "session coordinator"); - ensure_open(state.phase, self.label())?; - if state.vms.contains_key(&vm_id) { - return Err(OwnershipCoordinatorError::Duplicate { - scope: "VM", - id: format!("{}:{vm_id}", self.label()), - }); + let root = upgrade_root(&self.root, "session", &self.label())?; + let mut state = lock(&root.state, "ownership membership"); + { + let session = matching_session(&state, self)?; + ensure_open(session.phase, self.label())?; + if session.vms.contains_key(&vm_id) { + return Err(OwnershipCoordinatorError::Duplicate { + scope: "VM", + id: format!("{}:{vm_id}", self.label()), + }); + } + check_limit( + "VM coordinators per session", + session.vms.len(), + root.limits.max_vms_per_session, + VM_LIMIT_PATH, + )?; } - check_limit( - "VM coordinators per session", - state.vms.len(), - self.inner.limits.max_vms_per_session, - VM_LIMIT_PATH, - )?; - let vm = VmCoordinator { - parent: Arc::downgrade(&self.inner), - inner: Arc::new(VmCoordinatorInner { - connection_id: self.inner.connection_id.clone(), - session_id: self.inner.session_id.clone(), - vm_id: vm_id.clone(), - limits: self.inner.limits, - state: Mutex::new(VmCoordinatorState { - phase: CoordinatorPhase::Open, - operations: BTreeMap::new(), - next_operation_id: 1, - next_lifecycle_id: 1, - lifecycle: VmLifecycleState::Idle, - }), - changed: Notify::new(), + let generation = take_id(&mut state.next_generation); + let gate = Arc::new(VmLifecycleGate { + connection_id: self.connection_id.clone(), + session_id: self.session_id.clone(), + vm_id: vm_id.clone(), + limits: root.limits, + state: Mutex::new(VmGateState { + closing: false, + operations: BTreeMap::new(), + next_operation_id: 1, + next_lifecycle_id: 1, + lifecycle: VmLifecycleState::Idle, }), - }; - state.vms.insert(vm_id, vm.clone()); - Ok(vm) + changed: Notify::new(), + }); + matching_session_mut(&mut state, self)?.vms.insert( + vm_id.clone(), + VmRecord { + generation, + phase: CoordinatorPhase::Open, + gate: Arc::clone(&gate), + }, + ); + Ok(VmCoordinator { + root: Arc::downgrade(&root), + connection_id: self.connection_id.clone(), + connection_generation: self.connection_generation, + session_id: self.session_id.clone(), + session_generation: self.generation, + vm_id, + generation, + gate, + }) } pub(crate) fn vm(&self, vm_id: &str) -> Result { - lock(&self.inner.state, "session coordinator") + let root = upgrade_root(&self.root, "session", &self.label())?; + let state = lock(&root.state, "ownership membership"); + let session = matching_session(&state, self)?; + let vm = session .vms .get(vm_id) - .cloned() .ok_or_else(|| OwnershipCoordinatorError::NotFound { scope: "VM", id: format!("{}:{vm_id}", self.label()), - }) + })?; + Ok(VmCoordinator { + root: Arc::downgrade(&root), + connection_id: self.connection_id.clone(), + connection_generation: self.connection_generation, + session_id: self.session_id.clone(), + session_generation: self.generation, + vm_id: vm_id.to_owned(), + generation: vm.generation, + gate: Arc::clone(&vm.gate), + }) } + #[cfg(test)] pub(crate) fn snapshot(&self) -> EntityCoordinatorSnapshot { - let state = lock(&self.inner.state, "session coordinator"); + let Some(root) = self.root.upgrade() else { + return closed_entity_snapshot(); + }; + let state = lock(&root.state, "ownership membership"); + let Ok(session) = matching_session(&state, self) else { + return closed_entity_snapshot(); + }; EntityCoordinatorSnapshot { - phase: state.phase, + phase: session.phase, + active_operations: session + .vms + .values() + .map(|vm| vm.gate.active_operations()) + .sum(), + child_count: session.vms.len(), + } + } + + fn label(&self) -> String { + format!("session {}:{}", self.connection_id, self.session_id) + } +} + +impl VmCoordinator { + #[cfg(test)] + pub(crate) fn snapshot(&self) -> VmCoordinatorSnapshot { + let phase = self + .root + .upgrade() + .and_then(|root| { + let state = lock(&root.state, "ownership membership"); + matching_vm(&state, self).ok().map(|vm| vm.phase) + }) + .unwrap_or(CoordinatorPhase::Closed); + let state = lock(&self.gate.state, "VM lifecycle gate"); + VmCoordinatorSnapshot { + phase, active_operations: state.operations.len(), - child_count: state.vms.len(), + lifecycle: state.lifecycle.phase(), } } + #[cfg(test)] + fn begin_lifecycle( + &self, + cancellation: OperationCancellation, + ) -> Result { + let root = upgrade_root(&self.root, "VM", &self.label())?; + let state = lock(&root.state, "ownership membership"); + let vm = matching_vm(&state, self)?; + ensure_open(vm.phase, self.label())?; + self.gate.begin_lifecycle(cancellation) + } + pub(crate) fn begin_disposal( &self, reason: OperationCancellationReason, - ) -> Result { - let vms = { - let mut state = lock(&self.inner.state, "session coordinator"); - ensure_open(state.phase, self.label())?; - state.phase = CoordinatorPhase::Closing; - signal_all(&state.operations, reason); - state.vms.values().cloned().collect::>() - }; - for vm in &vms { - vm.begin_closing(reason); + operation_drain: Option, + ) -> Result { + let root = upgrade_root(&self.root, "VM", &self.label())?; + { + let mut state = lock(&root.state, "ownership membership"); + let vm = matching_vm_mut(&mut state, self)?; + ensure_open(vm.phase, self.label())?; + vm.phase = CoordinatorPhase::Closing; + vm.gate.begin_closing(reason); } - Ok(SessionDisposal { - session: self.clone(), - vms, + Ok(VmDisposal { + root, + connection_id: self.connection_id.clone(), + connection_generation: self.connection_generation, + session_id: self.session_id.clone(), + session_generation: self.session_generation, + vm_id: self.vm_id.clone(), + generation: self.generation, + gate: Arc::clone(&self.gate), + operation_drain, completed: false, }) } - fn force_closing(&self, reason: OperationCancellationReason) { - let vms = { - let mut state = lock(&self.inner.state, "session coordinator"); - if state.phase == CoordinatorPhase::Open { - state.phase = CoordinatorPhase::Closing; + fn label(&self) -> String { + format!( + "VM {}:{}:{}", + self.connection_id, self.session_id, self.vm_id + ) + } +} + +impl VmLifecycleGate { + fn label(&self) -> String { + format!( + "VM {}:{}:{}", + self.connection_id, self.session_id, self.vm_id + ) + } + + fn active_operations(&self) -> usize { + let state = lock(&self.state, "VM lifecycle gate"); + state.operations.len() + usize::from(!matches!(state.lifecycle, VmLifecycleState::Idle)) + } + + async fn wait_drained(&self) { + loop { + let changed = self.changed.notified(); + if self.active_operations() == 0 { + return; } - signal_all(&state.operations, reason); - state.vms.values().cloned().collect::>() - }; - for vm in &vms { - vm.begin_closing(reason); + changed.await; } } fn register_operation( - &self, + self: &Arc, cancellation: OperationCancellation, - ) -> Result { + ) -> Result { self.register_operation_with_class(cancellation, OperationAdmissionClass::Ordinary) } - fn register_operation_with_class( - &self, + fn register_internal( + self: &Arc, cancellation: OperationCancellation, - class: OperationAdmissionClass, - ) -> Result { + ) -> Result<(VmGatePermit, OperationAdmissionClass), OwnershipCoordinatorError> { if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } - let mut state = lock(&self.inner.state, "session coordinator"); - ensure_open(state.phase, self.label())?; + let mut state = lock(&self.state, "VM lifecycle gate"); + if state.closing { + return Err(OwnershipCoordinatorError::Closing { + scope: self.label(), + phase: CoordinatorPhase::Closing, + }); + } + let class = if matches!(state.lifecycle, VmLifecycleState::Active { .. }) + || operation_count(&state.operations, OperationAdmissionClass::InternalEvent) + >= self.limits.max_internal_event_operations_per_entity + { + OperationAdmissionClass::DeferredInternalEvent + } else { + OperationAdmissionClass::InternalEvent + }; check_operation_limit( &state.operations, class, - self.inner.limits, - "session-owned operations", - "session internal-event operations", + self.limits, + "VM-owned operations", + "VM internal-event operations", )?; let id = take_id(&mut state.next_operation_id); state.operations.insert( @@ -1010,59 +968,30 @@ impl SessionCoordinator { class, }, ); - Ok(SessionOperationRegistration { - inner: Arc::clone(&self.inner), - id, - }) - } - - fn label(&self) -> String { - format!( - "session {}:{}", - self.inner.connection_id, self.inner.session_id - ) - } -} - -impl VmCoordinator { - pub(crate) fn connection_id(&self) -> &str { - &self.inner.connection_id - } - - pub(crate) fn session_id(&self) -> &str { - &self.inner.session_id - } - - pub(crate) fn vm_id(&self) -> &str { - &self.inner.vm_id - } - - pub(crate) fn snapshot(&self) -> VmCoordinatorSnapshot { - let state = lock(&self.inner.state, "VM coordinator"); - VmCoordinatorSnapshot { - phase: state.phase, - active_operations: state.operations.len(), - lifecycle: state.lifecycle.phase(), - } - } - - fn register_operation( - &self, - cancellation: OperationCancellation, - ) -> Result { - self.register_operation_with_class(cancellation, OperationAdmissionClass::Ordinary) + Ok(( + VmGatePermit { + gate: Arc::clone(self), + id, + }, + class, + )) } fn register_operation_with_class( - &self, + self: &Arc, cancellation: OperationCancellation, class: OperationAdmissionClass, - ) -> Result { + ) -> Result { if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } - let mut state = lock(&self.inner.state, "VM coordinator"); - ensure_open(state.phase, self.label())?; + let mut state = lock(&self.state, "VM lifecycle gate"); + if state.closing { + return Err(OwnershipCoordinatorError::Closing { + scope: self.label(), + phase: CoordinatorPhase::Closing, + }); + } if class == OperationAdmissionClass::Ordinary && !matches!(state.lifecycle, VmLifecycleState::Idle) { @@ -1074,7 +1003,7 @@ impl VmCoordinator { check_operation_limit( &state.operations, class, - self.inner.limits, + self.limits, "VM-owned operations", "VM internal-event operations", )?; @@ -1086,22 +1015,27 @@ impl VmCoordinator { class, }, ); - Ok(VmOperationRegistration { - inner: Arc::clone(&self.inner), + Ok(VmGatePermit { + gate: Arc::clone(self), id, }) } fn begin_lifecycle( - &self, + self: &Arc, cancellation: OperationCancellation, ) -> Result { if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } let id = { - let mut state = lock(&self.inner.state, "VM coordinator"); - ensure_open(state.phase, self.label())?; + let mut state = lock(&self.state, "VM lifecycle gate"); + if state.closing { + return Err(OwnershipCoordinatorError::Closing { + scope: self.label(), + phase: CoordinatorPhase::Closing, + }); + } if !matches!(state.lifecycle, VmLifecycleState::Idle) { return Err(OwnershipCoordinatorError::LifecycleConflict { vm: self.label(), @@ -1116,7 +1050,7 @@ impl VmCoordinator { id }; Ok(VmLifecycleAdmission { - inner: Arc::clone(&self.inner), + gate: Arc::clone(self), id, cancellation, armed: true, @@ -1124,229 +1058,108 @@ impl VmCoordinator { } fn begin_closing(&self, reason: OperationCancellationReason) { - let mut state = lock(&self.inner.state, "VM coordinator"); - if state.phase == CoordinatorPhase::Open { - state.phase = CoordinatorPhase::Closing; - } + let mut state = lock(&self.state, "VM lifecycle gate"); + state.closing = true; signal_all(&state.operations, reason); state.lifecycle.signal(reason); - self.inner.changed.notify_one(); - } - - pub(crate) fn begin_disposal( - &self, - reason: OperationCancellationReason, - ) -> Result { - { - let state = lock(&self.inner.state, "VM coordinator"); - ensure_open(state.phase, self.label())?; - } - self.begin_closing(reason); - Ok(VmDisposal { - vm: self.clone(), - completed: false, - }) - } - - fn label(&self) -> String { - format!( - "VM {}:{}:{}", - self.inner.connection_id, self.inner.session_id, self.inner.vm_id - ) + self.changed.notify_waiters(); } } #[derive(Debug)] pub(crate) struct CoordinatorOperationPermit { vm_lifecycle: Option, - vm_operation: Option, - _extension_ordering: Option, - _session: Option, - _connection: ConnectionOperationRegistration, + vm_gate: Option, } impl CoordinatorOperationPermit { + #[cfg(test)] pub(crate) fn is_vm_lifecycle(&self) -> bool { self.vm_lifecycle.is_some() } - pub(crate) fn is_vm_operation(&self) -> bool { - self.vm_operation.is_some() - } - /// Promote one already-tracked internal event into active service /// capacity. A false result keeps the same ownership registrations live; /// disposal can therefore cancel and drain the event while it waits. pub(crate) fn try_activate_deferred_internal_event( &mut self, ) -> Result { - let vm = self.vm_operation.as_ref().ok_or_else(|| { - OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("VM-bound internal-event permit"), - actual: String::from("permit without VM operation registration"), - } - })?; - let session = - self._session + let permit = + self.vm_gate .as_ref() .ok_or_else(|| OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("session-bound internal-event permit"), - actual: String::from("permit without session registration"), + expected: String::from("VM-bound internal-event permit"), + actual: String::from("permit without VM gate admission"), })?; - - let mut vm_state = lock(&vm.inner.state, "VM coordinator"); - let mut session_state = lock(&session.inner.state, "session coordinator"); - let mut connection_state = lock(&self._connection.inner.state, "connection coordinator"); + let mut vm_state = lock(&permit.gate.state, "VM lifecycle gate"); let cancellation = vm_state .operations - .get(&vm.id) + .get(&permit.id) .ok_or_else(|| OwnershipCoordinatorError::NotFound { - scope: "VM operation", - id: vm.id.to_string(), + scope: "VM gate permit", + id: permit.id.to_string(), })? .cancellation .clone(); if let Some(reason) = cancellation.reason() { return Err(OwnershipCoordinatorError::Cancelled { reason }); } - ensure_open(vm_state.phase, format!("VM {}", vm.inner.vm_id))?; - ensure_open( - session_state.phase, - format!( - "session {}:{}", - session.inner.connection_id, session.inner.session_id - ), - )?; - ensure_open( - connection_state.phase, - format!("connection {}", self._connection.inner.connection_id), - )?; - - let classes = [ - vm_state - .operations - .get(&vm.id) - .map(|operation| operation.class), - session_state - .operations - .get(&session.id) - .map(|operation| operation.class), - connection_state - .operations - .get(&self._connection.id) - .map(|operation| operation.class), - ]; - if classes - .iter() - .all(|class| *class == Some(OperationAdmissionClass::InternalEvent)) - { + if vm_state.closing { + return Err(OwnershipCoordinatorError::Closing { + scope: permit.gate.label(), + phase: CoordinatorPhase::Closing, + }); + } + if matches!(vm_state.lifecycle, VmLifecycleState::Active { .. }) { + return Ok(false); + } + let class = vm_state + .operations + .get(&permit.id) + .map(|operation| operation.class); + if class == Some(OperationAdmissionClass::InternalEvent) { return Ok(true); } - if !classes - .iter() - .all(|class| *class == Some(OperationAdmissionClass::DeferredInternalEvent)) - { + if class != Some(OperationAdmissionClass::DeferredInternalEvent) { return Err(OwnershipCoordinatorError::OwnershipMismatch { - expected: String::from("matching deferred internal-event registrations"), - actual: format!("registration classes {classes:?}"), + expected: String::from("deferred internal-event registration"), + actual: format!("registration class {class:?}"), }); } - let limit = vm.inner.limits.max_internal_event_operations_per_entity; + let limit = permit.gate.limits.max_internal_event_operations_per_entity; let has_capacity = - operation_count(&vm_state.operations, OperationAdmissionClass::InternalEvent) < limit - && operation_count( - &session_state.operations, - OperationAdmissionClass::InternalEvent, - ) < limit - && operation_count( - &connection_state.operations, - OperationAdmissionClass::InternalEvent, - ) < limit; + operation_count(&vm_state.operations, OperationAdmissionClass::InternalEvent) < limit; if !has_capacity { - return Ok(false); - } - vm_state - .operations - .get_mut(&vm.id) - .expect("deferred VM registration checked above") - .class = OperationAdmissionClass::InternalEvent; - session_state - .operations - .get_mut(&session.id) - .expect("deferred session registration checked above") - .class = OperationAdmissionClass::InternalEvent; - connection_state - .operations - .get_mut(&self._connection.id) - .expect("deferred connection registration checked above") - .class = OperationAdmissionClass::InternalEvent; - Ok(true) - } -} - -#[derive(Debug)] -struct ConnectionOperationRegistration { - inner: Arc, - id: u64, -} - -impl Drop for ConnectionOperationRegistration { - fn drop(&mut self) { - let mut state = lock(&self.inner.state, "connection coordinator"); - if state.operations.remove(&self.id).is_some() { - self.inner.drained.notify_one(); - } - } -} - -#[derive(Debug)] -struct ExtensionOrderingRegistration { - inner: Arc, - key: (String, Vec), -} - -impl Drop for ExtensionOrderingRegistration { - fn drop(&mut self) { - lock(&self.inner.state, "connection coordinator") - .extension_ordering - .remove(&self.key); - } -} - -#[derive(Debug)] -struct SessionOperationRegistration { - inner: Arc, - id: u64, -} - -impl Drop for SessionOperationRegistration { - fn drop(&mut self) { - let mut state = lock(&self.inner.state, "session coordinator"); - if state.operations.remove(&self.id).is_some() { - self.inner.drained.notify_one(); + return Ok(false); } + vm_state + .operations + .get_mut(&permit.id) + .expect("deferred VM registration checked above") + .class = OperationAdmissionClass::InternalEvent; + Ok(true) } } #[derive(Debug)] -struct VmOperationRegistration { - inner: Arc, +struct VmGatePermit { + gate: Arc, id: u64, } -impl Drop for VmOperationRegistration { +impl Drop for VmGatePermit { fn drop(&mut self) { - let mut state = lock(&self.inner.state, "VM coordinator"); + let mut state = lock(&self.gate.state, "VM lifecycle gate"); if state.operations.remove(&self.id).is_some() { - self.inner.changed.notify_one(); + self.gate.changed.notify_waiters(); } } } #[derive(Debug)] struct VmLifecycleAdmission { - inner: Arc, + gate: Arc, id: u64, cancellation: OperationCancellation, armed: bool, @@ -1355,9 +1168,14 @@ struct VmLifecycleAdmission { impl VmLifecycleAdmission { async fn wait(mut self) -> Result { loop { - let changed = self.inner.changed.notified(); + // Create the notification future before checking state. A permit + // released after this point either changes the state we observe or + // wakes this registered listener, avoiding a check-then-sleep lost + // wakeup. `Notify` is only a hint: the state and lifecycle id remain + // authoritative, so spurious or coalesced wakes are harmless. + let changed = self.gate.changed.notified(); { - let mut state = lock(&self.inner.state, "VM coordinator"); + let mut state = lock(&self.gate.state, "VM lifecycle gate"); match &state.lifecycle { VmLifecycleState::Pending { id, .. } if *id == self.id => { if state.operations.is_empty() { @@ -1368,7 +1186,7 @@ impl VmLifecycleAdmission { }; self.armed = false; return Ok(VmLifecycleGuard { - inner: Arc::clone(&self.inner), + gate: Arc::clone(&self.gate), id: self.id, }); } @@ -1395,25 +1213,28 @@ impl VmLifecycleAdmission { impl Drop for VmLifecycleAdmission { fn drop(&mut self) { if self.armed { - release_lifecycle(&self.inner, self.id); + release_lifecycle(&self.gate, self.id); } } } #[derive(Debug)] struct VmLifecycleGuard { - inner: Arc, + gate: Arc, id: u64, } impl Drop for VmLifecycleGuard { fn drop(&mut self) { - release_lifecycle(&self.inner, self.id); + release_lifecycle(&self.gate, self.id); } } -fn release_lifecycle(inner: &Arc, id: u64) { - let mut state = lock(&inner.state, "VM coordinator"); +fn release_lifecycle(gate: &Arc, id: u64) { + let mut state = lock(&gate.state, "VM lifecycle gate"); + // The id prevents a stale dropped admission/guard from reopening a newer + // lifecycle generation. Both cancellation before activation and normal + // completion converge here through RAII. let matches_id = match &state.lifecycle { VmLifecycleState::Pending { id: active, .. } | VmLifecycleState::Active { id: active, .. } => *active == id, @@ -1421,129 +1242,67 @@ fn release_lifecycle(inner: &Arc, id: u64) { }; if matches_id { state.lifecycle = VmLifecycleState::Idle; - inner.changed.notify_one(); - } -} - -#[derive(Debug)] -pub(crate) struct SessionDisposal { - session: SessionCoordinator, - vms: Vec, - completed: bool, -} - -impl SessionDisposal { - pub(crate) async fn wait_drained(&self) { - loop { - let drained = self.session.inner.drained.notified(); - if lock(&self.session.inner.state, "session coordinator") - .operations - .is_empty() - { - return; - } - drained.await; - } - } - - pub(crate) fn complete(mut self) -> Result<(), OwnershipCoordinatorError> { - let active_operations = lock(&self.session.inner.state, "session coordinator") - .operations - .len(); - if active_operations != 0 { - return Err(OwnershipCoordinatorError::NotDrained { - scope: self.session.label(), - active_operations, - }); - } - for vm in &self.vms { - let mut state = lock(&vm.inner.state, "VM coordinator"); - let vm_active = state.operations.len() - + usize::from(!matches!(state.lifecycle, VmLifecycleState::Idle)); - if vm_active != 0 { - return Err(OwnershipCoordinatorError::NotDrained { - scope: vm.label(), - active_operations: vm_active, - }); - } - state.phase = CoordinatorPhase::Closed; - } - { - let mut state = lock(&self.session.inner.state, "session coordinator"); - state.phase = CoordinatorPhase::Closed; - state.vms.clear(); - } - if let Some(parent) = self.session.parent.upgrade() { - let mut state = lock(&parent.state, "connection coordinator"); - if state - .sessions - .get(&self.session.inner.session_id) - .is_some_and(|current| Arc::ptr_eq(¤t.inner, &self.session.inner)) - { - state.sessions.remove(&self.session.inner.session_id); - } - } - self.completed = true; - Ok(()) - } -} - -impl Drop for SessionDisposal { - fn drop(&mut self) { - if !self.completed { - tracing::warn!( - connection_id = %self.session.inner.connection_id, - session_id = %self.session.inner.session_id, - "ERR_AGENTOS_SESSION_DISPOSAL_INCOMPLETE: session remains Closing for a bounded retry" - ); - } + gate.changed.notify_waiters(); } } #[derive(Debug)] pub(crate) struct VmDisposal { - vm: VmCoordinator, + root: Arc, + connection_id: String, + connection_generation: u64, + session_id: String, + session_generation: u64, + vm_id: String, + generation: u64, + gate: Arc, + operation_drain: Option, completed: bool, } impl VmDisposal { pub(crate) async fn wait_drained(&self) { - loop { - let changed = self.vm.inner.changed.notified(); - let drained = { - let state = lock(&self.vm.inner.state, "VM coordinator"); - state.operations.is_empty() && matches!(state.lifecycle, VmLifecycleState::Idle) - }; - if drained { - return; - } - changed.await; + if let Some(drain) = &self.operation_drain { + drain.wait_drained().await; } + self.gate.wait_drained().await; } pub(crate) fn complete(mut self) -> Result<(), OwnershipCoordinatorError> { - { - let mut state = lock(&self.vm.inner.state, "VM coordinator"); - let active_operations = state.operations.len() - + usize::from(!matches!(state.lifecycle, VmLifecycleState::Idle)); + if let Some(drain) = &self.operation_drain { + let active_operations = drain.active_operations(); if active_operations != 0 { return Err(OwnershipCoordinatorError::NotDrained { - scope: self.vm.label(), + scope: self.gate.label(), active_operations, }); } - state.phase = CoordinatorPhase::Closed; } - if let Some(parent) = self.vm.parent.upgrade() { - let mut state = lock(&parent.state, "session coordinator"); - if state - .vms - .get(&self.vm.inner.vm_id) - .is_some_and(|current| Arc::ptr_eq(¤t.inner, &self.vm.inner)) - { - state.vms.remove(&self.vm.inner.vm_id); - } + let active_operations = self.gate.active_operations(); + if active_operations != 0 { + return Err(OwnershipCoordinatorError::NotDrained { + scope: self.gate.label(), + active_operations, + }); } + let mut state = lock(&self.root.state, "ownership membership"); + let connection = state + .connections + .get_mut(&self.connection_id) + .filter(|connection| connection.generation == self.connection_generation) + .ok_or_else(|| stale_generation("connection", &self.connection_id))?; + let session = connection + .sessions + .get_mut(&self.session_id) + .filter(|session| session.generation == self.session_generation) + .ok_or_else(|| stale_generation("session", &self.session_id))?; + let vm = session + .vms + .get(&self.vm_id) + .filter(|vm| vm.generation == self.generation && Arc::ptr_eq(&vm.gate, &self.gate)) + .ok_or_else(|| stale_generation("VM", &self.vm_id))?; + ensure_closing(vm.phase, self.gate.label())?; + session.vms.remove(&self.vm_id); self.completed = true; Ok(()) } @@ -1553,9 +1312,9 @@ impl Drop for VmDisposal { fn drop(&mut self) { if !self.completed { tracing::warn!( - connection_id = %self.vm.inner.connection_id, - session_id = %self.vm.inner.session_id, - vm_id = %self.vm.inner.vm_id, + connection_id = %self.connection_id, + session_id = %self.session_id, + vm_id = %self.vm_id, "ERR_AGENTOS_VM_DISPOSAL_INCOMPLETE: VM remains Closing for a bounded retry" ); } @@ -1564,75 +1323,39 @@ impl Drop for VmDisposal { #[derive(Debug)] pub(crate) struct ConnectionDisposal { - connection: ConnectionCoordinator, - sessions: Vec, + root: Arc, + connection_id: String, + generation: u64, + gates: Vec>, completed: bool, } impl ConnectionDisposal { pub(crate) async fn wait_drained(&self) { - loop { - let drained = self.connection.inner.drained.notified(); - if lock(&self.connection.inner.state, "connection coordinator") - .operations - .is_empty() - { - return; - } - drained.await; + for gate in &self.gates { + gate.wait_drained().await; } } pub(crate) fn complete(mut self) -> Result<(), OwnershipCoordinatorError> { - let active_operations = lock(&self.connection.inner.state, "connection coordinator") - .operations - .len(); + let active_operations = self.gates.iter().map(|gate| gate.active_operations()).sum(); if active_operations != 0 { return Err(OwnershipCoordinatorError::NotDrained { - scope: format!("connection {}", self.connection.inner.connection_id), + scope: format!("connection {}", self.connection_id), active_operations, }); } - for session in &self.sessions { - let mut session_state = lock(&session.inner.state, "session coordinator"); - if !session_state.operations.is_empty() { - return Err(OwnershipCoordinatorError::NotDrained { - scope: session.label(), - active_operations: session_state.operations.len(), - }); - } - for vm in session_state.vms.values() { - let mut vm_state = lock(&vm.inner.state, "VM coordinator"); - let vm_active = vm_state.operations.len() - + usize::from(!matches!(vm_state.lifecycle, VmLifecycleState::Idle)); - if vm_active != 0 { - return Err(OwnershipCoordinatorError::NotDrained { - scope: vm.label(), - active_operations: vm_active, - }); - } - vm_state.phase = CoordinatorPhase::Closed; - } - session_state.vms.clear(); - session_state.phase = CoordinatorPhase::Closed; - } - { - let mut state = lock(&self.connection.inner.state, "connection coordinator"); - state.sessions.clear(); - state.phase = CoordinatorPhase::Closed; - } - if let Some(root) = self.connection.root.upgrade() { - let mut state = lock(&root.state, "connection registry"); - if state - .connections - .get(&self.connection.inner.connection_id) - .is_some_and(|current| Arc::ptr_eq(¤t.inner, &self.connection.inner)) - { - state - .connections - .remove(&self.connection.inner.connection_id); - } - } + let mut state = lock(&self.root.state, "ownership membership"); + let connection = state + .connections + .get(&self.connection_id) + .filter(|connection| connection.generation == self.generation) + .ok_or_else(|| stale_generation("connection", &self.connection_id))?; + ensure_closing( + connection.phase, + format!("connection {}", self.connection_id), + )?; + state.connections.remove(&self.connection_id); self.completed = true; Ok(()) } @@ -1642,13 +1365,169 @@ impl Drop for ConnectionDisposal { fn drop(&mut self) { if !self.completed { tracing::warn!( - connection_id = %self.connection.inner.connection_id, + connection_id = %self.connection_id, "ERR_AGENTOS_CONNECTION_DISPOSAL_INCOMPLETE: connection remains Closing for a bounded retry" ); } } } +fn upgrade_root( + root: &Weak, + scope: &'static str, + id: &str, +) -> Result, OwnershipCoordinatorError> { + root.upgrade() + .ok_or_else(|| OwnershipCoordinatorError::NotFound { + scope, + id: id.to_owned(), + }) +} + +fn stale_generation(scope: &'static str, id: &str) -> OwnershipCoordinatorError { + OwnershipCoordinatorError::NotFound { + scope, + id: format!("{id} (stale generation)"), + } +} + +fn matching_connection<'a>( + state: &'a OwnershipCoordinatorState, + handle: &ConnectionCoordinator, +) -> Result<&'a ConnectionRecord, OwnershipCoordinatorError> { + state + .connections + .get(&handle.connection_id) + .filter(|connection| connection.generation == handle.generation) + .ok_or_else(|| stale_generation("connection", &handle.connection_id)) +} + +fn matching_connection_mut<'a>( + state: &'a mut OwnershipCoordinatorState, + handle: &ConnectionCoordinator, +) -> Result<&'a mut ConnectionRecord, OwnershipCoordinatorError> { + state + .connections + .get_mut(&handle.connection_id) + .filter(|connection| connection.generation == handle.generation) + .ok_or_else(|| stale_generation("connection", &handle.connection_id)) +} + +fn matching_session<'a>( + state: &'a OwnershipCoordinatorState, + handle: &SessionCoordinator, +) -> Result<&'a SessionRecord, OwnershipCoordinatorError> { + state + .connections + .get(&handle.connection_id) + .filter(|connection| connection.generation == handle.connection_generation) + .and_then(|connection| connection.sessions.get(&handle.session_id)) + .filter(|session| session.generation == handle.generation) + .ok_or_else(|| stale_generation("session", &handle.session_id)) +} + +fn matching_session_mut<'a>( + state: &'a mut OwnershipCoordinatorState, + handle: &SessionCoordinator, +) -> Result<&'a mut SessionRecord, OwnershipCoordinatorError> { + state + .connections + .get_mut(&handle.connection_id) + .filter(|connection| connection.generation == handle.connection_generation) + .and_then(|connection| connection.sessions.get_mut(&handle.session_id)) + .filter(|session| session.generation == handle.generation) + .ok_or_else(|| stale_generation("session", &handle.session_id)) +} + +#[cfg(test)] +fn matching_vm<'a>( + state: &'a OwnershipCoordinatorState, + handle: &VmCoordinator, +) -> Result<&'a VmRecord, OwnershipCoordinatorError> { + state + .connections + .get(&handle.connection_id) + .filter(|connection| connection.generation == handle.connection_generation) + .and_then(|connection| connection.sessions.get(&handle.session_id)) + .filter(|session| session.generation == handle.session_generation) + .and_then(|session| session.vms.get(&handle.vm_id)) + .filter(|vm| vm.generation == handle.generation && Arc::ptr_eq(&vm.gate, &handle.gate)) + .ok_or_else(|| stale_generation("VM", &handle.vm_id)) +} + +fn matching_vm_mut<'a>( + state: &'a mut OwnershipCoordinatorState, + handle: &VmCoordinator, +) -> Result<&'a mut VmRecord, OwnershipCoordinatorError> { + state + .connections + .get_mut(&handle.connection_id) + .filter(|connection| connection.generation == handle.connection_generation) + .and_then(|connection| connection.sessions.get_mut(&handle.session_id)) + .filter(|session| session.generation == handle.session_generation) + .and_then(|session| session.vms.get_mut(&handle.vm_id)) + .filter(|vm| vm.generation == handle.generation && Arc::ptr_eq(&vm.gate, &handle.gate)) + .ok_or_else(|| stale_generation("VM", &handle.vm_id)) +} + +fn validate_ownership_locked( + state: &OwnershipCoordinatorState, + ownership: &OwnershipScope, +) -> Result>, OwnershipCoordinatorError> { + let (connection_id, session_id, vm_id) = ownership_ids(ownership); + let connection = state.connections.get(connection_id).ok_or_else(|| { + OwnershipCoordinatorError::NotFound { + scope: "connection", + id: connection_id.to_owned(), + } + })?; + ensure_open(connection.phase, format!("connection {connection_id}"))?; + let Some(session_id) = session_id else { + return Ok(None); + }; + let session = + connection + .sessions + .get(session_id) + .ok_or_else(|| OwnershipCoordinatorError::NotFound { + scope: "session", + id: format!("{connection_id}:{session_id}"), + })?; + ensure_open( + session.phase, + format!("session {connection_id}:{session_id}"), + )?; + let Some(vm_id) = vm_id else { + return Ok(None); + }; + let vm = session + .vms + .get(vm_id) + .ok_or_else(|| OwnershipCoordinatorError::NotFound { + scope: "VM", + id: format!("{connection_id}:{session_id}:{vm_id}"), + })?; + ensure_open(vm.phase, format!("VM {connection_id}:{session_id}:{vm_id}"))?; + Ok(Some(Arc::clone(&vm.gate))) +} + +#[cfg(test)] +fn closed_entity_snapshot() -> EntityCoordinatorSnapshot { + EntityCoordinatorSnapshot { + phase: CoordinatorPhase::Closed, + active_operations: 0, + child_count: 0, + } +} + +fn ensure_closing(phase: CoordinatorPhase, scope: String) -> Result<(), OwnershipCoordinatorError> { + if phase == CoordinatorPhase::Closing { + Ok(()) + } else { + Err(OwnershipCoordinatorError::Closing { scope, phase }) + } +} + fn ownership_ids(ownership: &OwnershipScope) -> (&str, Option<&str>, Option<&str>) { match ownership { OwnershipScope::ConnectionOwnership(scope) => (&scope.connection_id, None, None), @@ -1663,43 +1542,26 @@ fn ownership_ids(ownership: &OwnershipScope) -> (&str, Option<&str>, Option<&str } } -fn validate_ordering_ownership( +fn validate_vm_concurrency_ownership( ownership: &OwnershipScope, - ordering: &RequestOrderingKey, + vm_concurrency: &VmConcurrencyClass, ) -> Result<(), OwnershipCoordinatorError> { - let (connection_id, session_id, vm_id) = ownership_ids(ownership); - let valid = match ordering { - RequestOrderingKey::Connection(key_connection) => key_connection == connection_id, - RequestOrderingKey::Session { - connection_id: key_connection, - session_id: key_session, - } => key_connection == connection_id && session_id.is_some_and(|id| id == key_session), - RequestOrderingKey::VmLifecycle { - connection_id: key_connection, - session_id: key_session, - vm_id: key_vm, + let valid = match vm_concurrency { + VmConcurrencyClass::OwnershipOnly => true, + VmConcurrencyClass::SharedVm | VmConcurrencyClass::ExclusiveVmLifecycle => { + matches!(ownership, OwnershipScope::VmOwnership(_)) } - | RequestOrderingKey::VmOperation { - connection_id: key_connection, - session_id: key_session, - vm_id: key_vm, - } => { - key_connection == connection_id - && session_id.is_some_and(|id| id == key_session) - && vm_id.is_some_and(|id| id == key_vm) - } - RequestOrderingKey::Extension { - connection_id: key_connection, - .. - } => key_connection == connection_id, - RequestOrderingKey::Unordered => true, }; if valid { Ok(()) } else { Err(OwnershipCoordinatorError::OwnershipMismatch { - expected: ownership_label(ownership), - actual: ordering_label(ordering), + expected: String::from("VM ownership"), + actual: format!( + "{} with {} VM concurrency", + ownership_label(ownership), + vm_concurrency_label(vm_concurrency) + ), }) } } @@ -1713,30 +1575,11 @@ fn ownership_label(ownership: &OwnershipScope) -> String { } } -fn ordering_label(ordering: &RequestOrderingKey) -> String { - match ordering { - RequestOrderingKey::Connection(connection) => connection.clone(), - RequestOrderingKey::Session { - connection_id, - session_id, - } => format!("{connection_id}/{session_id}"), - RequestOrderingKey::VmLifecycle { - connection_id, - session_id, - vm_id, - } - | RequestOrderingKey::VmOperation { - connection_id, - session_id, - vm_id, - } => format!("{connection_id}/{session_id}/{vm_id}"), - RequestOrderingKey::Extension { - namespace, - connection_id, - key, - .. - } => format!("{connection_id}/extension/{namespace}/{}-bytes", key.len()), - RequestOrderingKey::Unordered => String::from("unordered"), +fn vm_concurrency_label(vm_concurrency: &VmConcurrencyClass) -> String { + match vm_concurrency { + VmConcurrencyClass::OwnershipOnly => String::from("ownership-only"), + VmConcurrencyClass::SharedVm => String::from("shared-vm"), + VmConcurrencyClass::ExclusiveVmLifecycle => String::from("exclusive-vm-lifecycle"), } } @@ -1896,11 +1739,7 @@ mod tests { RequestOperationMetadata::new( OwnershipScope::vm(connection, session, vm), "VM operation", - RequestOrderingKey::VmOperation { - connection_id: connection.to_owned(), - session_id: session.to_owned(), - vm_id: vm.to_owned(), - }, + VmConcurrencyClass::SharedVm, ) } @@ -1908,24 +1747,7 @@ mod tests { RequestOperationMetadata::new( OwnershipScope::vm(connection, session, vm), "VM lifecycle", - RequestOrderingKey::VmLifecycle { - connection_id: connection.to_owned(), - session_id: session.to_owned(), - vm_id: vm.to_owned(), - }, - ) - } - - fn extension_metadata(connection: &str, key: &[u8]) -> RequestOperationMetadata { - RequestOperationMetadata::new( - OwnershipScope::connection(connection), - "extension operation", - RequestOrderingKey::Extension { - namespace: String::from("test.extension"), - connection_id: connection.to_owned(), - key: key.to_vec(), - policy: ExtensionOrderingPolicy::CoreExclusive, - }, + VmConcurrencyClass::ExclusiveVmLifecycle, ) } @@ -1954,6 +1776,23 @@ mod tests { drop(gate_a); } + #[tokio::test(flavor = "current_thread")] + async fn ordinary_operations_share_one_vm_without_serializing() { + let (coordinator, _, _, vm_a, _) = configured(); + let metadata = vm_metadata("connection-a", "session-a", "vm-a"); + let first = coordinator + .admit(&metadata, OperationCancellation::new()) + .await + .expect("first ordinary VM operation starts"); + let second = coordinator + .admit(&metadata, OperationCancellation::new()) + .await + .expect("second ordinary VM operation starts concurrently"); + assert_eq!(vm_a.snapshot().active_operations, 2); + drop(second); + drop(first); + } + #[tokio::test(flavor = "current_thread")] async fn lifecycle_waits_for_same_vm_operation_and_excludes_new_work() { let (coordinator, _, _, vm_a, _) = configured(); @@ -2064,7 +1903,19 @@ mod tests { let lifecycle = lifecycle .await .expect("lifecycle starts after internal event registration drains"); + let mut internal_during_active = coordinator + .admit_internal_vm_event(&metadata, OperationCancellation::new()) + .map(expect_internal_deferred) + .expect("internal progress remains durably deferred while lifecycle is active"); + assert!(!internal_during_active + .try_activate_deferred_internal_event() + .expect("active lifecycle remains deferred")); drop(lifecycle); + assert!(internal_during_active + .try_activate_deferred_internal_event() + .expect("deferred progress activates after lifecycle")); + assert_eq!(vm.snapshot().active_operations, 1); + drop(internal_during_active); let disposal_cancellation = OperationCancellation::new(); let disposal_tracked = coordinator @@ -2157,147 +2008,208 @@ mod tests { } #[tokio::test(flavor = "current_thread")] - async fn session_disposal_closes_admission_cancels_owned_work_and_drains() { - let (coordinator, connection, session, vm_a, _) = configured(); - let cancellation = OperationCancellation::new(); - let operation = coordinator - .admit( - &vm_metadata("connection-a", "session-a", "vm-a"), - cancellation.clone(), - ) - .await - .expect("owned operation starts"); - let disposal = coordinator - .begin_session_disposal( - &OwnershipScope::session("connection-a", "session-a"), - OperationCancellationReason::ConnectionClosed, - ) - .expect("session enters Closing"); - assert_eq!(session.snapshot().phase, CoordinatorPhase::Closing); - assert_eq!(vm_a.snapshot().phase, CoordinatorPhase::Closing); - assert_eq!( - cancellation.reason(), - Some(OperationCancellationReason::ConnectionClosed) - ); - - let rejected = coordinator - .admit( - &vm_metadata("connection-a", "session-a", "vm-a"), - OperationCancellation::new(), - ) - .await - .expect_err("closing session rejects new owned operations"); - assert_eq!(rejected.code(), "ERR_AGENTOS_COORDINATOR_CLOSING"); - - let mut drained = Box::pin(disposal.wait_drained()); - let waker = std::task::Waker::noop(); - let mut cx = Context::from_waker(waker); - assert!(matches!(drained.as_mut().poll(&mut cx), Poll::Pending)); - drop(operation); - drained.as_mut().await; - drop(drained); - disposal.complete().expect("complete drained disposal"); - assert!(connection.session("session-a").is_err()); - } - - #[tokio::test(flavor = "current_thread")] - async fn cross_connection_ordering_key_cannot_address_owned_vm() { + async fn vm_concurrency_requires_vm_ownership() { let (coordinator, _, _, _, _) = configured(); - let connection_b = coordinator - .register_connection("connection-b") - .expect("register connection B"); - let session_b = connection_b.open_session("session-a").expect("session B"); - session_b.open_vm("vm-a").expect("VM B"); - let metadata = RequestOperationMetadata::new( - OwnershipScope::vm("connection-b", "session-a", "vm-a"), - "forged operation", - RequestOrderingKey::VmOperation { - connection_id: String::from("connection-a"), - session_id: String::from("session-a"), - vm_id: String::from("vm-a"), - }, + OwnershipScope::connection("connection-a"), + "invalid shared VM operation", + VmConcurrencyClass::SharedVm, ); let error = coordinator .admit(&metadata, OperationCancellation::new()) .await - .expect_err("ordering key cannot cross connection ownership"); + .expect_err("VM concurrency requires VM ownership"); assert_eq!(error.code(), "ERR_AGENTOS_COORDINATOR_OWNERSHIP"); - assert_eq!(connection_b.snapshot().active_operations, 0); } #[tokio::test(flavor = "current_thread")] - async fn opaque_extension_ordering_keys_reject_only_same_connection_conflicts() { - let (coordinator, _, _, _, _) = configured(); - coordinator - .register_connection("connection-b") - .expect("register connection B"); - + async fn ownership_only_policy_tracks_scope_without_serializing_requests() { + let (coordinator, connection, session, _, _) = configured(); + let metadata = RequestOperationMetadata::new( + OwnershipScope::session("connection-a", "session-a"), + "session-owned work", + VmConcurrencyClass::OwnershipOnly, + ); let first = coordinator - .admit( - &extension_metadata("connection-a", b"same-key"), - OperationCancellation::new(), - ) + .admit(&metadata, OperationCancellation::new()) .await - .expect("first keyed operation starts"); - let same_key = coordinator - .admit( - &extension_metadata("connection-a", b"same-key"), - OperationCancellation::new(), - ) + .expect("first ownership-only request starts"); + let second = coordinator + .admit(&metadata, OperationCancellation::new()) .await - .expect_err("same connection and key conflict"); - assert_eq!(same_key.code(), "ERR_AGENTOS_ORDERING_CONFLICT"); + .expect("second ownership-only request starts concurrently"); + assert_eq!(connection.snapshot().active_operations, 0); + assert_eq!(session.snapshot().active_operations, 0); + drop(second); + drop(first); + } - let different_key = coordinator + #[tokio::test(flavor = "current_thread")] + async fn lifecycle_conflicts_and_closing_never_reopen_admission() { + let (coordinator, _, _, vm, _) = configured(); + let ordinary = coordinator .admit( - &extension_metadata("connection-a", b"different-key"), + &vm_metadata("connection-a", "session-a", "vm-a"), OperationCancellation::new(), ) .await - .expect("different key progresses concurrently"); - let different_connection = coordinator + .expect("ordinary operation starts"); + let pending = vm + .begin_lifecycle(OperationCancellation::new()) + .expect("first lifecycle becomes pending"); + let second = vm + .begin_lifecycle(OperationCancellation::new()) + .expect_err("second pending lifecycle is rejected"); + assert_eq!(second.code(), "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT"); + drop(pending); + assert_eq!(vm.snapshot().lifecycle, VmLifecyclePhase::Idle); + + let pending = vm + .begin_lifecycle(OperationCancellation::new()) + .expect("lifecycle can retry after cancellation"); + drop(ordinary); + let active = pending.wait().await.expect("lifecycle activates"); + let second = vm + .begin_lifecycle(OperationCancellation::new()) + .expect_err("second active lifecycle is rejected"); + assert_eq!(second.code(), "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT"); + + let disposal = coordinator + .begin_vm_disposal( + &OwnershipScope::vm("connection-a", "session-a", "vm-a"), + OperationCancellationReason::Explicit, + ) + .expect("closing begins while lifecycle is active"); + drop(active); + let rejected = coordinator .admit( - &extension_metadata("connection-b", b"same-key"), + &vm_metadata("connection-a", "session-a", "vm-a"), OperationCancellation::new(), ) .await - .expect("same key is isolated by connection"); + .expect_err("active lifecycle drop cannot reopen a closing VM"); + assert_eq!(rejected.code(), "ERR_AGENTOS_COORDINATOR_CLOSING"); + disposal.wait_drained().await; + disposal.complete().expect("closed gate drains"); + } - drop(different_connection); - drop(different_key); - drop(first); - coordinator - .admit( - &extension_metadata("connection-a", b"same-key"), - OperationCancellation::new(), + #[tokio::test(flavor = "current_thread")] + async fn stale_vm_generation_cannot_mutate_recreated_membership() { + let (coordinator, _, session, stale_vm, _) = configured(); + let disposal = coordinator + .begin_vm_disposal( + &OwnershipScope::vm("connection-a", "session-a", "vm-a"), + OperationCancellationReason::Explicit, ) - .await - .expect("key is released with terminal ownership permit"); + .expect("dispose first VM generation"); + disposal.wait_drained().await; + disposal.complete().expect("remove first VM generation"); + + let current_vm = session.open_vm("vm-a").expect("recreate textual VM id"); + assert_ne!(stale_vm.generation, current_vm.generation); + let stale_error = stale_vm + .begin_lifecycle(OperationCancellation::new()) + .expect_err("stale handle cannot mutate the replacement VM"); + assert_eq!(stale_error.code(), "ERR_AGENTOS_COORDINATOR_NOT_FOUND"); + assert_eq!(current_vm.snapshot().lifecycle, VmLifecyclePhase::Idle); + + // Even a stale gate generation token cannot release a later lifecycle. + let current = current_vm + .begin_lifecycle(OperationCancellation::new()) + .expect("current lifecycle begins"); + let current_id = match lock(¤t_vm.gate.state, "test VM gate").lifecycle { + VmLifecycleState::Pending { id, .. } => id, + _ => panic!("current lifecycle must be pending"), + }; + release_lifecycle(¤t_vm.gate, current_id.wrapping_sub(1)); + assert_eq!(current_vm.snapshot().lifecycle, VmLifecyclePhase::Pending); + drop(current); + assert_eq!(current_vm.snapshot().lifecycle, VmLifecyclePhase::Idle); } #[tokio::test(flavor = "current_thread")] - async fn opaque_extension_key_cannot_cross_connection_ownership() { - let (coordinator, _, _, _, _) = configured(); - let metadata = RequestOperationMetadata::new( - OwnershipScope::connection("connection-a"), - "forged extension operation", - RequestOrderingKey::Extension { - namespace: String::from("test.extension"), - connection_id: String::from("connection-b"), - key: b"key".to_vec(), - policy: ExtensionOrderingPolicy::CoreExclusive, - }, - ); - let error = coordinator - .admit(&metadata, OperationCancellation::new()) - .await - .expect_err("extension ordering key cannot cross connection ownership"); - assert_eq!(error.code(), "ERR_AGENTOS_COORDINATOR_OWNERSHIP"); + async fn deterministic_gate_model_matches_randomized_admit_drop_and_cancel_sequence() { + let (coordinator, _, _, vm, _) = configured(); + let metadata = vm_metadata("connection-a", "session-a", "vm-a"); + let mut ordinary = Vec::::new(); + let mut internal = Vec::::new(); + let mut pending = None::; + let mut active = None::; + let mut seed = 0x4d59_5df4_d0f3_3173_u64; + + for _ in 0..512 { + seed = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(1); + match (seed >> 32) % 8 { + 0 if pending.is_none() && active.is_none() => { + match vm.begin_lifecycle(OperationCancellation::new()) { + Ok(admission) => pending = Some(admission), + Err(error) => assert_eq!(error.code(), "ERR_AGENTOS_COORDINATOR_CLOSING"), + } + } + 1 if pending.is_some() => drop(pending.take()), + 2 if active.is_some() => drop(active.take()), + 3 if !ordinary.is_empty() => { + let index = seed as usize % ordinary.len(); + ordinary.swap_remove(index); + } + 4 if !internal.is_empty() => { + let index = seed as usize % internal.len(); + internal.swap_remove(index); + } + 5 => { + if let Ok(permit) = coordinator + .admit(&metadata, OperationCancellation::new()) + .await + { + ordinary.push(permit); + } + } + _ => { + if let Ok(admission) = + coordinator.admit_internal_vm_event(&metadata, OperationCancellation::new()) + { + match admission { + InternalVmEventAdmission::Admitted(permit) + | InternalVmEventAdmission::Deferred(permit) => internal.push(permit), + } + } + } + } + + if pending.is_some() && ordinary.is_empty() && internal.is_empty() { + active = Some( + pending + .take() + .expect("pending lifecycle") + .wait() + .await + .expect("drained lifecycle activates"), + ); + } + let snapshot = vm.snapshot(); + assert_eq!(snapshot.active_operations, ordinary.len() + internal.len()); + assert_eq!( + snapshot.lifecycle, + if active.is_some() { + VmLifecyclePhase::Active + } else if pending.is_some() { + VmLifecyclePhase::Pending + } else { + VmLifecyclePhase::Idle + } + ); + } + + drop(pending); + drop(active); + drop(ordinary); + drop(internal); + assert_eq!(vm.snapshot().active_operations, 0); + assert_eq!(vm.snapshot().lifecycle, VmLifecyclePhase::Idle); } #[test] - fn coordinator_membership_and_operation_state_are_bounded() { + fn coordinator_membership_state_is_bounded() { let coordinator = OwnershipCoordinator::new(OwnershipCoordinatorLimits { max_connections: 1, max_sessions_per_connection: 1, @@ -2314,15 +2226,6 @@ mod tests { assert_eq!(error.code(), "ERR_AGENTOS_COORDINATOR_LIMIT"); assert!(error.to_string().contains(CONNECTION_LIMIT_PATH)); - let active = connection - .register_operation(OperationCancellation::new()) - .expect("first connection operation"); - let error = connection - .register_operation(OperationCancellation::new()) - .expect_err("per-entity operation bound"); - assert!(error.to_string().contains(IN_FLIGHT_REQUEST_COUNT_PATH)); - drop(active); - let session = connection.open_session("session-a").expect("first session"); let error = connection .open_session("session-b") diff --git a/crates/native-sidecar/src/request_operations.rs b/crates/native-sidecar/src/request_operations.rs index 4cb48ef51b..1243912ff0 100644 --- a/crates/native-sidecar/src/request_operations.rs +++ b/crates/native-sidecar/src/request_operations.rs @@ -12,8 +12,6 @@ use std::sync::atomic::{AtomicU8, Ordering}; use std::sync::{Arc, Mutex, MutexGuard}; use tokio::sync::Notify; -use crate::extension::ExtensionOrderingPolicy; - pub(crate) const IN_FLIGHT_REQUEST_COUNT_PATH: &str = "runtime.protocol.maxInFlightRequests"; pub(crate) const IN_FLIGHT_REQUEST_BYTES_PATH: &str = "runtime.protocol.maxInFlightRequestBytes"; @@ -47,66 +45,53 @@ impl RequestOperationKey { } } -/// Narrow conflict domains carried with an independently executing request. -/// This is metadata only: entity coordinators enforce the ordering policy. +/// VM concurrency class for an independently executing request. +/// +/// Entity identity is deliberately absent: [`RequestOperationMetadata::ownership`] +/// is the sole source of connection/session/VM ownership. The coordinator uses +/// ownership for cancellation, disposal, limits, and active-operation tracking, +/// then applies this class only to the narrow VM exclusion that remains. #[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum RequestOrderingKey { - Connection(String), - Session { - connection_id: String, - session_id: String, - }, - VmLifecycle { - connection_id: String, - session_id: String, - vm_id: String, - }, - VmOperation { - connection_id: String, - session_id: String, - vm_id: String, - }, - Extension { - namespace: String, - connection_id: String, - key: Vec, - policy: ExtensionOrderingPolicy, - }, - Unordered, +pub(crate) enum VmConcurrencyClass { + /// Track the operation at its ownership scopes without adding a conflict + /// domain. Connection- and session-owned work uses this class and remains + /// concurrent with other admitted work at the same scope. + OwnershipOnly, + /// Ordinary VM work shares the VM with other ordinary operations. Its + /// registration is counted by the per-VM lifecycle gate so exclusive work + /// can wait for a precise, cancellation-aware drain. + SharedVm, + /// VM lifecycle work atomically closes ordinary admission, waits for work + /// already registered under `SharedVm`, and then runs alone. New shared + /// requests are rejected rather than queued while the gate is pending or + /// active. + ExclusiveVmLifecycle, } +/// Complete admission description for one independently executing request. +/// Ownership says which entity path the operation uses; the VM concurrency +/// class says whether that work enters the lifecycle gate. #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct RequestOperationMetadata { pub(crate) ownership: OwnershipScope, pub(crate) operation: String, - pub(crate) ordering_key: RequestOrderingKey, + pub(crate) vm_concurrency: VmConcurrencyClass, } impl RequestOperationMetadata { pub(crate) fn new( ownership: OwnershipScope, operation: impl Into, - ordering_key: RequestOrderingKey, + vm_concurrency: VmConcurrencyClass, ) -> Self { Self { ownership, operation: operation.into(), - ordering_key, + vm_concurrency, } } } -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -pub(crate) enum RequestOperationState { - Admitted, - Running, - Cancelling, - Completing, - Failed, - Shutdown, - Terminal, -} - #[derive(Clone, Copy, Debug, PartialEq, Eq)] #[repr(u8)] pub(crate) enum OperationCancellationReason { @@ -188,11 +173,11 @@ const PUBLICATION_RETAINED: u8 = 2; const PUBLICATION_TAKEN_OVER: u8 = 3; #[derive(Clone, Debug)] -pub(crate) struct TerminalResponseGuard { +pub(crate) struct ResponsePublicationGuard { state: Arc, } -impl TerminalResponseGuard { +impl ResponsePublicationGuard { fn new() -> Self { Self { state: Arc::new(AtomicU8::new(PUBLICATION_UNCLAIMED)), @@ -252,13 +237,12 @@ pub(crate) struct RequestOperationSnapshot { pub(crate) key: RequestOperationKey, pub(crate) metadata: RequestOperationMetadata, pub(crate) request_bytes: usize, - pub(crate) state: RequestOperationState, pub(crate) cancellation_reason: Option, pub(crate) terminal_claimed: bool, } #[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct RequestOperationRegistrySnapshot { +pub(crate) struct OperationTableSnapshot { pub(crate) in_flight_requests: usize, pub(crate) in_flight_request_bytes: usize, pub(crate) closed: Option, @@ -272,38 +256,61 @@ pub(crate) struct ForcedRequestOutcome { pub(crate) ownership: OwnershipScope, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum OperationClass { + Ordinary, + Progress, +} + #[derive(Debug)] struct OperationRecord { + class: OperationClass, generation: u64, metadata: RequestOperationMetadata, request_bytes: usize, - state: RequestOperationState, cancellation: OperationCancellation, - terminal: TerminalResponseGuard, + publication: ResponsePublicationGuard, } #[derive(Debug)] struct RegistryState { operations: BTreeMap, in_flight_request_bytes: usize, + in_flight_progress_bytes: usize, next_generation: u64, closed: Option, + progress_closed: Option, closed_connections: BTreeMap, + closed_scopes: Vec<(OwnershipScope, OperationCancellationReason)>, } #[derive(Debug)] -struct RequestOperationRegistryInner { +struct OperationTableInner { limits: RequestOperationLimits, + progress_limits: ProgressRequestLimits, state: Mutex, + changed: Notify, } #[derive(Clone, Debug)] -pub(crate) struct RequestOperationRegistry { - inner: Arc, +pub(crate) struct OperationTable { + inner: Arc, } -impl RequestOperationRegistry { +impl OperationTable { + #[cfg(test)] pub(crate) fn new(limits: RequestOperationLimits) -> Self { + let progress_limits = ProgressRequestLimits { + max_requests: limits.max_requests, + max_request_bytes: limits.max_request_bytes, + }; + Self::new_with_limits(limits, progress_limits) + } + + fn new_with_limits( + limits: RequestOperationLimits, + progress_limits: ProgressRequestLimits, + ) -> Self { assert!( limits.max_requests > 0, "request count limit must be positive" @@ -312,22 +319,47 @@ impl RequestOperationRegistry { limits.max_request_bytes > 0, "request byte limit must be positive" ); + assert!( + progress_limits.max_requests > 0, + "progress request count limit must be positive" + ); + assert!( + progress_limits.max_request_bytes > 0, + "progress request byte limit must be positive" + ); Self { - inner: Arc::new(RequestOperationRegistryInner { + inner: Arc::new(OperationTableInner { limits, + progress_limits, state: Mutex::new(RegistryState { operations: BTreeMap::new(), in_flight_request_bytes: 0, + in_flight_progress_bytes: 0, next_generation: 1, closed: None, + progress_closed: None, closed_connections: BTreeMap::new(), + closed_scopes: Vec::new(), }), + changed: Notify::new(), }), } } pub(crate) fn from_protocol_config(config: &RuntimeProtocolConfig) -> Self { - Self::new(RequestOperationLimits::from(config)) + Self::new_with_limits( + RequestOperationLimits::from(config), + ProgressRequestLimits::from(config), + ) + } + + /// Return the progress-admission projection over this exact operation + /// table. Both classes share request identity and publication ownership, + /// while retaining independently configured count and byte budgets. + pub(crate) fn progress_requests(&self) -> ProgressOperationView { + ProgressOperationView { + inner: Arc::clone(&self.inner), + } } pub(crate) fn admit( @@ -346,17 +378,17 @@ impl RequestOperationRegistry { let generation = state.next_generation; state.next_generation = state.next_generation.wrapping_add(1).max(1); let cancellation = OperationCancellation::new(); - let terminal = TerminalResponseGuard::new(); + let publication = ResponsePublicationGuard::new(); state.in_flight_request_bytes = requested_total; state.operations.insert( key.clone(), OperationRecord { + class: OperationClass::Ordinary, generation, metadata: metadata.clone(), request_bytes, - state: RequestOperationState::Admitted, cancellation: cancellation.clone(), - terminal: terminal.clone(), + publication: publication.clone(), }, ); drop(state); @@ -368,14 +400,14 @@ impl RequestOperationRegistry { metadata, request_bytes, cancellation, - terminal, + publication, released: false, }) } - /// Check admission before acquiring a terminal-output reservation. The - /// router is the sole admission producer, but [`Self::admit`] repeats this - /// check so this preflight can never weaken the registry's bounds. + /// Test-only admission probe. Production routing uses [`Self::admit`] once + /// and releases the resulting operation if output reservation fails. + #[cfg(test)] pub(crate) fn check_admission( &self, key: &RequestOperationKey, @@ -409,12 +441,27 @@ impl RequestOperationRegistry { reason, }); } + if let Some((scope, reason)) = state + .closed_scopes + .iter() + .find(|(scope, _)| scope_contains(scope, &metadata.ownership)) + { + return Err(RequestAdmissionError::OwnershipClosed { + scope: ownership_label(scope), + reason: *reason, + }); + } if state.operations.contains_key(&key) { return Err(RequestAdmissionError::DuplicateRequest { key: key.clone() }); } - if state.operations.len() >= self.inner.limits.max_requests { + let ordinary_count = state + .operations + .values() + .filter(|record| record.class == OperationClass::Ordinary) + .count(); + if ordinary_count >= self.inner.limits.max_requests { return Err(RequestAdmissionError::CountLimit { - current: state.operations.len(), + current: ordinary_count, requested: 1, limit: self.inner.limits.max_requests, }); @@ -437,6 +484,7 @@ impl RequestOperationRegistry { Ok(()) } + #[cfg(test)] pub(crate) fn cancel( &self, key: &RequestOperationKey, @@ -446,12 +494,14 @@ impl RequestOperationRegistry { let Some(record) = state.operations.get_mut(key) else { return CancelOperationResult::NotFound; }; - if record.state == RequestOperationState::Terminal { + if record.class != OperationClass::Ordinary { + return CancelOperationResult::NotFound; + } + if record.publication.is_claimed() { return CancelOperationResult::AlreadyTerminal; } let signalled = record.cancellation.signal(reason); if signalled { - advance_cancellation_state(record, reason); CancelOperationResult::Signalled } else { CancelOperationResult::AlreadySignalled( @@ -463,6 +513,7 @@ impl RequestOperationRegistry { } } + #[cfg(test)] pub(crate) fn close_connection( &self, connection_id: &str, @@ -476,10 +527,10 @@ impl RequestOperationRegistry { let mut signalled = 0; for (key, record) in &mut state.operations { if key.connection_id == connection_id - && record.state != RequestOperationState::Terminal + && record.class == OperationClass::Ordinary + && !record.publication.is_claimed() && record.cancellation.signal(reason) { - advance_cancellation_state(record, reason); signalled += 1; } } @@ -491,15 +542,57 @@ impl RequestOperationRegistry { state.closed.get_or_insert(reason); let mut signalled = 0; for record in state.operations.values_mut() { - if record.state != RequestOperationState::Terminal && record.cancellation.signal(reason) + if record.class == OperationClass::Ordinary + && !record.publication.is_claimed() + && record.cancellation.signal(reason) { - advance_cancellation_state(record, reason); signalled += 1; } } signalled } + /// Atomically close one ownership subtree and signal every operation in it + /// except the lifecycle request performing the close. Admission and scope + /// closure use the same short mutex, so no operation can slip between the + /// closed marker and the bounded-table scan. + pub(crate) fn close_scope( + &self, + scope: OwnershipScope, + reason: OperationCancellationReason, + excluded: Option<&RequestOperationKey>, + ) -> ScopeOperationDrain { + let mut state = self.lock_state(); + if !state + .closed_scopes + .iter() + .any(|(closed, _)| closed == &scope) + { + state.closed_scopes.push((scope.clone(), reason)); + } + for (key, record) in &state.operations { + if excluded != Some(key) && scope_contains(&scope, &record.metadata.ownership) { + record.cancellation.signal(reason); + } + } + drop(state); + self.inner.changed.notify_waiters(); + ScopeOperationDrain { + registry: self.clone(), + scope, + excluded: excluded.cloned(), + } + } + + /// Reopen an exact textual scope only after its old generation completed + /// and membership was explicitly recreated. Ancestor closures still win. + pub(crate) fn reopen_scope(&self, scope: &OwnershipScope) { + let mut state = self.lock_state(); + state.closed_scopes.retain(|(closed, _)| closed != scope); + drop(state); + self.inner.changed.notify_waiters(); + } + /// Claim terminal ownership for every unfinished request and remove all /// registry accounting in one critical section. The supervisor calls this /// only after the cooperative drain deadline; task-local operation handles @@ -511,18 +604,23 @@ impl RequestOperationRegistry { ) -> Vec { let mut state = self.lock_state(); state.closed.get_or_insert(reason); - let operations = std::mem::take(&mut state.operations); + let keys = state + .operations + .iter() + .filter(|(_, record)| record.class == OperationClass::Ordinary) + .map(|(key, _)| key.clone()) + .collect::>(); state.in_flight_request_bytes = 0; - operations + let outcomes = keys .into_iter() - .filter_map(|(key, mut record)| { + .filter_map(|key| state.operations.remove(&key).map(|record| (key, record))) + .filter_map(|(key, record)| { record.cancellation.signal(reason); - advance_cancellation_state(&mut record, reason); // A claimed-but-still-registered record may be draining its // finite ordinary event batch. Its terminal response already // exists, so removing accounting must not synthesize a second // terminal during shutdown. - if !record.terminal.try_take_over_unretained() { + if !record.publication.try_take_over_unretained() { return None; } Some(ForcedRequestOutcome { @@ -530,98 +628,60 @@ impl RequestOperationRegistry { ownership: record.metadata.ownership, }) }) - .collect() + .collect::>(); + drop(state); + self.inner.changed.notify_waiters(); + outcomes } - pub(crate) fn snapshot(&self) -> RequestOperationRegistrySnapshot { + pub(crate) fn snapshot(&self) -> OperationTableSnapshot { let state = self.lock_state(); - RequestOperationRegistrySnapshot { - in_flight_requests: state.operations.len(), + OperationTableSnapshot { + in_flight_requests: state + .operations + .values() + .filter(|record| record.class == OperationClass::Ordinary) + .count(), in_flight_request_bytes: state.in_flight_request_bytes, closed: state.closed, closed_connections: state.closed_connections.keys().cloned().collect(), operations: state .operations .iter() + .filter(|(_, record)| record.class == OperationClass::Ordinary) .map(|(key, record)| RequestOperationSnapshot { key: key.clone(), metadata: record.metadata.clone(), request_bytes: record.request_bytes, - state: record.state, cancellation_reason: record.cancellation.reason(), - terminal_claimed: record.terminal.is_claimed(), + terminal_claimed: record.publication.is_claimed(), }) .collect(), } } - fn transition( - &self, - key: &RequestOperationKey, - generation: u64, - next: RequestOperationState, - ) -> Result<(), OperationTransitionError> { + fn try_mark_terminal(&self, key: &RequestOperationKey, generation: u64) -> bool { let mut state = self.lock_state(); let Some(record) = state.operations.get_mut(key) else { - return Err(OperationTransitionError::NotFound(key.clone())); + return false; }; - if record.generation != generation { - return Err(OperationTransitionError::StaleGeneration(key.clone())); - } - if !valid_transition(record.state, next) { - return Err(OperationTransitionError::Invalid { - key: key.clone(), - current: record.state, - next, - }); + if record.class != OperationClass::Ordinary { + return false; } - record.state = next; - Ok(()) - } - - fn try_mark_terminal( - &self, - key: &RequestOperationKey, - generation: u64, - ) -> Result { - let mut state = self.lock_state(); - let Some(record) = state.operations.get_mut(key) else { - return Err(OperationTransitionError::NotFound(key.clone())); - }; if record.generation != generation { - return Err(OperationTransitionError::StaleGeneration(key.clone())); + return false; } - if record.state == RequestOperationState::Terminal { - debug_assert!(record.terminal.is_claimed()); - return Ok(false); + if !record.publication.try_claim() { + return false; } - if !valid_transition(record.state, RequestOperationState::Terminal) { - return Err(OperationTransitionError::Invalid { - key: key.clone(), - current: record.state, - next: RequestOperationState::Terminal, - }); - } - if !record.terminal.try_claim() { - // All claims go through this method, so a claimed guard and a - // terminal state are updated in the same registry critical section. - tracing::error!( - request_id = key.request_id, - connection_id = %key.connection_id, - "ERR_AGENTOS_REQUEST_TERMINAL_STATE: terminal response was claimed without a terminal registry state" - ); - return Ok(false); - } - record.state = RequestOperationState::Terminal; - Ok(true) + true } fn release(&self, key: &RequestOperationKey, generation: u64, request_bytes: usize) { let mut state = self.lock_state(); - let should_remove = state - .operations - .get(key) - .is_some_and(|record| record.generation == generation); + let should_remove = state.operations.get(key).is_some_and(|record| { + record.class == OperationClass::Ordinary && record.generation == generation + }); if should_remove { state.operations.remove(key); state.in_flight_request_bytes = state @@ -637,6 +697,7 @@ impl RequestOperationRegistry { ); 0 }); + self.inner.changed.notify_waiters(); } } @@ -650,6 +711,40 @@ impl RequestOperationRegistry { } } +/// Read-only drain view for one closed ownership subtree. It never retains the +/// registry lock across `.await`; `Notify` is only a wake hint and every wake +/// rechecks the bounded authoritative table. +#[derive(Clone, Debug)] +pub(crate) struct ScopeOperationDrain { + registry: OperationTable, + scope: OwnershipScope, + excluded: Option, +} + +impl ScopeOperationDrain { + pub(crate) async fn wait_drained(&self) { + loop { + let changed = self.registry.inner.changed.notified(); + if self.active_operations() == 0 { + return; + } + changed.await; + } + } + + pub(crate) fn active_operations(&self) -> usize { + self.registry + .lock_state() + .operations + .iter() + .filter(|(key, record)| { + self.excluded.as_ref() != Some(key) + && scope_contains(&self.scope, &record.metadata.ownership) + }) + .count() + } +} + fn ownership_connection_id(ownership: &OwnershipScope) -> &str { match ownership { OwnershipScope::ConnectionOwnership(scope) => &scope.connection_id, @@ -658,53 +753,57 @@ fn ownership_connection_id(ownership: &OwnershipScope) -> &str { } } -fn cancellation_state(reason: OperationCancellationReason) -> RequestOperationState { - match reason { - OperationCancellationReason::Explicit => RequestOperationState::Cancelling, - OperationCancellationReason::ConnectionClosed - | OperationCancellationReason::Shutdown - | OperationCancellationReason::TransportClosed => RequestOperationState::Shutdown, - } -} - -fn advance_cancellation_state(record: &mut OperationRecord, reason: OperationCancellationReason) { - let next = cancellation_state(reason); - if valid_transition(record.state, next) { - record.state = next; +fn ownership_label(ownership: &OwnershipScope) -> String { + match ownership { + OwnershipScope::ConnectionOwnership(scope) => scope.connection_id.clone(), + OwnershipScope::SessionOwnership(scope) => { + format!("{}/{}", scope.connection_id, scope.session_id) + } + OwnershipScope::VmOwnership(scope) => { + format!( + "{}/{}/{}", + scope.connection_id, scope.session_id, scope.vm_id + ) + } } } -fn valid_transition(current: RequestOperationState, next: RequestOperationState) -> bool { - use RequestOperationState as State; - matches!( - (current, next), +fn scope_contains(parent: &OwnershipScope, child: &OwnershipScope) -> bool { + match (parent, child) { ( - State::Admitted, - State::Running | State::Cancelling | State::Failed | State::Shutdown - ) | ( - State::Running, - State::Cancelling | State::Completing | State::Failed | State::Shutdown - ) | ( - State::Cancelling, - State::Completing | State::Failed | State::Shutdown - ) | (State::Completing, State::Failed | State::Terminal) - | (State::Failed, State::Terminal) - | ( - State::Shutdown, - State::Completing | State::Failed | State::Terminal - ) - ) + OwnershipScope::ConnectionOwnership(parent), + OwnershipScope::ConnectionOwnership(child), + ) => parent.connection_id == child.connection_id, + (OwnershipScope::ConnectionOwnership(parent), OwnershipScope::SessionOwnership(child)) => { + parent.connection_id == child.connection_id + } + (OwnershipScope::ConnectionOwnership(parent), OwnershipScope::VmOwnership(child)) => { + parent.connection_id == child.connection_id + } + (OwnershipScope::SessionOwnership(parent), OwnershipScope::SessionOwnership(child)) => { + parent.connection_id == child.connection_id && parent.session_id == child.session_id + } + (OwnershipScope::SessionOwnership(parent), OwnershipScope::VmOwnership(child)) => { + parent.connection_id == child.connection_id && parent.session_id == child.session_id + } + (OwnershipScope::VmOwnership(parent), OwnershipScope::VmOwnership(child)) => { + parent.connection_id == child.connection_id + && parent.session_id == child.session_id + && parent.vm_id == child.vm_id + } + _ => false, + } } #[derive(Debug)] pub(crate) struct RequestOperation { - registry: RequestOperationRegistry, + registry: OperationTable, key: RequestOperationKey, generation: u64, metadata: RequestOperationMetadata, request_bytes: usize, cancellation: OperationCancellation, - terminal: TerminalResponseGuard, + publication: ResponsePublicationGuard, released: bool, } @@ -717,6 +816,7 @@ impl RequestOperation { &self.metadata } + #[cfg(test)] pub(crate) fn request_bytes(&self) -> usize { self.request_bytes } @@ -725,19 +825,16 @@ impl RequestOperation { self.cancellation.clone() } - pub(crate) fn transition( - &self, - next: RequestOperationState, - ) -> Result<(), OperationTransitionError> { - self.registry.transition(&self.key, self.generation, next) + pub(crate) fn reopen_scope(&self, scope: &OwnershipScope) { + self.registry.reopen_scope(scope); } - pub(crate) fn try_mark_terminal(&self) -> Result { + pub(crate) fn try_mark_terminal(&self) -> bool { self.registry.try_mark_terminal(&self.key, self.generation) } pub(crate) fn mark_terminal_retained(&self) -> bool { - self.terminal.mark_retained() + self.publication.mark_retained() } pub(crate) fn release(mut self) { @@ -756,7 +853,7 @@ impl RequestOperation { impl Drop for RequestOperation { fn drop(&mut self) { - if !self.terminal.is_claimed() { + if !self.publication.is_claimed() { tracing::error!( request_id = self.key.request_id, connection_id = %self.key.connection_id, @@ -768,6 +865,7 @@ impl Drop for RequestOperation { } } +#[cfg(test)] #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum CancelOperationResult { Signalled, @@ -798,6 +896,10 @@ pub(crate) enum RequestAdmissionError { connection_id: String, reason: OperationCancellationReason, }, + OwnershipClosed { + scope: String, + reason: OperationCancellationReason, + }, OwnershipMismatch { key_connection_id: String, ownership_connection_id: String, @@ -810,9 +912,9 @@ impl RequestAdmissionError { Self::CountLimit { .. } => "ERR_AGENTOS_IN_FLIGHT_REQUEST_LIMIT", Self::ByteLimit { .. } => "ERR_AGENTOS_IN_FLIGHT_REQUEST_BYTE_LIMIT", Self::DuplicateRequest { .. } => "ERR_AGENTOS_DUPLICATE_REQUEST_ID", - Self::RegistryClosed { .. } | Self::ConnectionClosed { .. } => { - "ERR_AGENTOS_REQUEST_ADMISSION_CLOSED" - } + Self::RegistryClosed { .. } + | Self::ConnectionClosed { .. } + | Self::OwnershipClosed { .. } => "ERR_AGENTOS_REQUEST_ADMISSION_CLOSED", Self::OwnershipMismatch { .. } => "ERR_AGENTOS_REQUEST_OWNERSHIP_MISMATCH", } } @@ -865,6 +967,11 @@ impl fmt::Display for RequestAdmissionError { "{}: connection {connection_id} is closed ({reason:?})", self.code() ), + Self::OwnershipClosed { scope, reason } => write!( + formatter, + "{}: request ownership {scope} is closed ({reason:?})", + self.code() + ), Self::OwnershipMismatch { key_connection_id, ownership_connection_id, @@ -879,41 +986,6 @@ impl fmt::Display for RequestAdmissionError { impl std::error::Error for RequestAdmissionError {} -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) enum OperationTransitionError { - NotFound(RequestOperationKey), - StaleGeneration(RequestOperationKey), - Invalid { - key: RequestOperationKey, - current: RequestOperationState, - next: RequestOperationState, - }, -} - -impl fmt::Display for OperationTransitionError { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::NotFound(key) => write!( - formatter, - "request operation {}:{} is not registered", - key.connection_id, key.request_id - ), - Self::StaleGeneration(key) => write!( - formatter, - "request operation {}:{} has a stale generation", - key.connection_id, key.request_id - ), - Self::Invalid { key, current, next } => write!( - formatter, - "invalid request operation transition for {}:{}: {current:?} -> {next:?}", - key.connection_id, key.request_id - ), - } - } -} - -impl std::error::Error for OperationTransitionError {} - pub(crate) const PROGRESS_REQUEST_COUNT_PATH: &str = "runtime.protocol.maxProgressFrames"; pub(crate) const PROGRESS_REQUEST_BYTES_PATH: &str = "runtime.protocol.maxProgressBytes"; @@ -936,66 +1008,6 @@ impl From<&RuntimeProtocolConfig> for ProgressRequestLimits { } } -#[derive(Clone, Debug)] -pub(crate) struct ProgressAcknowledgementGuard { - state: Arc, -} - -impl ProgressAcknowledgementGuard { - fn new() -> Self { - Self { - state: Arc::new(AtomicU8::new(PUBLICATION_UNCLAIMED)), - } - } - - fn try_claim(&self) -> bool { - self.state - .compare_exchange( - PUBLICATION_UNCLAIMED, - PUBLICATION_PUBLISHING, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() - } - - fn mark_retained(&self) -> bool { - self.state - .compare_exchange( - PUBLICATION_PUBLISHING, - PUBLICATION_RETAINED, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() - } - - fn try_take_over_unretained(&self) -> bool { - loop { - let current = self.state.load(Ordering::Acquire); - if matches!(current, PUBLICATION_RETAINED | PUBLICATION_TAKEN_OVER) { - return false; - } - if self - .state - .compare_exchange( - current, - PUBLICATION_TAKEN_OVER, - Ordering::AcqRel, - Ordering::Acquire, - ) - .is_ok() - { - return true; - } - } - } - - fn is_claimed(&self) -> bool { - self.state.load(Ordering::Acquire) != PUBLICATION_UNCLAIMED - } -} - #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) struct ProgressRequestSnapshot { pub(crate) key: RequestOperationKey, @@ -1006,7 +1018,7 @@ pub(crate) struct ProgressRequestSnapshot { } #[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct ProgressRequestRegistrySnapshot { +pub(crate) struct ProgressOperationSnapshot { pub(crate) in_flight_requests: usize, pub(crate) in_flight_request_bytes: usize, pub(crate) closed: Option, @@ -1014,65 +1026,35 @@ pub(crate) struct ProgressRequestRegistrySnapshot { pub(crate) requests: Vec, } -#[derive(Debug)] -struct ProgressRequestRecord { - generation: u64, - ownership: OwnershipScope, - request_bytes: usize, - cancellation: OperationCancellation, - acknowledgement: ProgressAcknowledgementGuard, -} - -#[derive(Debug)] -struct ProgressRegistryState { - requests: BTreeMap, - in_flight_request_bytes: usize, - next_generation: u64, - closed: Option, - closed_connections: BTreeMap, -} - -#[derive(Debug)] -struct ProgressRequestRegistryInner { - limits: ProgressRequestLimits, - state: Mutex, -} - +/// Progress-facing projection over the shared operation table. This is not an +/// independent registry: ordinary and progress admission lock the same map. #[derive(Clone, Debug)] -pub(crate) struct ProgressRequestRegistry { - inner: Arc, +pub(crate) struct ProgressOperationView { + inner: Arc, } -impl ProgressRequestRegistry { +impl ProgressOperationView { + #[cfg(test)] pub(crate) fn new(limits: ProgressRequestLimits) -> Self { - assert!( - limits.max_requests > 0, - "progress request count limit must be positive" - ); - assert!( - limits.max_request_bytes > 0, - "progress request byte limit must be positive" - ); - Self { - inner: Arc::new(ProgressRequestRegistryInner { - limits, - state: Mutex::new(ProgressRegistryState { - requests: BTreeMap::new(), - in_flight_request_bytes: 0, - next_generation: 1, - closed: None, - closed_connections: BTreeMap::new(), - }), - }), - } + OperationTable::new_with_limits( + RequestOperationLimits { + max_requests: limits.max_requests, + max_request_bytes: limits.max_request_bytes, + }, + limits, + ) + .progress_requests() } + #[cfg(test)] pub(crate) fn from_protocol_config(config: &RuntimeProtocolConfig) -> Self { - Self::new(ProgressRequestLimits::from(config)) + OperationTable::from_protocol_config(config).progress_requests() } - /// Preflight before acquiring progress-output capacity. [`Self::admit`] - /// repeats the check and remains authoritative. + /// Test-only admission probe. Production routing uses + /// [`Self::admit_owned`] once and releases the resulting operation if + /// output reservation fails. + #[cfg(test)] pub(crate) fn check_admission( &self, key: &RequestOperationKey, @@ -1082,6 +1064,7 @@ impl ProgressRequestRegistry { self.check_admission_locked(&state, key, request_bytes) } + #[cfg(test)] pub(crate) fn admit( &self, key: RequestOperationKey, @@ -1105,23 +1088,38 @@ impl ProgressRequestRegistry { ); let mut state = self.lock_state(); self.check_admission_locked(&state, &key, request_bytes)?; + if let Some((scope, reason)) = state + .closed_scopes + .iter() + .find(|(scope, _)| scope_contains(scope, &ownership)) + { + return Err(ProgressRequestAdmissionError::OwnershipClosed { + scope: ownership_label(scope), + reason: *reason, + }); + } let requested_total = state - .in_flight_request_bytes + .in_flight_progress_bytes .checked_add(request_bytes) .expect("progress admission preflight checked the request byte total"); let generation = state.next_generation; state.next_generation = state.next_generation.wrapping_add(1).max(1); let cancellation = OperationCancellation::new(); - let acknowledgement = ProgressAcknowledgementGuard::new(); - state.in_flight_request_bytes = requested_total; - state.requests.insert( + let publication = ResponsePublicationGuard::new(); + state.in_flight_progress_bytes = requested_total; + state.operations.insert( key.clone(), - ProgressRequestRecord { + OperationRecord { + class: OperationClass::Progress, generation, - ownership, + metadata: RequestOperationMetadata::new( + ownership, + "direct progress", + VmConcurrencyClass::OwnershipOnly, + ), request_bytes, cancellation: cancellation.clone(), - acknowledgement: acknowledgement.clone(), + publication: publication.clone(), }, ); drop(state); @@ -1131,21 +1129,26 @@ impl ProgressRequestRegistry { key, generation, request_bytes, + #[cfg(test)] cancellation, - acknowledgement, + publication, released: false, }) } + #[cfg(test)] pub(crate) fn cancel( &self, key: &RequestOperationKey, reason: OperationCancellationReason, ) -> ProgressCancelResult { let state = self.lock_state(); - let Some(record) = state.requests.get(key) else { + let Some(record) = state.operations.get(key) else { return ProgressCancelResult::NotFound; }; + if record.class != OperationClass::Progress { + return ProgressCancelResult::NotFound; + } if record.cancellation.signal(reason) { ProgressCancelResult::Signalled } else { @@ -1158,6 +1161,7 @@ impl ProgressRequestRegistry { } } + #[cfg(test)] pub(crate) fn close_connection( &self, connection_id: &str, @@ -1169,19 +1173,21 @@ impl ProgressRequestRegistry { .entry(connection_id.to_owned()) .or_insert(reason); state - .requests + .operations .iter() .filter(|(key, _)| key.connection_id == connection_id) + .filter(|(_, record)| record.class == OperationClass::Progress) .filter(|(_, record)| record.cancellation.signal(reason)) .count() } pub(crate) fn close(&self, reason: OperationCancellationReason) -> usize { let mut state = self.lock_state(); - state.closed.get_or_insert(reason); + state.progress_closed.get_or_insert(reason); state - .requests + .operations .values() + .filter(|record| record.class == OperationClass::Progress) .filter(|record| record.cancellation.signal(reason)) .count() } @@ -1192,13 +1198,14 @@ impl ProgressRequestRegistry { pub(crate) fn signal_all(&self, reason: OperationCancellationReason) -> usize { let state = self.lock_state(); state - .requests + .operations .values() + .filter(|record| record.class == OperationClass::Progress) .filter(|record| record.cancellation.signal(reason)) .count() } - /// Exactly-once counterpart to [`RequestOperationRegistry::force_terminalize`] + /// Exactly-once counterpart to [`OperationTable::force_terminalize`] /// for direct progress requests. Returned descriptors must receive a /// synthetic acknowledgement after task-local reservations are released. pub(crate) fn force_acknowledge( @@ -1206,43 +1213,57 @@ impl ProgressRequestRegistry { reason: OperationCancellationReason, ) -> Vec { let mut state = self.lock_state(); - state.closed.get_or_insert(reason); - let requests = std::mem::take(&mut state.requests); - state.in_flight_request_bytes = 0; - requests + state.progress_closed.get_or_insert(reason); + let keys = state + .operations + .iter() + .filter(|(_, record)| record.class == OperationClass::Progress) + .map(|(key, _)| key.clone()) + .collect::>(); + state.in_flight_progress_bytes = 0; + let outcomes = keys .into_iter() + .filter_map(|key| state.operations.remove(&key).map(|record| (key, record))) .filter_map(|(key, record)| { record.cancellation.signal(reason); // A claimed acknowledgement may remain registered while its // finite ordinary event batch drains. Remove accounting but // never synthesize a duplicate acknowledgement. - if !record.acknowledgement.try_take_over_unretained() { + if !record.publication.try_take_over_unretained() { return None; } Some(ForcedRequestOutcome { key, - ownership: record.ownership, + ownership: record.metadata.ownership, }) }) - .collect() + .collect::>(); + drop(state); + self.inner.changed.notify_waiters(); + outcomes } - pub(crate) fn snapshot(&self) -> ProgressRequestRegistrySnapshot { + pub(crate) fn snapshot(&self) -> ProgressOperationSnapshot { let state = self.lock_state(); - ProgressRequestRegistrySnapshot { - in_flight_requests: state.requests.len(), - in_flight_request_bytes: state.in_flight_request_bytes, - closed: state.closed, + ProgressOperationSnapshot { + in_flight_requests: state + .operations + .values() + .filter(|record| record.class == OperationClass::Progress) + .count(), + in_flight_request_bytes: state.in_flight_progress_bytes, + closed: state.progress_closed, closed_connections: state.closed_connections.keys().cloned().collect(), requests: state - .requests + .operations .iter() + .filter(|(_, record)| record.class == OperationClass::Progress) .map(|(key, record)| ProgressRequestSnapshot { key: key.clone(), - ownership: record.ownership.clone(), + ownership: record.metadata.ownership.clone(), request_bytes: record.request_bytes, cancellation_reason: record.cancellation.reason(), - acknowledgement_claimed: record.acknowledgement.is_claimed(), + acknowledgement_claimed: record.publication.is_claimed(), }) .collect(), } @@ -1250,11 +1271,11 @@ impl ProgressRequestRegistry { fn check_admission_locked( &self, - state: &ProgressRegistryState, + state: &RegistryState, key: &RequestOperationKey, request_bytes: usize, ) -> Result<(), ProgressRequestAdmissionError> { - if let Some(reason) = state.closed { + if let Some(reason) = state.progress_closed { return Err(ProgressRequestAdmissionError::RegistryClosed { reason }); } if let Some(reason) = state.closed_connections.get(&key.connection_id).copied() { @@ -1263,29 +1284,34 @@ impl ProgressRequestRegistry { reason, }); } - if state.requests.contains_key(key) { + if state.operations.contains_key(key) { return Err(ProgressRequestAdmissionError::DuplicateRequest { key: key.clone() }); } - if state.requests.len() >= self.inner.limits.max_requests { + let progress_count = state + .operations + .values() + .filter(|record| record.class == OperationClass::Progress) + .count(); + if progress_count >= self.inner.progress_limits.max_requests { return Err(ProgressRequestAdmissionError::CountLimit { - current: state.requests.len(), + current: progress_count, requested: 1, - limit: self.inner.limits.max_requests, + limit: self.inner.progress_limits.max_requests, }); } let requested_total = state - .in_flight_request_bytes + .in_flight_progress_bytes .checked_add(request_bytes) .ok_or(ProgressRequestAdmissionError::ByteLimit { - current: state.in_flight_request_bytes, + current: state.in_flight_progress_bytes, requested: request_bytes, - limit: self.inner.limits.max_request_bytes, + limit: self.inner.progress_limits.max_request_bytes, })?; - if requested_total > self.inner.limits.max_request_bytes { + if requested_total > self.inner.progress_limits.max_request_bytes { return Err(ProgressRequestAdmissionError::ByteLimit { - current: state.in_flight_request_bytes, + current: state.in_flight_progress_bytes, requested: request_bytes, - limit: self.inner.limits.max_request_bytes, + limit: self.inner.progress_limits.max_request_bytes, }); } Ok(()) @@ -1293,29 +1319,29 @@ impl ProgressRequestRegistry { fn release(&self, key: &RequestOperationKey, generation: u64, request_bytes: usize) { let mut state = self.lock_state(); - let should_remove = state - .requests - .get(key) - .is_some_and(|record| record.generation == generation); + let should_remove = state.operations.get(key).is_some_and(|record| { + record.class == OperationClass::Progress && record.generation == generation + }); if should_remove { - state.requests.remove(key); - state.in_flight_request_bytes = state - .in_flight_request_bytes + state.operations.remove(key); + state.in_flight_progress_bytes = state + .in_flight_progress_bytes .checked_sub(request_bytes) .unwrap_or_else(|| { tracing::error!( request_id = key.request_id, connection_id = %key.connection_id, request_bytes, - retained_bytes = state.in_flight_request_bytes, + retained_bytes = state.in_flight_progress_bytes, "ERR_AGENTOS_PROGRESS_ADMISSION_ACCOUNTING: progress request byte reservation underflow" ); 0 }); + self.inner.changed.notify_waiters(); } } - fn lock_state(&self) -> MutexGuard<'_, ProgressRegistryState> { + fn lock_state(&self) -> MutexGuard<'_, RegistryState> { self.inner.state.lock().unwrap_or_else(|poisoned| { tracing::error!( "ERR_AGENTOS_PROGRESS_REGISTRY_POISONED: recovering progress-request registry state" @@ -1327,34 +1353,38 @@ impl ProgressRequestRegistry { #[derive(Debug)] pub(crate) struct ProgressRequest { - registry: ProgressRequestRegistry, + registry: ProgressOperationView, key: RequestOperationKey, generation: u64, request_bytes: usize, + #[cfg(test)] cancellation: OperationCancellation, - acknowledgement: ProgressAcknowledgementGuard, + publication: ResponsePublicationGuard, released: bool, } impl ProgressRequest { + #[cfg(test)] pub(crate) fn key(&self) -> &RequestOperationKey { &self.key } + #[cfg(test)] pub(crate) fn request_bytes(&self) -> usize { self.request_bytes } + #[cfg(test)] pub(crate) fn cancellation(&self) -> OperationCancellation { self.cancellation.clone() } pub(crate) fn try_acknowledge(&self) -> bool { - self.acknowledgement.try_claim() + self.publication.try_claim() } pub(crate) fn mark_acknowledgement_retained(&self) -> bool { - self.acknowledgement.mark_retained() + self.publication.mark_retained() } pub(crate) fn release(mut self) { @@ -1373,7 +1403,7 @@ impl ProgressRequest { impl Drop for ProgressRequest { fn drop(&mut self) { - if !self.acknowledgement.is_claimed() { + if !self.publication.is_claimed() { tracing::error!( request_id = self.key.request_id, connection_id = %self.key.connection_id, @@ -1384,6 +1414,7 @@ impl Drop for ProgressRequest { } } +#[cfg(test)] #[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum ProgressCancelResult { Signalled, @@ -1413,6 +1444,10 @@ pub(crate) enum ProgressRequestAdmissionError { connection_id: String, reason: OperationCancellationReason, }, + OwnershipClosed { + scope: String, + reason: OperationCancellationReason, + }, } impl ProgressRequestAdmissionError { @@ -1421,9 +1456,9 @@ impl ProgressRequestAdmissionError { Self::CountLimit { .. } => "ERR_AGENTOS_PROGRESS_REQUEST_LIMIT", Self::ByteLimit { .. } => "ERR_AGENTOS_PROGRESS_REQUEST_BYTE_LIMIT", Self::DuplicateRequest { .. } => "ERR_AGENTOS_DUPLICATE_PROGRESS_REQUEST_ID", - Self::RegistryClosed { .. } | Self::ConnectionClosed { .. } => { - "ERR_AGENTOS_PROGRESS_ADMISSION_CLOSED" - } + Self::RegistryClosed { .. } + | Self::ConnectionClosed { .. } + | Self::OwnershipClosed { .. } => "ERR_AGENTOS_PROGRESS_ADMISSION_CLOSED", } } @@ -1475,6 +1510,11 @@ impl fmt::Display for ProgressRequestAdmissionError { "{}: progress request connection {connection_id} is closed ({reason:?})", self.code() ), + Self::OwnershipClosed { scope, reason } => write!( + formatter, + "{}: progress request ownership {scope} is closed ({reason:?})", + self.code() + ), } } } @@ -1497,25 +1537,19 @@ mod tests { RequestOperationMetadata::new( ownership(connection_id), "test", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ) } - fn registry(max_requests: usize, max_request_bytes: usize) -> RequestOperationRegistry { - RequestOperationRegistry::new(RequestOperationLimits { + fn registry(max_requests: usize, max_request_bytes: usize) -> OperationTable { + OperationTable::new(RequestOperationLimits { max_requests, max_request_bytes, }) } fn finish(operation: RequestOperation) { - operation - .transition(RequestOperationState::Running) - .expect("mark running"); - operation - .transition(RequestOperationState::Completing) - .expect("mark completing"); - assert!(operation.try_mark_terminal().expect("mark terminal")); + assert!(operation.try_mark_terminal()); operation.release(); } @@ -1608,10 +1642,7 @@ mod tests { registry.cancel(operation.key(), OperationCancellationReason::Shutdown), CancelOperationResult::AlreadySignalled(OperationCancellationReason::Explicit) ); - operation - .transition(RequestOperationState::Completing) - .expect("cancel completion"); - assert!(operation.try_mark_terminal().expect("cancel terminal")); + assert!(operation.try_mark_terminal()); operation.release(); } @@ -1656,8 +1687,8 @@ mod tests { RequestAdmissionError::RegistryClosed { .. } )); - assert!(a.try_mark_terminal().expect("connection close terminal")); - assert!(b.try_mark_terminal().expect("shutdown terminal")); + assert!(a.try_mark_terminal()); + assert!(b.try_mark_terminal()); a.release(); b.release(); } @@ -1674,20 +1705,12 @@ mod tests { session_id: String::from("session-a"), }), "gated", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 17, ) .expect("admit gated request"); - operation - .transition(RequestOperationState::Running) - .expect("mark running"); - operation - .transition(RequestOperationState::Completing) - .expect("begin completion before output publication"); - assert!(operation - .try_mark_terminal() - .expect("claim normal terminal publication")); + assert!(operation.try_mark_terminal()); let cancellation = operation.cancellation(); let forced = registry.force_terminalize(OperationCancellationReason::Shutdown); @@ -1717,13 +1740,7 @@ mod tests { let operation = registry .admit(RequestOperationKey::new("a", 10), metadata("a"), 11) .expect("admit retained request"); - operation - .transition(RequestOperationState::Running) - .expect("mark running"); - operation - .transition(RequestOperationState::Completing) - .expect("begin completion"); - assert!(operation.try_mark_terminal().expect("claim terminal")); + assert!(operation.try_mark_terminal()); assert!(operation.mark_terminal_retained()); assert!(registry @@ -1735,7 +1752,7 @@ mod tests { #[test] fn terminal_response_guard_is_exactly_once_across_racing_clones() { - let guard = TerminalResponseGuard::new(); + let guard = ResponsePublicationGuard::new(); let clones = (0..16).map(|_| guard.clone()).collect::>(); let winners = clones .into_iter() @@ -1748,60 +1765,27 @@ mod tests { } #[test] - fn operation_preserves_metadata_and_state_transitions() { + fn operation_preserves_metadata_and_has_one_publication_claim() { let registry = registry(1, 100); let operation = registry .admit( RequestOperationKey::new("a", 1), RequestOperationMetadata::new( - ownership("a"), + OwnershipScope::vm("a", "session", "vm"), "configure_vm", - RequestOrderingKey::VmLifecycle { - connection_id: String::from("a"), - session_id: String::from("session"), - vm_id: String::from("vm"), - }, + VmConcurrencyClass::ExclusiveVmLifecycle, ), 17, ) .expect("admit operation"); assert_eq!(operation.request_bytes(), 17); assert_eq!(operation.metadata().operation, "configure_vm"); - operation - .transition(RequestOperationState::Running) - .expect("running"); - operation - .transition(RequestOperationState::Completing) - .expect("completing"); - assert!(operation.try_mark_terminal().expect("terminal")); - assert!(!operation.try_mark_terminal().expect("duplicate terminal")); assert_eq!( - registry.snapshot().operations[0].state, - RequestOperationState::Terminal + operation.metadata().vm_concurrency, + VmConcurrencyClass::ExclusiveVmLifecycle ); - operation.release(); - } - - #[test] - fn invalid_terminal_transition_does_not_consume_the_terminal_claim() { - let registry = registry(1, 100); - let operation = registry - .admit(RequestOperationKey::new("a", 1), metadata("a"), 10) - .expect("admit operation"); - - let error = operation - .try_mark_terminal() - .expect_err("admitted work must not skip completion"); - assert!(matches!(error, OperationTransitionError::Invalid { .. })); - assert!(!registry.snapshot().operations[0].terminal_claimed); - - operation - .transition(RequestOperationState::Running) - .expect("running"); - operation - .transition(RequestOperationState::Completing) - .expect("completing"); - assert!(operation.try_mark_terminal().expect("terminal")); + assert!(operation.try_mark_terminal()); + assert!(!operation.try_mark_terminal()); operation.release(); } @@ -1811,13 +1795,7 @@ mod tests { let operation = registry .admit(RequestOperationKey::new("a", 1), metadata("a"), 10) .expect("admit operation"); - operation - .transition(RequestOperationState::Running) - .expect("running"); - operation - .transition(RequestOperationState::Completing) - .expect("completing"); - assert!(operation.try_mark_terminal().expect("terminal")); + assert!(operation.try_mark_terminal()); assert_eq!( registry.cancel(operation.key(), OperationCancellationReason::Shutdown), @@ -1825,46 +1803,31 @@ mod tests { ); assert_eq!(registry.close(OperationCancellationReason::Shutdown), 0); let snapshot = registry.snapshot(); - assert_eq!( - snapshot.operations[0].state, - RequestOperationState::Terminal - ); assert_eq!(snapshot.operations[0].cancellation_reason, None); operation.release(); } #[test] - fn cancellation_during_completion_signals_without_regressing_state() { + fn cancellation_before_publication_is_visible_to_the_operation() { let registry = registry(1, 100); let operation = registry .admit(RequestOperationKey::new("a", 1), metadata("a"), 10) .expect("admit operation"); - operation - .transition(RequestOperationState::Running) - .expect("running"); - operation - .transition(RequestOperationState::Completing) - .expect("completing"); - assert_eq!( registry.cancel(operation.key(), OperationCancellationReason::Explicit), CancelOperationResult::Signalled ); let snapshot = registry.snapshot(); - assert_eq!( - snapshot.operations[0].state, - RequestOperationState::Completing - ); assert_eq!( snapshot.operations[0].cancellation_reason, Some(OperationCancellationReason::Explicit) ); - assert!(operation.try_mark_terminal().expect("terminal")); + assert!(operation.try_mark_terminal()); operation.release(); } - fn progress_registry(max_requests: usize, max_request_bytes: usize) -> ProgressRequestRegistry { - ProgressRequestRegistry::new(ProgressRequestLimits { + fn progress_registry(max_requests: usize, max_request_bytes: usize) -> ProgressOperationView { + ProgressOperationView::new(ProgressRequestLimits { max_requests, max_request_bytes, }) @@ -1877,8 +1840,8 @@ mod tests { protocol.max_in_flight_request_bytes = 1; protocol.max_progress_frames = 2; protocol.max_progress_bytes = 10; - let ordinary_registry = RequestOperationRegistry::from_protocol_config(&protocol); - let registry = ProgressRequestRegistry::from_protocol_config(&protocol); + let ordinary_registry = OperationTable::from_protocol_config(&protocol); + let registry = ordinary_registry.progress_requests(); let first_key = RequestOperationKey::new("a", 1); registry .check_admission(&first_key, 5) @@ -1912,6 +1875,137 @@ mod tests { second.release(); } + #[test] + fn request_identity_is_unique_across_ordinary_and_progress_classes() { + let ordinary = OperationTable::new_with_limits( + RequestOperationLimits { + max_requests: 2, + max_request_bytes: 16, + }, + ProgressRequestLimits { + max_requests: 2, + max_request_bytes: 16, + }, + ); + let progress = ordinary.progress_requests(); + let key = RequestOperationKey::new("a", 7); + let ordinary_request = ordinary + .admit(key.clone(), metadata("a"), 1) + .expect("ordinary request"); + assert!(matches!( + progress.admit(key.clone(), 1), + Err(ProgressRequestAdmissionError::DuplicateRequest { .. }) + )); + assert!(ordinary_request.try_mark_terminal()); + ordinary_request.release(); + + let progress_request = progress + .admit(key.clone(), 1) + .expect("request id is reusable after release"); + assert!(matches!( + ordinary.admit(key, metadata("a"), 1), + Err(RequestAdmissionError::DuplicateRequest { .. }) + )); + assert!(progress_request.try_acknowledge()); + progress_request.release(); + } + + #[tokio::test] + async fn vm_scope_close_cancels_both_classes_and_excludes_disposer() { + let ordinary = OperationTable::new_with_limits( + RequestOperationLimits { + max_requests: 4, + max_request_bytes: 64, + }, + ProgressRequestLimits { + max_requests: 4, + max_request_bytes: 64, + }, + ); + let progress = ordinary.progress_requests(); + let vm = OwnershipScope::vm("a", "session", "vm"); + let disposer = ordinary + .admit( + RequestOperationKey::new("a", 1), + RequestOperationMetadata::new( + vm.clone(), + "dispose", + VmConcurrencyClass::ExclusiveVmLifecycle, + ), + 1, + ) + .expect("dispose request"); + let work = ordinary + .admit( + RequestOperationKey::new("a", 2), + RequestOperationMetadata::new(vm.clone(), "read", VmConcurrencyClass::SharedVm), + 1, + ) + .expect("ordinary VM work"); + let progress_work = progress + .admit_owned(RequestOperationKey::new("a", 3), vm.clone(), 1) + .expect("progress VM work"); + let work_cancel = work.cancellation(); + let progress_cancel = progress_work.cancellation(); + + let drain = ordinary.close_scope( + vm.clone(), + OperationCancellationReason::Explicit, + Some(disposer.key()), + ); + assert_eq!(drain.active_operations(), 2); + assert_eq!( + work_cancel.reason(), + Some(OperationCancellationReason::Explicit) + ); + assert_eq!( + progress_cancel.reason(), + Some(OperationCancellationReason::Explicit) + ); + assert_eq!(disposer.cancellation().reason(), None); + assert!(matches!( + ordinary.admit( + RequestOperationKey::new("a", 4), + RequestOperationMetadata::new(vm.clone(), "late", VmConcurrencyClass::SharedVm), + 1, + ), + Err(RequestAdmissionError::OwnershipClosed { .. }) + )); + let other_vm = ordinary + .admit( + RequestOperationKey::new("a", 5), + RequestOperationMetadata::new( + OwnershipScope::vm("a", "session", "other"), + "other VM", + VmConcurrencyClass::SharedVm, + ), + 1, + ) + .expect("other VM remains independent"); + + assert!(work.try_mark_terminal()); + work.release(); + assert!(progress_work.try_acknowledge()); + progress_work.release(); + drain.wait_drained().await; + assert_eq!(drain.active_operations(), 0); + + ordinary.reopen_scope(&vm); + let reopened = ordinary + .admit( + RequestOperationKey::new("a", 6), + RequestOperationMetadata::new(vm, "reopened", VmConcurrencyClass::SharedVm), + 1, + ) + .expect("new VM generation reopens exact scope"); + assert!(reopened.try_mark_terminal()); + reopened.release(); + assert!(other_vm.try_mark_terminal()); + other_vm.release(); + assert!(disposer.try_mark_terminal()); + disposer.release(); + } + #[test] fn progress_byte_admission_releases_via_raii() { let registry = progress_registry(4, 5); diff --git a/crates/native-sidecar/src/service.rs b/crates/native-sidecar/src/service.rs index cf01c26fff..a00e097b00 100644 --- a/crates/native-sidecar/src/service.rs +++ b/crates/native-sidecar/src/service.rs @@ -4711,6 +4711,41 @@ where } fn reject_error(&self, request: &RequestFrame, error: &SidecarError) -> ResponseFrame { + if let SidecarError::RequestAdmission { + code, + message, + configuration_path, + retryable, + errno, + } = error + { + let vm_id = match &request.ownership { + OwnershipScope::VmOwnership(owner) => Some(owner.vm_id.clone()), + OwnershipScope::ConnectionOwnership(_) | OwnershipScope::SessionOwnership(_) => { + None + } + }; + return self.respond( + request, + ResponsePayload::Rejected(RejectedResponse { + code: (*code).to_owned(), + message: message.clone(), + limit_name: None, + configured_limit: None, + current_usage: None, + requested: Some(1), + unit: Some(String::from("requests")), + scope: Some(String::from("vm")), + vm_id, + session_generation: None, + capability_id: None, + operation: Some(String::from("vm.lifecycleAdmission")), + configuration_path: configuration_path.map(str::to_owned), + retryable: Some(*retryable), + errno: Some((*errno).to_owned()), + }), + ); + } let SidecarError::ResourceLimit(limit) = error else { return self.reject(request, error_code(error), &error.to_string()); }; diff --git a/crates/native-sidecar/src/state.rs b/crates/native-sidecar/src/state.rs index 825915b3cb..f76cb4edce 100644 --- a/crates/native-sidecar/src/state.rs +++ b/crates/native-sidecar/src/state.rs @@ -568,6 +568,13 @@ impl Default for NativeSidecarConfig { #[derive(Debug, Clone, PartialEq, Eq)] pub enum SidecarError { ResourceLimit(agentos_runtime::accounting::LimitError), + RequestAdmission { + code: &'static str, + message: String, + configuration_path: Option<&'static str>, + retryable: bool, + errno: &'static str, + }, InvalidState(String), ProtocolVersionMismatch(String), BridgeVersionMismatch(String), @@ -586,6 +593,7 @@ impl fmt::Display for SidecarError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::ResourceLimit(error) => error.fmt(f), + Self::RequestAdmission { message, .. } => f.write_str(message), Self::InvalidState(message) | Self::ProtocolVersionMismatch(message) | Self::BridgeVersionMismatch(message) diff --git a/crates/native-sidecar/src/stdio.rs b/crates/native-sidecar/src/stdio.rs index 4caaf602f7..e60fe7b354 100644 --- a/crates/native-sidecar/src/stdio.rs +++ b/crates/native-sidecar/src/stdio.rs @@ -4,12 +4,13 @@ use crate::extension_services::{ CompletedExtensionServiceCommand, ExtensionServiceCommand, OwnedProcessEventService, PreparedExtensionServiceCommand, RoutedExtensionServices, VmEventAdmissionResult, }; -use crate::ownership_coordinator::{CoordinatorOperationPermit, OwnershipCoordinator, VmDisposal}; +use crate::ownership_coordinator::{ + CoordinatorOperationPermit, OwnershipCoordinator, OwnershipCoordinatorError, VmDisposal, +}; use crate::request_operations::{ - ForcedRequestOutcome, OperationCancellationReason, ProgressRequest, - ProgressRequestAdmissionError, ProgressRequestRegistry, RequestAdmissionError, - RequestOperation, RequestOperationKey, RequestOperationMetadata, RequestOperationRegistry, - RequestOperationState, RequestOrderingKey, + ForcedRequestOutcome, OperationCancellationReason, OperationTable, ProgressOperationView, + ProgressRequest, ProgressRequestAdmissionError, RequestAdmissionError, RequestOperation, + RequestOperationKey, RequestOperationMetadata, VmConcurrencyClass, }; use crate::service::CompletedExtensionRequest; use crate::service::{CompletedRequest, PreparedMembershipCommit, PreparedRequest}; @@ -449,6 +450,7 @@ struct ProtocolOutputQueueState { rejection: VecDeque, terminal: VecDeque, observability: VecDeque, + next_control_class: u8, open: bool, terminal_error: Option, } @@ -463,11 +465,21 @@ impl ProtocolOutputQueueState { } fn pop_control(&mut self) -> Option { - self.progress - .pop_front() - .or_else(|| self.rejection.pop_front()) - .or_else(|| self.terminal.pop_front()) - .or_else(|| self.observability.pop_front()) + for offset in 0..4 { + let class = (self.next_control_class + offset) % 4; + let frame = match class { + 0 => self.progress.pop_front(), + 1 => self.rejection.pop_front(), + 2 => self.terminal.pop_front(), + 3 => self.observability.pop_front(), + _ => unreachable!("control class is reduced modulo four"), + }; + if frame.is_some() { + self.next_control_class = (class + 1) % 4; + return frame; + } + } + None } } @@ -492,6 +504,7 @@ impl ProtocolOutputQueue { rejection: VecDeque::new(), terminal: VecDeque::new(), observability: VecDeque::new(), + next_control_class: 0, open: true, terminal_error: None, }), @@ -1551,22 +1564,27 @@ async fn run_async( // separate from the runtime producer edge makes the protocol coordinator // the sole consumer that can drain execution-engine queues. let routed_process_event_notify = Arc::new(Notify::new()); - let request_operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let request_operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = request_operations.progress_requests(); let extension_service_capacity = protocol .max_in_flight_requests .saturating_add(protocol.max_control_frames) .max(1); let (extension_service_tx, extension_service_rx) = channel::(extension_service_capacity); + let progress_service_capacity = protocol.max_progress_frames.max(1); + let (progress_service_tx, progress_service_rx) = + channel::(progress_service_capacity); let (extension_service_completion_tx, extension_service_completion_rx) = channel::(extension_service_capacity); - let extension_services: Arc = - Arc::new(RoutedExtensionServices::new_with_process_event_broker( + let extension_services: Arc = Arc::new( + RoutedExtensionServices::new_with_process_event_broker_and_progress( extension_service_tx, + progress_service_tx, Arc::clone(&routed_process_event_notify), sidecar.process_event_broker(), - )); + ), + ); sidecar.set_extension_services(Arc::clone(&extension_services)); let (extension_completion_tx, extension_completion_rx) = channel::(extension_service_capacity); @@ -1751,6 +1769,7 @@ async fn run_async( stdin_control_rx, shutdown_rx, extension_service_rx, + progress_service_rx, extension_service_completion_tx, extension_service_completion_rx, extension_completion_tx, @@ -1777,14 +1796,15 @@ struct ProtocolEngine { sidecar: NativeSidecar, extension_services: Arc, ownership_coordinator: OwnershipCoordinator, - request_operations: RequestOperationRegistry, - progress_requests: ProgressRequestRegistry, + request_operations: OperationTable, + progress_requests: ProgressOperationView, callback_transport: Arc, frame_writer: ProtocolFrameWriter, stdin_rx: Receiver, String>>, stdin_control_rx: Receiver, shutdown_rx: Receiver, extension_service_rx: Receiver, + progress_service_rx: Receiver, extension_service_completion_tx: Sender, extension_service_completion_rx: Receiver, extension_completion_tx: Sender, @@ -1813,6 +1833,7 @@ async fn run_protocol_engine(engine: ProtocolEngine) -> Result<(), Box Result<(), Box>(); let mut active_connections = sidecar.connections.keys().cloned().collect::>(); let mut extension_service_tasks = JoinSet::<()>::new(); + let mut progress_service_tasks = JoinSet::<()>::new(); let mut extension_tasks = JoinSet::<()>::new(); let mut request_tasks = JoinSet::<()>::new(); let mut output_tasks = JoinSet::>::new(); @@ -1851,6 +1873,7 @@ async fn run_protocol_engine(engine: ProtocolEngine) -> Result<(), Box Result<(), Box Result<(), Box Result<(), Box Result<(), Box { + let Some(command) = maybe_progress_service else { + begin_protocol_transport_failure( + String::from("progress extension service command channel closed while router was active"), + &mut drain_state, + shutdown_grace, + &request_operations, + &progress_requests, + &callback_transport, + &frame_writer, + false, + &mut stdin_closed, + &mut control_ingress_closed, + &mut shutdown_ingress_closed, + &mut stdin_rx, + &mut stdin_control_rx, + &mut shutdown_rx, + ); + continue 'protocol; + }; + let prepared = prepare_extension_service_command( + &mut sidecar, + &ownership_coordinator, + command, + ); + schedule_extension_service_command( + prepared, + &extension_service_completion_tx, + &mut progress_service_tasks, + ); + } maybe_service = extension_service_rx.recv(), if extension_service_tasks.len() < extension_service_capacity => { let Some(command) = maybe_service else { begin_protocol_transport_failure( @@ -2566,6 +2626,28 @@ async fn run_protocol_engine(engine: ProtocolEngine) -> Result<(), Box { + if let Some(Err(error)) = maybe_task { + begin_protocol_transport_failure( + format!( + "ERR_AGENTOS_PROGRESS_SERVICE_SUPERVISOR_TASK: completion monitor failed: {error}" + ), + &mut drain_state, + shutdown_grace, + &request_operations, + &progress_requests, + &callback_transport, + &frame_writer, + false, + &mut stdin_closed, + &mut control_ingress_closed, + &mut shutdown_ingress_closed, + &mut stdin_rx, + &mut stdin_control_rx, + &mut shutdown_rx, + ); + } + } maybe_task = extension_tasks.join_next(), if !extension_tasks.is_empty() => { if let Some(Err(error)) = maybe_task { begin_protocol_transport_failure( @@ -2728,6 +2810,7 @@ async fn run_protocol_engine(engine: ProtocolEngine) -> Result<(), Box, - operations: &RequestOperationRegistry, - progress_requests: &ProgressRequestRegistry, + operations: &OperationTable, + progress_requests: &ProgressOperationView, close_progress_admission: bool, ) -> bool { if let Some(state) = drain_state { @@ -2801,8 +2884,8 @@ fn begin_protocol_transport_failure( message: String, drain_state: &mut Option, grace: Duration, - operations: &RequestOperationRegistry, - progress_requests: &ProgressRequestRegistry, + operations: &OperationTable, + progress_requests: &ProgressOperationView, callback_transport: &FrameSidecarRequestTransport, writer: &ProtocolFrameWriter, output_failed: bool, @@ -2846,11 +2929,12 @@ async fn finalize_protocol_drain( reason: OperationCancellationReason, output_grace: Duration, terminal_fallback_bytes: usize, - operations: &RequestOperationRegistry, - progress_requests: &ProgressRequestRegistry, + operations: &OperationTable, + progress_requests: &ProgressOperationView, callback_transport: &FrameSidecarRequestTransport, writer: &ProtocolFrameWriter, extension_service_tasks: &mut JoinSet<()>, + progress_service_tasks: &mut JoinSet<()>, extension_service_completion_rx: &mut Receiver, extension_tasks: &mut JoinSet<()>, extension_completion_rx: &mut Receiver, @@ -2869,6 +2953,8 @@ async fn finalize_protocol_drain( extension_service_tasks.abort_all(); while extension_service_tasks.join_next().await.is_some() {} + progress_service_tasks.abort_all(); + while progress_service_tasks.join_next().await.is_some() {} while let Ok(completion) = extension_service_completion_rx.try_recv() { drop(completion); } @@ -3056,8 +3142,8 @@ fn route_protocol_frame( accounted_frame: AccountedProtocolFrame, sidecar: &mut NativeSidecar, services: &Arc, - operations: &RequestOperationRegistry, - progress_requests: &ProgressRequestRegistry, + operations: &OperationTable, + progress_requests: &ProgressOperationView, ownership_coordinator: &OwnershipCoordinator, completion_tx: &Sender, request_completion_tx: &Sender, @@ -3087,15 +3173,23 @@ fn route_protocol_frame( ownership_connection_id(&request.ownership), request.request_id, ); - if let Err(error) = progress_requests.check_admission(&key, request_bytes) { - publish_progress_admission_rejection(write_tx, request, &error)?; - return Ok(()); - } + let progress_request = match progress_requests.admit_owned( + key, + request.ownership.clone(), + request_bytes, + ) { + Ok(progress_request) => progress_request, + Err(error) => { + publish_progress_admission_rejection(write_tx, request, &error)?; + return Ok(()); + } + }; let progress_fallback_bytes = terminal_fallback_bytes.min(write_tx.progress_budget.config.max_bytes); let reservation = match write_tx.try_reserve_progress(progress_fallback_bytes) { Ok(reservation) => reservation, Err(error) => { + drop(progress_request); publish_request_rejection( write_tx, request, @@ -3108,29 +3202,26 @@ fn route_protocol_frame( return Ok(()); } }; - match progress_requests.admit_owned(key, request.ownership.clone(), request_bytes) { - Ok(progress_request) => (None, Some(progress_request), reservation), - Err(error) => { - drop(reservation); - publish_progress_admission_rejection(write_tx, request, &error)?; - return Ok(()); - } - } + (None, Some(progress_request), reservation) } else { let metadata = request_operation_metadata(&request, &sidecar.extensions); let key = RequestOperationKey::new( ownership_connection_id(&request.ownership), request.request_id, ); - if let Err(error) = operations.check_admission(&key, &metadata, request_bytes) { - publish_request_admission_rejection(write_tx, request, &error)?; - return Ok(()); - } + let operation = match operations.admit(key, metadata, request_bytes) { + Ok(operation) => operation, + Err(error) => { + publish_request_admission_rejection(write_tx, request, &error)?; + return Ok(()); + } + }; let terminal_reservation = match write_tx .try_reserve_terminal(terminal_fallback_bytes) { Ok(reservation) => reservation, Err(error) => { + drop(operation); publish_request_rejection( write_tx, request, @@ -3143,20 +3234,9 @@ fn route_protocol_frame( return Ok(()); } }; - match operations.admit(key, metadata, request_bytes) { - Ok(operation) => (Some(operation), None, terminal_reservation), - Err(error) => { - drop(terminal_reservation); - publish_request_admission_rejection(write_tx, request, &error)?; - return Ok(()); - } - } + (Some(operation), None, terminal_reservation) }; - if let Some(operation) = &operation { - operation.transition(RequestOperationState::Running)?; - } - let prepared = match sidecar .prepare_extension_request_wire(request.clone(), Arc::clone(services)) { @@ -3201,7 +3281,7 @@ fn route_protocol_frame( progress_request, coordinator_permit: None, output_reservation, - result: Err(SidecarError::InvalidState(error.to_string())), + result: Err(coordinator_admission_error(error)), }; if completion_tx.send(completion).await.is_err() { tracing::error!( @@ -3326,14 +3406,17 @@ fn route_protocol_frame( return Ok(()); } }; - let disposal = match ownership_coordinator - .begin_vm_disposal(&request.ownership, cancellation_reason) - { + let disposal = match ownership_coordinator.begin_vm_disposal_with_operations( + &request.ownership, + cancellation_reason, + operations, + operation.key(), + ) { Ok(disposal) => disposal, Err(error) => { let dispatch = sidecar.reject_wire_request_error( request.clone(), - &SidecarError::InvalidState(error.to_string()), + &coordinator_admission_error(error), )?; schedule_dispatch_output( dispatch, @@ -3449,6 +3532,7 @@ fn route_protocol_frame( fn reap_protocol_tasks_nowait( extension_service_tasks: &mut JoinSet<()>, + progress_service_tasks: &mut JoinSet<()>, extension_tasks: &mut JoinSet<()>, request_tasks: &mut JoinSet<()>, output_tasks: &mut JoinSet>, @@ -3459,6 +3543,10 @@ fn reap_protocol_tasks_nowait( extension_service_tasks, "ERR_AGENTOS_EXTENSION_SERVICE_SUPERVISOR_TASK", )?; + reap_unit_tasks_nowait( + progress_service_tasks, + "ERR_AGENTOS_PROGRESS_SERVICE_SUPERVISOR_TASK", + )?; reap_unit_tasks_nowait( extension_tasks, "ERR_AGENTOS_REQUEST_SUPERVISOR_TASK: extension", @@ -3541,8 +3629,8 @@ fn schedule_prepared_request( operation, coordinator_permit: None, output_reservation, - result: DetachedRequestResult::Generic(Err(SidecarError::InvalidState( - error.to_string(), + result: DetachedRequestResult::Generic(Err(coordinator_admission_error( + error, ))), }; if request_completion_tx.send(completion).await.is_err() { @@ -3620,9 +3708,7 @@ fn schedule_prepared_create_vm( operation, coordinator_permit: None, output_reservation, - result: DetachedRequestResult::Create(Err(SidecarError::InvalidState( - error.to_string(), - ))), + result: DetachedRequestResult::Create(Err(coordinator_admission_error(error))), }; if request_completion_tx.send(completion).await.is_err() { tracing::error!( @@ -3904,7 +3990,12 @@ fn finish_request( } }; if update_membership { - update_ownership_membership(ownership_coordinator, &request, &dispatch.response.payload)?; + update_ownership_membership( + ownership_coordinator, + &operation, + &request, + &dispatch.response.payload, + )?; } drop(coordinator_permit); schedule_dispatch_output( @@ -3928,7 +4019,7 @@ fn schedule_dispatch_output( operation: Option, progress_request: Option, output_reservation: ProtocolReservation, - failed: bool, + _failed: bool, write_tx: &ProtocolFrameWriter, output_tasks: &mut JoinSet>, active_sessions: &mut BTreeSet, @@ -3964,17 +4055,7 @@ fn schedule_dispatch_output( let operation = operation.as_ref().ok_or_else(|| { String::from("ordinary request completion lost its operation reservation") })?; - operation - .transition(if failed { - RequestOperationState::Failed - } else { - RequestOperationState::Completing - }) - .map_err(|error| error.to_string())?; - if !operation - .try_mark_terminal() - .map_err(|error| error.to_string())? - { + if !operation.try_mark_terminal() { return Err(String::from( "ERR_AGENTOS_DUPLICATE_TERMINAL_RESPONSE: request terminal response was already claimed", )); @@ -4057,9 +4138,11 @@ fn rearm_event_ready(sender: &Sender<()>) -> Result<(), io::Error> { fn request_operation_metadata( request: &RequestFrame, - extensions: &BTreeMap>, + _extensions: &BTreeMap>, ) -> RequestOperationMetadata { - let connection_id = ownership_connection_id(&request.ownership).to_owned(); + // Lifecycle requests mutate or require a stable view of VM-wide state, so + // they use the exclusive gate. Every other core VM request is shared by + // default; this is deliberately not a per-request ordering matrix. let vm_lifecycle = matches!( &request.payload, RequestPayload::DisposeVmRequest(_) @@ -4073,47 +4156,25 @@ fn request_operation_metadata( | RequestPayload::SnapshotRootFilesystemRequest(_) | RequestPayload::LinkPackageRequest(_) ); - let ordering_key = match &request.payload { - RequestPayload::ExtEnvelope(envelope) => extensions - .get(&envelope.namespace) - .and_then(|extension| { - extension.request_ordering_key(&request.ownership, &envelope.payload) - }) - .map(|key| RequestOrderingKey::Extension { - namespace: envelope.namespace.clone(), - connection_id: connection_id.clone(), - key, - policy: extensions - .get(&envelope.namespace) - .expect("extension key came from registered extension") - .request_ordering_policy(&request.ownership, &envelope.payload), - }) - .unwrap_or(RequestOrderingKey::Unordered), + let vm_concurrency = match &request.payload { + // Extension payloads are opaque to the core. ACP and other extensions + // own any protocol-specific conflict state in their route handlers. + RequestPayload::ExtEnvelope(_) => VmConcurrencyClass::OwnershipOnly, _ => match &request.ownership { - OwnershipScope::ConnectionOwnership(_) => { - RequestOrderingKey::Connection(connection_id.clone()) + OwnershipScope::ConnectionOwnership(_) | OwnershipScope::SessionOwnership(_) => { + VmConcurrencyClass::OwnershipOnly } - OwnershipScope::SessionOwnership(scope) => RequestOrderingKey::Session { - connection_id: connection_id.clone(), - session_id: scope.session_id.clone(), - }, - OwnershipScope::VmOwnership(scope) if vm_lifecycle => RequestOrderingKey::VmLifecycle { - connection_id: connection_id.clone(), - session_id: scope.session_id.clone(), - vm_id: scope.vm_id.clone(), - }, - OwnershipScope::VmOwnership(scope) => RequestOrderingKey::VmOperation { - connection_id: connection_id.clone(), - session_id: scope.session_id.clone(), - vm_id: scope.vm_id.clone(), - }, + OwnershipScope::VmOwnership(_) if vm_lifecycle => { + VmConcurrencyClass::ExclusiveVmLifecycle + } + OwnershipScope::VmOwnership(_) => VmConcurrencyClass::SharedVm, }, }; let operation = match &request.payload { RequestPayload::ExtEnvelope(envelope) => format!("extension:{}", envelope.namespace), _ => String::from("sidecar_request"), }; - RequestOperationMetadata::new(request.ownership.clone(), operation, ordering_key) + RequestOperationMetadata::new(request.ownership.clone(), operation, vm_concurrency) } fn ownership_connection_id(ownership: &OwnershipScope) -> &str { @@ -4124,6 +4185,16 @@ fn ownership_connection_id(ownership: &OwnershipScope) -> &str { } } +fn coordinator_admission_error(error: OwnershipCoordinatorError) -> SidecarError { + SidecarError::RequestAdmission { + code: error.code(), + message: error.to_string(), + configuration_path: error.configuration_path(), + retryable: error.retryable(), + errno: error.errno(), + } +} + fn publish_request_admission_rejection( writer: &ProtocolFrameWriter, request: RequestFrame, @@ -4167,7 +4238,8 @@ fn publish_request_admission_rejection( Some(String::from("EEXIST")), ), RequestAdmissionError::RegistryClosed { .. } - | RequestAdmissionError::ConnectionClosed { .. } => ( + | RequestAdmissionError::ConnectionClosed { .. } + | RequestAdmissionError::OwnershipClosed { .. } => ( None, None, None, @@ -4253,7 +4325,8 @@ fn publish_progress_admission_rejection( Some(String::from("EEXIST")), ), ProgressRequestAdmissionError::RegistryClosed { .. } - | ProgressRequestAdmissionError::ConnectionClosed { .. } => ( + | ProgressRequestAdmissionError::ConnectionClosed { .. } + | ProgressRequestAdmissionError::OwnershipClosed { .. } => ( None, None, None, @@ -4377,24 +4450,39 @@ async fn cleanup_connections( fn update_ownership_membership( coordinator: &OwnershipCoordinator, + operation: &RequestOperation, request: &RequestFrame, response: &ResponsePayload, ) -> Result<(), io::Error> { let result = match response { ResponsePayload::AuthenticatedResponse(response) => { + let scope = OwnershipScope::connection(&response.connection_id); + operation.reopen_scope(&scope); ensure_connection_membership(coordinator, &response.connection_id) } - ResponsePayload::SessionOpenedResponse(response) => ensure_session_membership( - coordinator, - &response.owner_connection_id, - &response.session_id, - ), + ResponsePayload::SessionOpenedResponse(response) => { + let scope = + OwnershipScope::session(&response.owner_connection_id, &response.session_id); + operation.reopen_scope(&scope); + ensure_session_membership( + coordinator, + &response.owner_connection_id, + &response.session_id, + ) + } ResponsePayload::VmCreatedResponse(response) => match &request.ownership { - OwnershipScope::SessionOwnership(scope) => coordinator - .connection(&scope.connection_id) - .and_then(|connection| connection.session(&scope.session_id)) - .and_then(|session| session.open_vm(response.vm_id.clone())) - .map(|_| ()), + OwnershipScope::SessionOwnership(scope) => { + operation.reopen_scope(&OwnershipScope::vm( + &scope.connection_id, + &scope.session_id, + &response.vm_id, + )); + coordinator + .connection(&scope.connection_id) + .and_then(|connection| connection.session(&scope.session_id)) + .and_then(|session| session.open_vm(response.vm_id.clone())) + .map(|_| ()) + } ownership => Err( crate::ownership_coordinator::OwnershipCoordinatorError::OwnershipMismatch { expected: String::from("session ownership for CreateVmRequest"), @@ -5160,7 +5248,7 @@ mod tests { } #[test] - fn request_metadata_distinguishes_vm_lifecycle_from_vm_operations() { + fn request_metadata_maps_ownership_and_vm_concurrency_independently() { let extensions = BTreeMap::new(); let lifecycle = request_frame( 1, @@ -5170,12 +5258,8 @@ mod tests { }), ); assert!(matches!( - request_operation_metadata(&lifecycle, &extensions).ordering_key, - RequestOrderingKey::VmLifecycle { - connection_id, - session_id, - vm_id, - } if connection_id == "conn" && session_id == "session" && vm_id == "vm" + request_operation_metadata(&lifecycle, &extensions).vm_concurrency, + VmConcurrencyClass::ExclusiveVmLifecycle )); let operation = request_frame( @@ -5184,13 +5268,23 @@ mod tests { RequestPayload::GetProcessSnapshotRequest, ); assert!(matches!( - request_operation_metadata(&operation, &extensions).ordering_key, - RequestOrderingKey::VmOperation { - connection_id, - session_id, - vm_id, - } if connection_id == "conn" && session_id == "session" && vm_id == "vm" + request_operation_metadata(&operation, &extensions).vm_concurrency, + VmConcurrencyClass::SharedVm )); + + let session = request_frame( + 3, + session_ownership("conn", "session"), + RequestPayload::CreateVmRequest(wire::CreateVmRequest::legacy_test_config( + wire::GuestRuntimeKind::JavaScript, + Default::default(), + Default::default(), + None, + )), + ); + let metadata = request_operation_metadata(&session, &extensions); + assert_eq!(metadata.ownership, session_ownership("conn", "session")); + assert_eq!(metadata.vm_concurrency, VmConcurrencyClass::OwnershipOnly); } #[test] @@ -5226,6 +5320,7 @@ mod tests { const WAVES: usize = 32; let mut extension_service_tasks = JoinSet::new(); + let mut progress_service_tasks = JoinSet::new(); let mut extension_tasks = JoinSet::new(); let mut request_tasks = JoinSet::new(); let mut output_tasks = JoinSet::new(); @@ -5308,6 +5403,7 @@ mod tests { reap_protocol_tasks_nowait( &mut extension_service_tasks, + &mut progress_service_tasks, &mut extension_tasks, &mut request_tasks, &mut output_tasks, @@ -5317,6 +5413,7 @@ mod tests { .expect("pre-select reap succeeds"); assert!(extension_service_tasks.is_empty()); + assert!(progress_service_tasks.is_empty()); assert!(extension_tasks.is_empty()); assert!(request_tasks.is_empty()); assert!(output_tasks.is_empty()); @@ -5350,8 +5447,7 @@ mod tests { runtime.context(), ) .expect("test sidecar"); - let operations = - RequestOperationRegistry::from_protocol_config(&config.runtime.protocol); + let operations = OperationTable::from_protocol_config(&config.runtime.protocol); let ownership_coordinator = OwnershipCoordinator::from_runtime_config(&config.runtime); ownership_coordinator @@ -5395,9 +5491,6 @@ mod tests { 1, ) .expect("admit generic request"); - operation - .transition(RequestOperationState::Running) - .expect("mark generic request running"); schedule_prepared_request( prepared, request, @@ -5493,9 +5586,6 @@ mod tests { 1, ) .expect("admit panicking request"); - operation - .transition(RequestOperationState::Running) - .expect("mark panicking request running"); schedule_prepared_request( prepared, request, @@ -5630,10 +5720,9 @@ mod tests { .open_vm(vm_id.clone()) .expect("test coordinator VM"); - let operations = - RequestOperationRegistry::from_protocol_config(&config.runtime.protocol); + let operations = OperationTable::from_protocol_config(&config.runtime.protocol); let progress_requests = - ProgressRequestRegistry::from_protocol_config(&config.runtime.protocol); + ProgressOperationView::from_protocol_config(&config.runtime.protocol); let (writer, output) = test_frame_writer_with_inflight(8, 2); let ingress_budget = test_protocol_budget(4, 4096, "dispose route ingress"); let (service_tx, _service_rx) = channel(4); @@ -5692,9 +5781,6 @@ mod tests { 1, ) .expect("admit gated VM operation"); - operation - .transition(RequestOperationState::Running) - .expect("start gated VM operation"); let operation_cancellation = operation.cancellation(); schedule_prepared_request( prepared_operation, @@ -5895,10 +5981,9 @@ mod tests { ownership_coordinator .register_connection("conn-independent-service") .expect("register independent connection"); - let operations = - RequestOperationRegistry::from_protocol_config(&config.runtime.protocol); + let operations = OperationTable::from_protocol_config(&config.runtime.protocol); let progress_requests = - ProgressRequestRegistry::from_protocol_config(&config.runtime.protocol); + ProgressOperationView::from_protocol_config(&config.runtime.protocol); let ingress_budget = test_protocol_budget(4, 4096, "test request ingress"); let (writer, output) = test_frame_writer_with_inflight(8, 2); let (service_tx, _service_rx) = channel(4); @@ -6104,11 +6189,7 @@ mod tests { &RequestOperationMetadata::new( crate::protocol::OwnershipScope::vm(connection_id, session_id, &vm_id), "saturate ordinary VM admission", - RequestOrderingKey::VmOperation { - connection_id: connection_id.to_owned(), - session_id: session_id.to_owned(), - vm_id: vm_id.clone(), - }, + VmConcurrencyClass::SharedVm, ), crate::request_operations::OperationCancellation::new(), ) @@ -6361,8 +6442,8 @@ export async function loadPyodide() { } let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ingress_budget = test_protocol_budget(4, 4096, "root Python VFS ingress"); let (writer, output) = test_frame_writer_with_inflight(8, 2); let (service_tx, _service_rx) = channel(2); @@ -6820,8 +6901,8 @@ export async function loadPyodide() { ) .expect("test sidecar"); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership_coordinator = OwnershipCoordinator::from_runtime_config( &NativeSidecarConfig::default().runtime, ); @@ -6992,8 +7073,8 @@ export async function loadPyodide() { ) .expect("test sidecar"); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership_coordinator = OwnershipCoordinator::from_runtime_config( &NativeSidecarConfig::default().runtime, ); @@ -7175,8 +7256,8 @@ export async function loadPyodide() { ) .expect("test sidecar"); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership_coordinator = OwnershipCoordinator::from_runtime_config( &NativeSidecarConfig::default().runtime, ); @@ -7393,10 +7474,9 @@ export async function loadPyodide() { let mut count_protocol = agentos_runtime::RuntimeProtocolConfig::default(); count_protocol.max_in_flight_requests = 1; count_protocol.max_in_flight_request_bytes = 8; - let count_operations = - RequestOperationRegistry::from_protocol_config(&count_protocol); + let count_operations = OperationTable::from_protocol_config(&count_protocol); let count_progress_requests = - ProgressRequestRegistry::from_protocol_config(&count_protocol); + ProgressOperationView::from_protocol_config(&count_protocol); for request_id in [30, 31] { let request = ProtocolFrame::RequestFrame(request_frame( request_id, @@ -7466,10 +7546,9 @@ export async function loadPyodide() { let mut byte_protocol = agentos_runtime::RuntimeProtocolConfig::default(); byte_protocol.max_in_flight_requests = 2; byte_protocol.max_in_flight_request_bytes = 1; - let byte_operations = - RequestOperationRegistry::from_protocol_config(&byte_protocol); + let byte_operations = OperationTable::from_protocol_config(&byte_protocol); let byte_progress_requests = - ProgressRequestRegistry::from_protocol_config(&byte_protocol); + ProgressOperationView::from_protocol_config(&byte_protocol); for request_id in [40, 41] { let request = ProtocolFrame::RequestFrame(request_frame( request_id, @@ -7550,8 +7629,8 @@ export async function loadPyodide() { ) .expect("test sidecar"); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership_coordinator = OwnershipCoordinator::from_runtime_config( &NativeSidecarConfig::default().runtime, ); @@ -8410,8 +8489,8 @@ export async function loadPyodide() { async fn shutdown_takeover_before_retention_publishes_one_terminal_and_progress_frame() { let (writer, output) = test_frame_writer_with_inflight(8, 1); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership = connection_ownership("takeover-before-retention"); let operation = operations @@ -8420,18 +8499,12 @@ export async function loadPyodide() { RequestOperationMetadata::new( ownership.clone(), "terminal-race", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 11, ) .expect("admit terminal race"); - operation - .transition(RequestOperationState::Running) - .expect("mark terminal race running"); - operation - .transition(RequestOperationState::Completing) - .expect("mark terminal race completing"); - assert!(operation.try_mark_terminal().expect("claim terminal")); + assert!(operation.try_mark_terminal()); let terminal_reservation = writer .try_reserve_terminal(writer.terminal_budget.config.max_bytes) .expect("reserve original terminal"); @@ -8565,8 +8638,8 @@ export async function loadPyodide() { async fn broker_retention_before_shutdown_prevents_terminal_and_progress_takeover() { let (writer, output) = test_frame_writer_with_inflight(8, 1); let protocol = agentos_runtime::RuntimeProtocolConfig::default(); - let operations = RequestOperationRegistry::from_protocol_config(&protocol); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::from_protocol_config(&protocol); + let progress_requests = ProgressOperationView::from_protocol_config(&protocol); let ownership = connection_ownership("retention-before-takeover"); let operation = operations @@ -8575,18 +8648,12 @@ export async function loadPyodide() { RequestOperationMetadata::new( ownership.clone(), "terminal-race", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 11, ) .expect("admit terminal race"); - operation - .transition(RequestOperationState::Running) - .expect("mark terminal race running"); - operation - .transition(RequestOperationState::Completing) - .expect("mark terminal race completing"); - assert!(operation.try_mark_terminal().expect("claim terminal")); + assert!(operation.try_mark_terminal()); let terminal_reservation = writer .try_reserve_terminal(writer.terminal_budget.config.max_bytes) .expect("reserve terminal"); @@ -8656,10 +8723,10 @@ export async function loadPyodide() { tokio::task::LocalSet::new() .run_until(async { let (writer, output) = test_frame_writer_with_inflight(8, 1); - let operations = RequestOperationRegistry::from_protocol_config( + let operations = OperationTable::from_protocol_config( &agentos_runtime::RuntimeProtocolConfig::default(), ); - let progress_requests = ProgressRequestRegistry::from_protocol_config( + let progress_requests = ProgressOperationView::from_protocol_config( &agentos_runtime::RuntimeProtocolConfig::default(), ); let ordinary_ownership = session_ownership("shutdown-conn", "session-a"); @@ -8669,14 +8736,11 @@ export async function loadPyodide() { RequestOperationMetadata::new( ordinary_ownership.clone(), "gated-shutdown", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 17, ) .expect("admit gated ordinary request"); - operation - .transition(RequestOperationState::Running) - .expect("mark gated request running"); let cancellation = operation.cancellation(); let terminal_reservation = writer .try_reserve_terminal(writer.terminal_budget.config.max_bytes) @@ -8730,7 +8794,7 @@ export async function loadPyodide() { &RequestOperationMetadata::new( ordinary_ownership.clone(), "late", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 1, ), @@ -8765,6 +8829,7 @@ export async function loadPyodide() { let (_completion_tx, mut completion_rx) = channel(2); let (_service_completion_tx, mut service_completion_rx) = channel(2); let mut service_tasks = JoinSet::new(); + let mut progress_service_tasks = JoinSet::new(); let (_request_completion_tx, mut request_completion_rx) = channel(2); let mut request_tasks = JoinSet::new(); let mut output_tasks = JoinSet::new(); @@ -8778,6 +8843,7 @@ export async function loadPyodide() { &callback_transport, &writer, &mut service_tasks, + &mut progress_service_tasks, &mut service_completion_rx, &mut extension_tasks, &mut completion_rx, @@ -8836,10 +8902,10 @@ export async function loadPyodide() { tokio::task::LocalSet::new() .run_until(async { let (writer, output) = test_frame_writer_with_inflight(8, 1); - let operations = RequestOperationRegistry::from_protocol_config( + let operations = OperationTable::from_protocol_config( &agentos_runtime::RuntimeProtocolConfig::default(), ); - let progress_requests = ProgressRequestRegistry::from_protocol_config( + let progress_requests = ProgressOperationView::from_protocol_config( &agentos_runtime::RuntimeProtocolConfig::default(), ); let ownership = connection_ownership("disconnect-conn"); @@ -8849,14 +8915,11 @@ export async function loadPyodide() { RequestOperationMetadata::new( ownership.clone(), "gated-disconnect", - RequestOrderingKey::Unordered, + VmConcurrencyClass::OwnershipOnly, ), 17, ) .expect("admit disconnected ordinary request"); - operation - .transition(RequestOperationState::Running) - .expect("mark disconnected request running"); let terminal_reservation = writer .try_reserve_terminal(writer.terminal_budget.config.max_bytes) .expect("reserve original terminal outcome"); @@ -8915,6 +8978,7 @@ export async function loadPyodide() { let (_completion_tx, mut completion_rx) = channel(2); let (_service_completion_tx, mut service_completion_rx) = channel(2); let mut service_tasks = JoinSet::new(); + let mut progress_service_tasks = JoinSet::new(); let (_request_completion_tx, mut request_completion_rx) = channel(2); let mut request_tasks = JoinSet::new(); let mut output_tasks = JoinSet::new(); @@ -8930,6 +8994,7 @@ export async function loadPyodide() { &callback_transport, &writer, &mut service_tasks, + &mut progress_service_tasks, &mut service_completion_rx, &mut extension_tasks, &mut completion_rx, @@ -9127,6 +9192,37 @@ export async function loadPyodide() { output.close(); } + #[tokio::test] + async fn protocol_output_progress_burst_cannot_starve_terminal_response() { + let (writer, output) = test_frame_writer(16); + writer + .try_send_progress(queue_test_sidecar_request(-1)) + .expect("queue first progress frame"); + writer + .try_send_progress(queue_test_sidecar_request(-2)) + .expect("queue second progress frame"); + writer + .try_send(queue_test_response(3)) + .expect("queue terminal response"); + + let first = decode_test_output(output.recv_control().await.expect("first output")); + let second = decode_test_output(output.recv_control().await.expect("second output")); + let third = decode_test_output(output.recv_control().await.expect("third output")); + assert!(matches!( + first, + ProtocolFrame::SidecarRequestFrame(frame) if frame.request_id == -1 + )); + assert!(matches!( + second, + ProtocolFrame::ResponseFrame(frame) if frame.request_id == 3 + )); + assert!(matches!( + third, + ProtocolFrame::SidecarRequestFrame(frame) if frame.request_id == -2 + )); + output.close(); + } + #[test] fn combined_stdio_preserves_control_priority_ahead_of_ordinary() { let (writer, output) = test_frame_writer(16); diff --git a/crates/native-sidecar/src/stdio/request_concurrency_tests.rs b/crates/native-sidecar/src/stdio/request_concurrency_tests.rs index f25f0d3224..a051d6bda1 100644 --- a/crates/native-sidecar/src/stdio/request_concurrency_tests.rs +++ b/crates/native-sidecar/src/stdio/request_concurrency_tests.rs @@ -89,13 +89,6 @@ impl Extension for GatedExtension { } } - fn request_ordering_key(&self, _ownership: &OwnershipScope, payload: &[u8]) -> Option> { - let command = std::str::from_utf8(payload).ok()?; - let keyed = command.strip_prefix("key:")?; - let (key, _) = keyed.split_once(':')?; - Some(key.as_bytes().to_vec()) - } - fn handle_request<'a>( &'a self, _ctx: crate::ExtensionContext, @@ -173,9 +166,11 @@ struct ProtocolLoopHarness { ingress_budget: ProtocolBudget, control_budget: ProtocolBudget, extension_routes: Arc>>, + extension_services: Arc, + ordinary_service_capacity: usize, ownership_coordinator: OwnershipCoordinator, - operations: RequestOperationRegistry, - progress_requests: ProgressRequestRegistry, + operations: OperationTable, + progress_requests: ProgressOperationView, } impl ProtocolLoopHarness { @@ -234,12 +229,11 @@ impl ProtocolLoopHarness { } } - let operations = - RequestOperationRegistry::new(crate::request_operations::RequestOperationLimits { - max_requests, - max_request_bytes, - }); - let progress_requests = ProgressRequestRegistry::from_protocol_config(&protocol); + let operations = OperationTable::new(crate::request_operations::RequestOperationLimits { + max_requests, + max_request_bytes, + }); + let progress_requests = operations.progress_requests(); let output = Arc::new(ProtocolOutputQueue::new( protocol.max_egress_frames, protocol.max_control_frames, @@ -269,12 +263,17 @@ impl ProtocolLoopHarness { .saturating_add(protocol.max_control_frames) .max(1); let (service_tx, service_rx) = channel(service_capacity); - let extension_services: Arc = - Arc::new(RoutedExtensionServices::new_with_process_event_broker( + let (progress_service_tx, progress_service_rx) = + channel(protocol.max_progress_frames.max(1)); + let routed_extension_services = Arc::new( + RoutedExtensionServices::new_with_process_event_broker_and_progress( service_tx, + progress_service_tx, Arc::clone(&routed_process_event_notify), sidecar.process_event_broker(), - )); + ), + ); + let extension_services: Arc = routed_extension_services.clone(); let (extension_completion_tx, extension_completion_rx) = channel(service_capacity); let (extension_service_completion_tx, extension_service_completion_rx) = channel(service_capacity); @@ -291,6 +290,8 @@ impl ProtocolLoopHarness { ingress_budget, control_budget, extension_routes, + extension_services: routed_extension_services, + ordinary_service_capacity: service_capacity, ownership_coordinator: ownership_coordinator.clone(), operations: operations.clone(), progress_requests: progress_requests.clone(), @@ -308,6 +309,7 @@ impl ProtocolLoopHarness { stdin_control_rx: control_rx, shutdown_rx, extension_service_rx: service_rx, + progress_service_rx, extension_service_completion_tx, extension_service_completion_rx, extension_completion_tx, @@ -430,6 +432,112 @@ fn extension_request( ) } +fn guest_filesystem_request( + request_id: RequestId, + connection_id: &str, + session_id: &str, + vm_id: &str, + operation: wire::GuestFilesystemOperation, + path: &str, + content: Option<&str>, +) -> RequestFrame { + request_frame( + request_id, + vm_ownership(connection_id, session_id, vm_id), + RequestPayload::GuestFilesystemCallRequest(wire::GuestFilesystemCallRequest { + operation, + path: path.to_owned(), + destination_path: None, + target: None, + content: content.map(str::to_owned), + encoding: None, + recursive: false, + max_depth: None, + mode: None, + uid: None, + gid: None, + atime_ms: None, + mtime_ms: None, + len: None, + offset: None, + }), + ) +} + +async fn start_protocol_loop_with_vm( + state: Arc, +) -> ( + ProtocolLoopHarness, + tokio::task::JoinHandle>>, + String, + String, + String, +) { + let (harness, engine) = ProtocolLoopHarness::build(state, &[], 8, 16 * 1024); + let engine_task = tokio::task::spawn_local(run_protocol_engine(engine)); + harness.send_request( + request_frame( + 1, + connection_ownership("client-hint"), + RequestPayload::AuthenticateRequest(wire::AuthenticateRequest { + client_name: String::from("VM concurrency regression"), + auth_token: String::new(), + protocol_version: wire::PROTOCOL_VERSION, + bridge_version: agentos_bridge::bridge_contract().version, + }), + ), + 1, + ); + let authenticated = harness.response().await; + let ResponsePayload::AuthenticatedResponse(authenticated) = authenticated.payload else { + panic!("expected authenticated response"); + }; + let connection_id = authenticated.connection_id; + harness.send_request( + request_frame( + 2, + connection_ownership(&connection_id), + RequestPayload::OpenSessionRequest(wire::OpenSessionRequest { + placement: wire::SidecarPlacement::SidecarPlacementShared( + wire::SidecarPlacementShared { pool: None }, + ), + metadata: Default::default(), + }), + ), + 1, + ); + let opened = harness.response().await; + let ResponsePayload::SessionOpenedResponse(opened) = opened.payload else { + panic!("expected session-opened response"); + }; + let session_id = opened.session_id; + let vm_config: agentos_vm_config::CreateVmConfig = + serde_json::from_value(serde_json::json!({ "permissions": { "fs": "allow" } })) + .expect("build filesystem-enabled VM config"); + harness.send_request( + request_frame( + 3, + session_ownership(&connection_id, &session_id), + RequestPayload::CreateVmRequest(wire::CreateVmRequest { + runtime: wire::GuestRuntimeKind::JavaScript, + config: serde_json::to_string(&vm_config).expect("serialize test VM config"), + }), + ), + 1, + ); + let created = harness.response().await; + let ResponsePayload::VmCreatedResponse(created) = created.payload else { + panic!("expected VM-created response"); + }; + ( + harness, + engine_task, + connection_id, + session_id, + created.vm_id, + ) +} + fn protocol_loop_binding_process( vm: &mut crate::state::VmState, label: &str, @@ -569,6 +677,244 @@ async fn request_concurrency_real_loop_starts_two_blocking_sessions_together() { .await; } +#[tokio::test(flavor = "current_thread")] +async fn request_concurrency_real_loop_same_vm_read_and_write_share_the_gate() { + tokio::task::LocalSet::new() + .run_until(async { + let state = Arc::new(GatedExtensionState::default()); + let (harness, engine_task, connection_id, session_id, vm_id) = + start_protocol_loop_with_vm(state).await; + + harness.send_request( + guest_filesystem_request( + 4, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::WriteFile, + "/seed.txt", + Some("seed"), + ), + 1, + ); + assert!(matches!( + harness.response().await.payload, + ResponsePayload::GuestFilesystemResultResponse(_) + )); + + // Retain one real ordinary permit while both protocol requests + // run. If the gate serialized ordinary work, neither could finish + // until this permit was dropped. + let held = harness + .ownership_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership(&connection_id, &session_id, &vm_id), + "held same-VM ordinary operation", + VmConcurrencyClass::SharedVm, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + .expect("hold one ordinary VM permit"); + harness.send_request( + guest_filesystem_request( + 5, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::ReadFile, + "/seed.txt", + None, + ), + 1, + ); + harness.send_request( + guest_filesystem_request( + 6, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::WriteFile, + "/parallel.txt", + Some("parallel"), + ), + 1, + ); + let responses = [harness.response().await, harness.response().await]; + assert_eq!( + responses + .iter() + .map(|response| response.request_id) + .collect::>(), + BTreeSet::from([5, 6]), + ); + assert!(responses.iter().all(|response| matches!( + response.payload, + ResponsePayload::GuestFilesystemResultResponse(_) + ))); + drop(held); + finish_cleanly(&harness, engine_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn request_concurrency_real_loop_sleeping_prompt_does_not_hold_same_vm_gate() { + tokio::task::LocalSet::new() + .run_until(async { + let state = Arc::new(GatedExtensionState::default()); + let (harness, engine_task, connection_id, session_id, vm_id) = + start_protocol_loop_with_vm(Arc::clone(&state)).await; + harness.send_request( + guest_filesystem_request( + 4, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::WriteFile, + "/during-prompt.txt", + Some("readable"), + ), + 1, + ); + assert_eq!(harness.response().await.request_id, 4); + + harness.send_request( + extension_request(5, &connection_id, &session_id, "block:prompt-vm"), + 1, + ); + state.gate("prompt-vm").wait_started().await; + harness.send_request( + guest_filesystem_request( + 6, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::ReadFile, + "/during-prompt.txt", + None, + ), + 1, + ); + harness.send_request( + guest_filesystem_request( + 7, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::WriteFile, + "/also-during-prompt.txt", + Some("writable"), + ), + 1, + ); + let filesystem = [harness.response().await, harness.response().await]; + assert_eq!( + filesystem + .iter() + .map(|response| response.request_id) + .collect::>(), + BTreeSet::from([6, 7]), + "same-VM filesystem work must finish before the sleeping prompt", + ); + state.gate("prompt-vm").release(); + assert_eq!(harness.response().await.request_id, 5); + finish_cleanly(&harness, engine_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn request_concurrency_real_loop_configure_waits_and_rejects_new_same_vm_work() { + tokio::task::LocalSet::new() + .run_until(async { + let state = Arc::new(GatedExtensionState::default()); + let (harness, engine_task, connection_id, session_id, vm_id) = + start_protocol_loop_with_vm(state).await; + let held = harness + .ownership_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership(&connection_id, &session_id, &vm_id), + "held pre-configure VM operation", + VmConcurrencyClass::SharedVm, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + .expect("hold an earlier ordinary operation"); + harness.send_request( + request_frame( + 4, + vm_ownership(&connection_id, &session_id, &vm_id), + RequestPayload::ConfigureVmRequest(wire::ConfigureVmRequest { + mounts: Vec::new(), + software: Vec::new(), + permissions: None, + module_access_cwd: None, + instructions: Vec::new(), + projected_modules: Vec::new(), + command_permissions: Default::default(), + loopback_exempt_ports: Vec::new(), + packages: Vec::new(), + packages_mount_at: String::new(), + bootstrap_commands: Vec::new(), + binding_shim_commands: Vec::new(), + }), + ), + 1, + ); + let vm = harness + .ownership_coordinator + .connection(&connection_id) + .expect("connection coordinator") + .session(&session_id) + .expect("session coordinator") + .vm(&vm_id) + .expect("VM coordinator"); + tokio::time::timeout(TEST_TIMEOUT, async { + while vm.snapshot().lifecycle + != crate::ownership_coordinator::VmLifecyclePhase::Pending + { + tokio::task::yield_now().await; + } + }) + .await + .expect("configure enters pending lifecycle state"); + + harness.send_request( + guest_filesystem_request( + 5, + &connection_id, + &session_id, + &vm_id, + wire::GuestFilesystemOperation::Stat, + "/", + None, + ), + 1, + ); + let rejected = harness.response().await; + assert_eq!(rejected.request_id, 5); + let ResponsePayload::RejectedResponse(rejection) = rejected.payload else { + panic!("same-VM work must be rejected while configure is pending"); + }; + assert_eq!(rejection.code, "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT"); + assert_eq!(rejection.retryable, Some(true)); + + drop(held); + let configured = harness.response().await; + assert_eq!(configured.request_id, 4); + assert!(matches!( + configured.payload, + ResponsePayload::VmConfiguredResponse(_) + )); + finish_cleanly(&harness, engine_task).await; + }) + .await; +} + #[tokio::test(flavor = "current_thread")] async fn request_concurrency_real_loop_pipelined_membership_commits_before_dependent_routing() { tokio::task::LocalSet::new() @@ -780,11 +1126,7 @@ async fn request_concurrency_real_loop_vm_a_disposal_does_not_delay_vm_b_query() &RequestOperationMetadata::new( vm_ownership("conn-1", "session-1", &vm_a), "held VM-A operation", - RequestOrderingKey::VmOperation { - connection_id: String::from("conn-1"), - session_id: String::from("session-1"), - vm_id: vm_a.clone(), - }, + VmConcurrencyClass::SharedVm, ), held_cancellation.clone(), ) @@ -824,6 +1166,14 @@ async fn request_concurrency_real_loop_vm_a_disposal_does_not_delay_vm_b_query() "disposal must signal the operation it is waiting to drain", ); + harness.send_request( + request_frame( + 36, + vm_ownership("conn-1", "session-1", &vm_a), + RequestPayload::GetProcessSnapshotRequest, + ), + 1, + ); harness.send_request( request_frame( 35, @@ -832,11 +1182,11 @@ async fn request_concurrency_real_loop_vm_a_disposal_does_not_delay_vm_b_query() ), 1, ); - let independent = harness.response().await; - assert_eq!( - independent.request_id, 35, - "VM-B generic work must complete while VM-A disposal is gated", - ); + let pending_responses = [harness.response().await, harness.response().await]; + let independent = pending_responses + .iter() + .find(|response| response.request_id == 35) + .expect("VM-B response completes during VM-A drain"); assert!( matches!( &independent.payload, @@ -845,7 +1195,23 @@ async fn request_concurrency_real_loop_vm_a_disposal_does_not_delay_vm_b_query() "unexpected VM-B response: {:?}", independent.payload ); - + let same_vm = pending_responses + .iter() + .find(|response| response.request_id == 36) + .expect("same-VM request receives a typed conflict"); + let ResponsePayload::RejectedResponse(rejection) = &same_vm.payload else { + panic!("same-VM request must be rejected while disposal is pending"); + }; + assert!( + matches!( + rejection.code.as_str(), + "ERR_AGENTOS_REQUEST_ADMISSION_CLOSED" + | "ERR_AGENTOS_COORDINATOR_CLOSING" + | "ERR_AGENTOS_VM_LIFECYCLE_CONFLICT" + ), + "unexpected same-VM lifecycle conflict: {}", + rejection.code, + ); drop(held_vm_a); let disposed = harness.response().await; assert_eq!(disposed.request_id, 34); @@ -1309,7 +1675,7 @@ async fn targeted_public_waiter_cannot_consume_internal_process_pump_wake() { } #[tokio::test(flavor = "current_thread")] -async fn request_concurrency_real_loop_enforces_only_matching_opaque_extension_keys() { +async fn request_concurrency_real_loop_leaves_opaque_route_exclusion_to_extension() { tokio::task::LocalSet::new() .run_until(async { let state = Arc::new(GatedExtensionState::default()); @@ -1336,18 +1702,11 @@ async fn request_concurrency_real_loop_enforces_only_matching_opaque_extension_k ); let responses = [harness.response().await, harness.response().await]; - let conflict = responses + let same_key = responses .iter() .find(|response| response.request_id == 16) - .expect("same opaque key receives a terminal conflict"); - let ResponsePayload::RejectedResponse(rejection) = &conflict.payload else { - panic!("same opaque key must be rejected"); - }; - assert!( - rejection.message.contains("ERR_AGENTOS_ORDERING_CONFLICT"), - "conflict remains typed: {}", - rejection.message - ); + .expect("same opaque key progresses independently"); + assert_eq!(response_payload(same_key), b"conflict"); let independent = responses .iter() .find(|response| response.request_id == 17) @@ -1574,6 +1933,257 @@ async fn request_concurrency_real_loop_cancel_bypasses_saturated_ordinary_admiss .await; } +#[tokio::test(flavor = "current_thread")] +async fn request_concurrency_real_loop_progress_service_bypasses_full_ordinary_service_queue() { + tokio::task::LocalSet::new() + .run_until(async { + let state = Arc::new(GatedExtensionState::default()); + let (harness, engine) = ProtocolLoopHarness::build(state, &[], 1, 4096); + let mut ordinary = Vec::with_capacity(harness.ordinary_service_capacity); + let waker = std::task::Waker::noop(); + let mut context = std::task::Context::from_waker(waker); + for _ in 0..harness.ordinary_service_capacity { + let mut request = harness.extension_services.acp_termination_grace(); + assert!(matches!( + request.as_mut().poll(&mut context), + std::task::Poll::Pending + )); + ordinary.push(request); + } + + // The ordinary service receiver has not started and its physical + // queue is full. WriteStdin is the adapter-cancellation path; it + // must still acquire the independently bounded progress lane. + let mut progress = harness.extension_services.write_stdin( + vm_ownership("missing", "missing", "missing"), + wire::WriteStdinRequest { + process_id: String::from("missing"), + chunk: b"cancel\n".to_vec(), + }, + ); + assert!(matches!( + progress.as_mut().poll(&mut context), + std::task::Poll::Pending + )); + + let engine_task = tokio::task::spawn_local(run_protocol_engine(engine)); + let progress_error = tokio::time::timeout(TEST_TIMEOUT, progress) + .await + .expect("progress service reaches the running protocol loop") + .expect_err("missing VM returns a normal routed service error"); + assert!( + !progress_error + .to_string() + .contains("ERR_AGENTOS_PROGRESS_SERVICE_LIMIT"), + "progress was admitted through its reserved service lane: {progress_error}", + ); + for request in ordinary { + tokio::time::timeout(TEST_TIMEOUT, request) + .await + .expect("ordinary service request drains") + .expect("termination-grace service succeeds"); + } + finish_cleanly(&harness, engine_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn request_concurrency_bounded_multi_vm_lifecycle_progress_load_finishes_exactly_once() { + tokio::task::LocalSet::new() + .run_until(async { + let state = Arc::new(GatedExtensionState::default()); + let sessions: &[&str] = &["route-a", "route-b", "route-c", "route-d"]; + let (harness, engine) = ProtocolLoopHarness::build( + Arc::clone(&state), + &[("conn", sessions)], + 64, + 64 * 1024, + ); + let connection = harness + .ownership_coordinator + .connection("conn") + .expect("load-test connection coordinator"); + let vm_a = connection + .session("route-a") + .expect("load-test route A") + .open_vm("vm-a") + .expect("open load-test VM A"); + let vm_b = connection + .session("route-b") + .expect("load-test route B") + .open_vm("vm-b") + .expect("open load-test VM B"); + let held_a = harness + .ownership_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership("conn", "route-a", "vm-a"), + "delayed load operation A", + VmConcurrencyClass::SharedVm, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + .expect("delay VM A ordinary completion"); + let held_b = harness + .ownership_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership("conn", "route-b", "vm-b"), + "delayed load operation B", + VmConcurrencyClass::SharedVm, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + .expect("delay VM B ordinary completion"); + let lifecycle_a_coordinator = harness.ownership_coordinator.clone(); + let lifecycle_a = tokio::task::spawn_local(async move { + lifecycle_a_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership("conn", "route-a", "vm-a"), + "periodic lifecycle A", + VmConcurrencyClass::ExclusiveVmLifecycle, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + }); + let lifecycle_b_coordinator = harness.ownership_coordinator.clone(); + let lifecycle_b = tokio::task::spawn_local(async move { + lifecycle_b_coordinator + .admit( + &RequestOperationMetadata::new( + vm_ownership("conn", "route-b", "vm-b"), + "periodic lifecycle B", + VmConcurrencyClass::ExclusiveVmLifecycle, + ), + crate::request_operations::OperationCancellation::new(), + ) + .await + }); + tokio::time::timeout(TEST_TIMEOUT, async { + while vm_a.snapshot().lifecycle + != crate::ownership_coordinator::VmLifecyclePhase::Pending + || vm_b.snapshot().lifecycle + != crate::ownership_coordinator::VmLifecyclePhase::Pending + { + tokio::task::yield_now().await; + } + }) + .await + .expect("both independent VM lifecycle requests become pending"); + let engine_task = tokio::task::spawn_local(run_protocol_engine(engine)); + + for index in 0..8_i64 { + let session = sessions[index as usize % sessions.len()]; + harness.send_request( + extension_request( + 1_000 + index, + "conn", + session, + &format!("block:load-{index}"), + ), + 1, + ); + } + for index in 0..8_i64 { + state.gate(&format!("load-{index}")).wait_started().await; + } + // Deterministic frame-fuzz case: the first ordinary request is + // still live when a progress-class frame reuses its connection ID. + // The shared table rejects the duplicate without delivering the + // cancel payload to the original operation. + harness.send_request( + extension_request(1_000, "conn", "route-a", "cancel:load-0"), + 1, + ); + let duplicate = harness.response().await; + assert_eq!(duplicate.request_id, 1_000); + let ResponsePayload::RejectedResponse(duplicate) = duplicate.payload else { + panic!("cross-class duplicate frame must be rejected"); + }; + assert_eq!(duplicate.code, "ERR_AGENTOS_DUPLICATE_PROGRESS_REQUEST_ID"); + for index in 0..16_i64 { + let session = sessions[index as usize % sessions.len()]; + harness.send_request( + extension_request(1_100 + index, "conn", session, &format!("echo-{index}")), + 1, + ); + } + for index in 0..8_i64 { + let session = sessions[index as usize % sessions.len()]; + harness.send_request( + extension_request( + 1_200 + index, + "conn", + session, + &format!("cancel:load-{index}"), + ), + 1, + ); + } + + drop(held_a); + let lifecycle_a = tokio::time::timeout(TEST_TIMEOUT, lifecycle_a) + .await + .expect("VM A lifecycle completion deadline") + .expect("VM A lifecycle task joined") + .expect("VM A lifecycle activates independently"); + assert_eq!( + vm_b.snapshot().lifecycle, + crate::ownership_coordinator::VmLifecyclePhase::Pending, + "VM A lifecycle completion must not alter VM B", + ); + drop(lifecycle_a); + drop(held_b); + let lifecycle_b = tokio::time::timeout(TEST_TIMEOUT, lifecycle_b) + .await + .expect("VM B lifecycle completion deadline") + .expect("VM B lifecycle task joined") + .expect("VM B lifecycle activates independently"); + drop(lifecycle_b); + + let responses = tokio::time::timeout(TEST_TIMEOUT, async { + let mut ids = BTreeSet::new(); + for _ in 0..32 { + let response = harness.response().await; + assert!( + ids.insert(response.request_id), + "duplicate terminal response" + ); + } + ids + }) + .await + .expect("bounded mixed ordinary/progress load completes"); + let expected = (1_000..1_008) + .chain(1_100..1_116) + .chain(1_200..1_208) + .collect::>(); + assert_eq!(responses, expected); + assert_eq!(harness.operations.snapshot().in_flight_requests, 0); + assert_eq!(harness.operations.snapshot().in_flight_request_bytes, 0); + assert_eq!(harness.progress_requests.snapshot().in_flight_requests, 0); + assert_eq!( + harness.progress_requests.snapshot().in_flight_request_bytes, + 0 + ); + assert_eq!( + vm_a.snapshot().lifecycle, + crate::ownership_coordinator::VmLifecyclePhase::Idle, + ); + assert_eq!( + vm_b.snapshot().lifecycle, + crate::ownership_coordinator::VmLifecyclePhase::Idle, + ); + finish_cleanly(&harness, engine_task).await; + }) + .await; +} + #[tokio::test(flavor = "current_thread")] async fn request_concurrency_real_loop_retains_admission_through_backpressured_event_batches() { tokio::task::LocalSet::new() diff --git a/crates/native-sidecar/tests/architecture_guards.rs b/crates/native-sidecar/tests/architecture_guards.rs index f656c3fae2..da75030a71 100644 --- a/crates/native-sidecar/tests/architecture_guards.rs +++ b/crates/native-sidecar/tests/architecture_guards.rs @@ -704,13 +704,80 @@ fn native_sidecar_has_no_prompt_specific_interrupt_workaround() { std::fs::read_to_string(root.join("crates/native-sidecar/src/extension.rs")) .expect("read generic extension contract"); assert!( - extension_contract.contains("fn request_ordering_key("), - "generic extension contract must expose an opaque ordering-key hook" + !extension_contract.contains("request_ordering_key"), + "generic extension routing must not recreate a core-owned ordering matrix" ); +} + +#[test] +fn vm_concurrency_classification_stays_narrow_and_auditable() { + let root = repo_root(); + let stdio = std::fs::read_to_string(root.join("crates/native-sidecar/src/stdio.rs")) + .expect("read native-sidecar stdio source"); + let classification = stdio + .split("fn request_operation_metadata") + .nth(1) + .and_then(|tail| tail.split("fn ownership_connection_id").next()) + .expect("locate request VM concurrency classification"); + for lifecycle_request in [ + "DisposeVmRequest", + "BootstrapRootFilesystemRequest", + "ConfigureVmRequest", + "CreateLayerRequest", + "SealLayerRequest", + "ImportSnapshotRequest", + "ExportSnapshotRequest", + "CreateOverlayRequest", + "SnapshotRootFilesystemRequest", + "LinkPackageRequest", + ] { + assert!( + classification.contains(lifecycle_request), + "VM lifecycle request {lifecycle_request} requires explicit exclusive classification" + ); + } assert!( - !extension_contract.contains("agentos_protocol"), - "generic extension ordering must not depend on ACP protocol types" - ); + classification.contains("RequestPayload::ExtEnvelope(_)") + && classification.contains("VmConcurrencyClass::OwnershipOnly") + && classification.contains("VmConcurrencyClass::SharedVm") + && classification.contains("VmConcurrencyClass::ExclusiveVmLifecycle"), + "classification must keep extension, ordinary VM, and lifecycle behavior explicit" + ); + + let operations = + std::fs::read_to_string(root.join("crates/native-sidecar/src/request_operations.rs")) + .expect("read request operation table source"); + let class = operations + .split("enum VmConcurrencyClass") + .nth(1) + .and_then(|tail| tail.split("/// Complete admission description").next()) + .expect("locate VM concurrency class"); + for variant in ["OwnershipOnly", "SharedVm", "ExclusiveVmLifecycle"] { + assert!( + class.contains(variant), + "missing VM concurrency class {variant}" + ); + } + assert!( + !class.contains("Extension") && !operations.contains("extension_conflicts"), + "native-sidecar must not recreate extension-specific conflict domains" + ); + + let ownership = + std::fs::read_to_string(root.join("crates/native-sidecar/src/ownership_coordinator.rs")) + .expect("read ownership coordinator source"); + for rationale in [ + "This is not a standard-library or Tokio `RwLock`", + "forbidden design is a lock held across request execution", + "reject", + "internal event", + "generation", + ] { + assert!( + ownership.contains(rationale), + "VM lifecycle gate documentation is missing rationale marker {rationale}" + ); + } } #[test] diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index 984311e9e6..01b6d81b2b 100644 --- a/crates/runtime/src/lib.rs +++ b/crates/runtime/src/lib.rs @@ -71,6 +71,7 @@ const DEFAULT_MAX_TERMINAL_TASK_REPORTS: usize = 4_096; const DEFAULT_VM_EXECUTOR_TEARDOWN_TIMEOUT_MS: u64 = 5_000; pub const DEFAULT_PROTOCOL_MAX_INGRESS_FRAMES: usize = 128; pub const DEFAULT_PROTOCOL_MAX_INGRESS_BYTES: usize = 64 * 1024 * 1024; +pub const DEFAULT_PROTOCOL_MAX_SESSIONS_PER_CONNECTION: usize = 4_096; pub const DEFAULT_PROTOCOL_MAX_CONTROL_FRAMES: usize = 1_024; pub const DEFAULT_PROTOCOL_MAX_CONTROL_BYTES: usize = 64 * 1024 * 1024; pub const DEFAULT_PROTOCOL_MAX_EGRESS_FRAMES: usize = 4_096; @@ -120,6 +121,8 @@ pub fn is_runtime_worker_thread() -> bool { pub struct RuntimeProtocolConfig { pub max_ingress_frames: usize, pub max_ingress_bytes: usize, + /// Live session membership retained for one authenticated connection. + pub max_sessions_per_connection: usize, pub max_control_frames: usize, pub max_control_bytes: usize, pub max_egress_frames: usize, @@ -154,6 +157,7 @@ impl Default for RuntimeProtocolConfig { Self { max_ingress_frames: DEFAULT_PROTOCOL_MAX_INGRESS_FRAMES, max_ingress_bytes: DEFAULT_PROTOCOL_MAX_INGRESS_BYTES, + max_sessions_per_connection: DEFAULT_PROTOCOL_MAX_SESSIONS_PER_CONNECTION, max_control_frames: DEFAULT_PROTOCOL_MAX_CONTROL_FRAMES, max_control_bytes: DEFAULT_PROTOCOL_MAX_CONTROL_BYTES, max_egress_frames: DEFAULT_PROTOCOL_MAX_EGRESS_FRAMES, @@ -581,6 +585,10 @@ impl RuntimeConfig { "runtime.protocol.maxPendingResponses", self.protocol.max_pending_responses, ), + ( + "runtime.protocol.maxSessionsPerConnection", + self.protocol.max_sessions_per_connection, + ), ( "runtime.protocol.maxPendingResponseBytes", self.protocol.max_pending_response_bytes, @@ -1845,6 +1853,7 @@ mod tests { assert_zero_rejected!(max_in_flight_requests, "maxInFlightRequests"); assert_zero_rejected!(max_in_flight_request_bytes, "maxInFlightRequestBytes"); + assert_zero_rejected!(max_sessions_per_connection, "maxSessionsPerConnection"); assert_zero_rejected!(max_terminal_frames, "maxTerminalFrames"); assert_zero_rejected!(max_terminal_bytes, "maxTerminalBytes"); assert_zero_rejected!(terminal_fallback_bytes, "terminalFallbackBytes"); diff --git a/request-concurrency-fix-prompt.md b/request-concurrency-fix-prompt.md index 80c0229913..6ea6121861 100644 --- a/request-concurrency-fix-prompt.md +++ b/request-concurrency-fix-prompt.md @@ -114,7 +114,8 @@ a P0 progress test from passing. registration remains live so teardown can cancel and drain it; that registration is not an exclusive state lease and does not serialize ordinary operations. -- [x] Same-entity conflicts have an explicit ordering key and bounded admission. +- [x] Same-VM lifecycle conflicts use one bounded per-VM gate; ordinary work is + concurrent by default and ACP owns its route-level exclusion. - [x] Every admitted request has exactly one terminal response. - [x] Terminal responses preserve request ID and ownership and may complete out of order. @@ -198,6 +199,7 @@ Add explicit runtime protocol configuration: - `runtime.protocol.maxInFlightRequests` - `runtime.protocol.maxInFlightRequestBytes` +- `runtime.protocol.maxSessionsPerConnection` - `runtime.protocol.maxTerminalFrames` - `runtime.protocol.maxTerminalBytes` - `runtime.protocol.terminalFallbackBytes` @@ -220,17 +222,17 @@ behind an application FIFO. ### Operation record -Each admitted request has one logical tracked record split across the registry, -its operation handle, and the task supervisor: +Each admitted request has one logical tracked record in the shared operation +table plus an RAII handle retained by the task supervisor: ```rust struct OperationRecord { generation: u64, - metadata: RequestOperationMetadata, // ownership + ordering + class: OperationClass, // Ordinary | Progress + metadata: RequestOperationMetadata, // ownership + VM concurrency request_bytes: usize, - state: RequestOperationState, cancellation: CancellationToken, - terminal: TerminalResponseGuard, + publication: ResponsePublicationGuard, } // RequestOperation carries the registry key/generation and admission @@ -238,20 +240,11 @@ struct OperationRecord { // output reservation remain outside the registry record. ``` -Required state transitions: - -```text -Admitted -> Running -> Completing -> Terminal - -> Cancelling -> Completing -> Terminal -Admitted/Running -> Failed -> Terminal -Admitted/Running -> Shutdown -> Completing -> Terminal -``` - -The terminal guard atomically permits exactly one terminal response. A late -completion after cancellation is recorded and discarded without producing a -second response. Dropping an unfinished record is a hard logged invariant -failure and must synthesize a typed terminal error when transport is still -available. +The table deliberately does not mirror task-execution states. Its publication +guard atomically permits exactly one terminal response or progress +acknowledgement. A late completion after cancellation is discarded without a +second publication. Dropping an unfinished record is a hard logged invariant +failure; bounded shutdown may take over an unretained publication right. Request IDs are unique within a connection, so the registry key is `(connection_id, request_id)`. A duplicate in-flight ID receives a typed @@ -280,8 +273,8 @@ terminal response so disposal cannot remove state underneath an active task. - [x] Add in-flight request count and byte configuration with validation. - [x] Implement count/byte admission reservations. -- [x] Implement the operation registry and duplicate-ID protection for ordinary - admitted work. Progress requests use the separate P0.3 registry. +- [x] Implement one operation table and duplicate-ID protection across ordinary + and progress admitted work, with independent class budgets. - [x] Implement tracked task completion and terminal-response guard. - [x] Convert ingress request handling from inline await to register-and-start. - [x] Preserve request ID and ownership on out-of-order completion. @@ -325,76 +318,42 @@ recreated the same serialization. Production `ExtensionContext` now owns an The target separates cloneable service handles from entity-owned mutable state. -### Process and connection state - -A small connection/session coordinator owns: - -- connection authentication and version state, -- session membership, -- VM membership indexes, -- request-operation ownership indexes, -- disposal state. - -Its commands must not perform guest execution, filesystem/network I/O, adapter -I/O, output writes, or unbounded waits. It may mutate indexes and return -cloneable entity handles. - -Connection/session disposal changes the entity state to `Closing` before -signalling owned operations. New requests for a closing entity are rejected. - -### VM state - -Each live VM has two complementary objects. Cloneable, thread-affine -`VmHandle(Rc>)` values provide short `try_read` and -`try_command` state sections. A `VmCoordinator` owns bounded operation and -lifecycle admission. Different handles and coordinators make progress -independently, and no state borrow crosses an external await. - -P0 ordering classes: - -```rust -enum RequestOrderingKey { - Connection(String), - Session { connection_id: String, session_id: String }, - VmLifecycle { connection_id: String, session_id: String, vm_id: String }, - VmOperation { connection_id: String, session_id: String, vm_id: String }, - Extension { - namespace: String, - connection_id: String, - key: Vec, - policy: ExtensionOrderingPolicy, - }, - Unordered, -} -``` - -Ordering semantics: - -- Authentication and connection/session membership changes use the relevant - connection/session coordinator. -- VM create is ordered with session disposal and VM membership insertion. -- Configure, dispose, layer topology changes, root snapshot/import/export, and - package linking use `VmLifecycle`. -- Filesystem, kernel, execution, stdin, process inspection/control, and VM fetch - use `VmOperation`. -- A VM lifecycle-exclusive command excludes VM operations for that VM. -- Ordinary VM operations are admitted independently but enter the VM - coordinator only for their state critical sections. -- Extension operations use an extension-provided opaque ordering key. Core does - not decode the extension payload. Core-exclusive keys reject an overlapping - operation with a typed bounded conflict; extensions may explicitly retain - conflict enforcement when their protocol requires a richer typed response. -- ACP uses its durable session ID as an extension-managed ordering key. Its - bounded route guard returns the existing typed `session_busy` policy for a - second active prompt in the same ACP session; prompts in different ACP - sessions run concurrently. - -A request-level operation may issue multiple VM commands over its lifetime. It -must not retain a mutable VM-state borrow or exclusive command guard between -commands. It does retain its nonexclusive ownership registration until terminal -completion. This lets teardown find the operation without preventing -`readFile` from completing while an ACP prompt waits on adapter output in the -same VM. +### Membership, operation ownership, and VM state + +One short-held metadata mutex owns the generation-bearing connection, session, +and VM membership tree. Entity records are `Open` or `Closing`; creation, +admission validation, and disposal linearize through this one lock. It never +protects guest execution, filesystem/network I/O, adapter I/O, output writes, +or an external wait. + +One bounded operation table is authoritative for every inbound request on a +connection. Ordinary and progress requests share `(connection_id, request_id)` +identity and one publication primitive, while retaining separate count and byte +budgets. Disposal closes a textual ownership scope and scans this bounded table +to cancel and drain matching work; there are no shadow connection/session/VM +operation maps. + +Each live VM has one independent lifecycle gate. Ordinary VM service calls take +nonexclusive permits and run concurrently. Configure, dispose, layer topology +changes, root snapshot/import/export, overlay creation, and package linking take +an exclusive lifecycle permit. `Idle -> Pending` immediately rejects later +ordinary admissions; the lifecycle becomes active only after earlier ordinary +and bounded internal-settlement permits drain. Internal events may settle while +pending, but remain durably deferred and cannot mutate the VM while lifecycle +work is active. + +Extension envelopes are opaque and take no generic core ordering key. ACP keeps +its bounded per-route state machine, so a second prompt on one ACP route receives +the existing typed `session_busy` response while prompts on different routes run +concurrently. A long prompt holds neither the membership lock nor a VM gate; +only its short VM service calls enter the gate. This is why same-VM `readFile` +and `writeFile` continue while the adapter prompt waits. + +The gate is explicit rather than an `RwLock`: lifecycle-pending work must be +rejected instead of queued invisibly, internal settlement has a separate bound, +and disposal needs cancellation generations and visible active counts. The +gate's mutex protects only non-suspending counter transitions and is released +before waiting. ### Event broker @@ -433,11 +392,11 @@ channel and the process-event broker. Service calls return owned `'static` futures or immediate results, communicate with coordinators, and never expose internal maps or stdio/fd details. -The generic `Extension` contract remains namespace-based and opaque. Its -implemented hooks are `request_class() -> ExtensionRequestClass::{Ordinary, -Progress}`, `request_ordering_key()`, and `request_ordering_policy()`. -Core supplies reserved progress admission; ACP alone decodes its opaque payload -and signals its keyed route state. +The generic `Extension` contract remains namespace-based and opaque. Its only +concurrency hook is +`request_class() -> ExtensionRequestClass::{Ordinary, Progress}`. Core supplies +reserved progress admission; ACP alone decodes its opaque payload and enforces +its per-route state machine. Native-sidecar must not import or decode `agentos-protocol`. @@ -460,7 +419,8 @@ progress blocker. - [x] Retain claimed internal event work durably when service admission is full; do not drop, duplicate, or hot-requeue it. - [x] Make extension request futures independently tracked and spawnable. -- [x] Add extension-owned opaque ordering/progress classification. +- [x] Add extension-owned progress classification while leaving route + exclusion inside ACP. - [x] Preserve all ownership validation at the coordinator boundary. - [x] Preserve explicit same-session ACP `session_busy` behavior. - [x] Prevent lifecycle operations from racing conflicting same-VM operations. @@ -512,16 +472,14 @@ to an unbounded collection. ### Generic extension progress routing ACP cancellation and permission responses remain encoded as opaque extension -requests. The extension classifies these frames through a generic hook. The -router uses the returned namespace plus opaque target key to find or signal the -active extension operation. +requests. The extension classifies these frames through one generic hook. The +router gives them reserved progress admission and invokes ACP without decoding +the payload or owning its route key. The implemented generic hooks are: ```rust fn request_class(...) -> ExtensionRequestClass; // Ordinary | Progress -fn request_ordering_key(...) -> Option>; -fn request_ordering_policy(...) -> ExtensionOrderingPolicy; ``` Core gives `Progress` independent reserved admission and invokes the opaque @@ -642,13 +600,14 @@ Tests: - [x] ACP test: one adapter has at most one response-loop consumer while different routes can have concurrent consumers. -Progress work uses its own reserved-lane-bounded, connection-scoped registry. -Its handle is retained until broker publication, claims exactly one -acknowledgement, and rejects a duplicate live request ID without affecting the -original. Shutdown closes ordinary admission, signals tracked work, continues -progress/control routing through the grace period, force-terminalizes any -unfinished operation exactly once, aborts only after takeover, drains control -output, and reports failed delivery. +Progress work is a class-specific view over the one operation table. It keeps +independent count/byte budgets but shares connection-scoped request identity and +the publication primitive with ordinary work. WriteStdin, CloseStdin, and +KillProcess use a separate bounded progress service channel so adapter +cancellation cannot queue behind ordinary extension services. Shutdown closes +ordinary admission, signals tracked work, continues progress/control routing +through the grace period, force-terminalizes unfinished work exactly once, +aborts only after takeover, drains control output, and reports failed delivery. ## P0.4 — Nonblocking output broker diff --git a/vm-lifecycle-gate-simplification-spec.md b/vm-lifecycle-gate-simplification-spec.md new file mode 100644 index 0000000000..f28e0f7cdb --- /dev/null +++ b/vm-lifecycle-gate-simplification-spec.md @@ -0,0 +1,639 @@ +# Native-sidecar ownership and VM lifecycle gate simplification + +Status: implemented and validated; unrelated repository-baseline gate failures +are recorded below + +Owner: agentOS native sidecar + +Baseline: the request-concurrency stack through +`refactor(native-sidecar): separate ownership from conflict policy` + +This document is the implementation contract and completion checklist. A box +may be checked only after the implementation exists and its corresponding test +passes. The implementing agent must record the commands and results in the +final validation section before declaring the work complete. + +## Objective + +Keep the concurrency behavior fixed by `request-concurrency-fix-prompt.md`, but +replace the overlapping request, ownership, conflict, and VM lifecycle +coordination machinery with the smallest explicit model that provides the same +safety. + +The completed design has: + +1. one bounded operation table for every inbound request on a connection, + including ordinary and progress requests; +2. one short-held ownership/membership state lock, never held across `.await`; +3. one small lifecycle gate per VM; +4. ACP-owned per-route exclusion, without generic core extension-ordering + machinery; and +5. no duplicate per-connection, per-session, and per-VM operation maps. + +The implementation must be a simplification. Production code in the ownership, +request-admission, and lifecycle subsystem should be materially smaller after +the change. Adding behavioral tests is expected and does not count against that +goal. + +## Why this change is needed + +The current implementation correctly prevents a long ACP prompt from holding +the protocol ingress loop, but represents one request in several overlapping +places: + +- `RequestOperationRegistry` or `ProgressRequestRegistry`; +- connection ownership operations; +- session ownership operations; +- VM ownership operations; +- extension conflict registrations; +- task/completion supervision; and +- terminal/progress publication guards. + +The ownership coordinator also combines four different concerns: + +- entity membership; +- disposal state; +- operation tracking; +- VM lifecycle exclusion. + +That makes cancellation and disposal hard to audit and caused ordinary and +progress request IDs to be checked in separate registries. For example, an +ordinary request `42` and progress request `42` can currently be admitted on +the same connection even though request IDs are connection-scoped. + +The lifecycle exclusion itself is necessary. The hierarchy around it is not. + +## Non-negotiable invariants + +- [x] Production ingress only decodes, validates, reserves, registers, and + starts work. It never awaits business execution, VM activity, output + capacity, adapter I/O, or event drainage. +- [x] Ordinary work is concurrent by default. +- [x] Different VMs never share a lifecycle gate. +- [x] A lifecycle mutation runs only after conflicting work already admitted + for that VM has drained. +- [x] Once lifecycle admission becomes pending, new ordinary work for that VM + is rejected with a typed, retryable conflict; it is not placed in a + hidden waiter queue. +- [x] Progress requests, registered sidecar responses, cancellation, permission + responses, shutdown, and transport failure do not require ordinary + request admission. +- [x] Progress-critical internal events needed to settle already-admitted work + can continue during lifecycle drain, remain bounded, and are included in + the drain count. +- [x] No internal event mutates a VM while its lifecycle gate is active. +- [x] A long ACP prompt does not hold the VM lifecycle gate. Only its short VM + service operations enter the gate. +- [x] Each `(connection_id, request_id)` identifies at most one live inbound + request, regardless of whether it is ordinary or progress-classified. +- [x] Every admitted request retains exactly one response-publication right and + releases all count/byte/gate accounting on every terminal path. +- [x] Entity disposal closes admission before signalling work, waits without + holding a state lock, and cannot remove state underneath an admitted + operation. +- [x] Every queue, table, retained byte collection, and waiter set remains + bounded with a typed error naming the limit and configuration path. +- [x] No mutex or `RefCell` borrow crosses `.await`. +- [x] No global `Arc>` or equivalent is introduced. +- [x] Native-sidecar remains opaque to ACP payloads. + +## Target architecture + +```text +fd 0 / fd 3 readers + | + v ++------------------------+ +| Protocol router | +| - classify | +| - reserve output | +| - admit operation | +| - start task | ++-----------+------------+ + | + v ++------------------------+ +| Operation table | one short metadata mutex +| - global request IDs | +| - ordinary/progress | +| - count/byte budgets | +| - ownership + phase | +| - cancellation | +| - drain notification | ++-----------+------------+ + | + | VM-scoped work only + v ++------------------------+ +| VmLifecycleGate | one independent gate per VM +| - ordinary permits | +| - internal permits | +| - lifecycle permit | ++------------------------+ +``` + +ACP route state remains inside the ACP extension. Output reservations and the +physical output broker remain separate from the operation table. + +## 1. One operation table + +Replace the separate ordinary and progress registries with one operation table +keyed by: + +```rust +struct OperationKey { + connection_id: String, + request_id: RequestId, +} +``` + +The table may retain separate configured budgets for ordinary and progress +classes, but duplicate-ID detection is global across both classes. + +The conceptual record is: + +```rust +struct OperationRecord { + generation: u64, + class: OperationClass, // Ordinary | Progress + ownership: OwnershipScope, + operation: String, + request_bytes: usize, + cancellation: OperationCancellation, + publication: ResponsePublicationGuard, +} +``` + +Requirements: + +- [x] Replace `RequestOperationRegistry` and `ProgressRequestRegistry` with one + authoritative map, or make one a thin class-specific view over a single + authoritative map. +- [x] Ordinary and progress budgets remain independently bounded. +- [x] Duplicate detection happens before either class reserves or starts work. +- [x] The same numeric request ID remains valid on different connections. +- [x] Replace `TerminalResponseGuard` and `ProgressAcknowledgementGuard` with + one response-publication primitive parameterized only by response class. +- [x] Keep exactly-once takeover for bounded shutdown and transport failure. +- [x] Remove operation lifecycle states that are used only to mirror task + execution. Retain only state that affects admission, cancellation, + publication correctness, or externally useful diagnostics. +- [x] Scope cancellation and drain scan the single bounded table. Do not add + secondary per-entity operation maps merely to avoid scanning at most the + configured in-flight operation bound. +- [x] Ordinary shutdown may close ordinary admission while progress admission + remains open for the bounded drain. +- [x] Table closure and entity closure use generations so an old disposal + cannot cancel a later entity reusing the same textual ID. + +## 2. Ownership and membership state + +Use one short-held metadata lock for connection, session, and VM membership. +This lock is allowed because it protects only bounded in-memory transitions and +is always released before external work or waiting. + +The simplest recommended representation places entity phases and the bounded +operation map under the same metadata mutex. If the implementation keeps +membership and operation storage behind separate mutexes, it must document one +lock order, make admission/disposal one linearizable transaction through a +generation-bearing lease, and add a deterministic test for the race. Merely +checking `Open`, releasing the membership lock, and registering an operation +later is not safe. + +The ownership state must not contain a second copy of every active operation. +The operation table is authoritative for active work and cancellation. + +Requirements: + +- [x] Entity records have an explicit generation and `Open | Closing` phase. +- [x] Request admission validates the complete connection/session/VM ownership + path against one coherent snapshot. +- [x] Entity creation and disposal are linearized against admission. +- [x] Disposal marks the entity `Closing` before signalling matching operation + records. +- [x] Disposal waits on operation-table and lifecycle-gate drain notification + without holding the membership lock. +- [x] Entity state is removed only after matching operations and gate permits + have drained. +- [x] A stale permit or completion from an earlier entity generation cannot + mutate a newly created entity with the same textual ID. +- [x] Delete `ConnectionOperationRegistration`, + `SessionOperationRegistration`, `VmOperationRegistration`, and their + independent operation-ID counters. +- [x] Delete nested lock sequences used solely to register the same operation + at connection, session, and VM scopes. +- [x] Do not reuse `runtime.protocol.maxInFlightRequests` as a session + membership limit. Reuse an existing semantically correct bound, remove a + redundant retained collection, or add a dedicated documented limit. + +## 3. Per-VM lifecycle gate + +Each live VM owns one independent `VmLifecycleGate`. The gate contains only the +state required to exclude lifecycle mutations from conflicting VM work. + +Conceptual API: + +```rust +impl VmLifecycleGate { + fn try_enter_ordinary( + &self, + cancellation: OperationCancellation, + ) -> Result; + + fn try_enter_internal( + &self, + cancellation: OperationCancellation, + ) -> Result; + + async fn begin_lifecycle( + &self, + cancellation: OperationCancellation, + ) -> Result; +} +``` + +Conceptual state: + +```rust +struct VmGateState { + phase: VmGatePhase, + closing: bool, + ordinary_active: usize, + internal_active: usize, + next_generation: u64, +} + +enum VmGatePhase { + Idle, + Pending { generation: u64 }, + Active { generation: u64 }, +} +``` + +Exact state semantics: + +```text +Idle + ordinary -> increment ordinary_active + internal -> increment internal_active + lifecycle -> Pending and close new ordinary admission + +Pending + new ordinary -> typed retryable conflict + bounded internal settlement work -> admitted or durably deferred + ordinary_active == 0 && internal_active == 0 -> Active + +Active + ordinary -> typed retryable conflict + internal -> remains in its durable source; never mutates the VM + matching lifecycle permit drops -> Idle + +closing == true (orthogonal to Idle/Pending/Active) + all new gate admission -> typed shutdown/disposal error + existing permits drain and notify disposal + cancellation/permit drop must never reopen admission +``` + +Requirements: + +- [x] `Idle -> Pending` is the linearization point that closes ordinary VM + admission. +- [x] Only one lifecycle request may be pending or active. A second receives a + typed conflict rather than waiting. +- [x] The lifecycle waiter registers notification before checking drain state, + preventing lost wakeups. +- [x] Lifecycle cancellation while pending restores `Idle` only when the + cancellation belongs to the current gate generation and the gate is not + closing. A closing gate remains unavailable. +- [x] Dropping an active lifecycle permit restores `Idle` exactly once when the + gate is open. It never clears the orthogonal closing state. +- [x] Dropping an ordinary or internal permit decrements exactly one counter + and wakes a pending lifecycle waiter when the gate may advance. +- [x] Counter underflow, stale generation, and poisoned-state recovery are hard + logged invariant failures. +- [x] Ordinary and internal admission remains bounded independently. +- [x] Internal work admitted during `Pending` is limited to the existing + settlement/event path. Public requests cannot label themselves internal. +- [x] Durable internal events are not removed from their source unless the + gate or a tracked deferred permit owns them. +- [x] A pending lifecycle operation cannot spin while internal capacity is + unavailable. +- [x] A lifecycle gate for VM A never reads, writes, or notifies VM B's gate. + +### Why this is not an `RwLock` + +The production doc comment on `VmLifecycleGate` must explain all of the +following: + +- a standard-library lock cannot cross `.await` without blocking a runtime + worker; +- ordinary admission must reject after lifecycle becomes pending rather than + join an implicit waiter queue; +- progress-critical internal settlement work must remain admissible during + `Pending` and must be counted in the lifecycle drain; +- disposal needs explicit cancellation, generations, and active counts; and +- the mutex inside the gate protects only non-suspending state transitions and + is released before waiting. + +The comment must also state that a short global metadata mutex is acceptable: +the forbidden design is a lock held across request execution, not a lock around +bounded map/counter transitions. + +## 4. Request classification + +Replace generalized conflict policy with one narrow VM concurrency class: + +```rust +enum VmConcurrencyClass { + None, + Ordinary, + ExclusiveLifecycle, +} +``` + +The names may differ, but the model must have no extension-specific conflict +variant. + +Classification requirements: + +- [x] Connection- and session-owned operations use `None`. +- [x] Extension envelope requests use `None`; ACP enforces its own route-level + response-loop exclusion. +- [x] Short VM service calls made by extensions use `Ordinary`. +- [x] Internal VM settlement/event work uses the dedicated internal gate API, + not `Ordinary` and not a public request class. +- [x] Ordinary core VM work uses `Ordinary`. +- [x] The following remain `ExclusiveLifecycle` unless an operation owner + provides a narrower, tested safety argument: + - dispose VM; + - bootstrap root filesystem; + - configure VM; + - create or seal layer; + - import or export snapshot; + - create overlay; + - snapshot root filesystem; + - link package. +- [x] Classification is implemented in one auditable location or on the + request type itself. Do not introduce a pairwise conflict matrix. +- [x] A newly added lifecycle request must require an explicit classification + in a compile-time exhaustive match or architecture test. + +## 5. Delete dead extension ordering machinery + +ACP is the only production extension currently returning an ordering key, and +it returns `ExtensionManaged`; the core coordinator therefore performs no +exclusion for that key. Retain ACP's route state and remove the unused generic +layer around it. + +- [x] Delete `ExtensionOrderingPolicy`. +- [x] Delete `Extension::request_ordering_key`. +- [x] Delete `Extension::request_ordering_policy`. +- [x] Delete `ConflictPolicy::Extension` or its replacement equivalent. +- [x] Delete connection-level `extension_conflicts` state and registration + guards. +- [x] Delete tests and architecture guards that require the removed hooks. +- [x] Preserve `Extension::request_class` so an opaque extension can identify + progress requests without native-sidecar decoding its payload. +- [x] Preserve ACP's per-route state machine and typed `session_busy` response. +- [x] Preserve concurrent prompts on different ACP routes. + +The following source audit must return no production matches: + +```bash +rg -n \ + 'ExtensionOrderingPolicy|request_ordering_key|request_ordering_policy|ConflictPolicy::Extension|extension_conflicts' \ + crates/native-sidecar/src crates/agentos-sidecar/src +``` + +## 6. Progress must remain end-to-end independent + +This refactor must not simplify away the reserved progress path. It must also +verify that progress is not reserved only at ingress/output while silently +joining an ordinary internal service queue. + +- [x] ACP cancel and permission response bypass ordinary request count/byte + admission and the VM ordinary gate. +- [x] Registered `SidecarResponseFrame` routing remains direct to its waiter. +- [x] Shutdown and terminal transport failure remain directly routable. +- [x] If ACP cancellation must issue `WriteStdin`, the service scheduler has + bounded progress-reserved admission or another direct bounded path. +- [x] Saturating ordinary extension-service work cannot prevent an admitted + cancellation from reaching the adapter and receiving its acknowledgement. +- [x] Progress saturation returns a typed limit error without consuming the + target operation's terminal response right. +- [x] Continuous progress traffic cannot starve already-retained terminal + responses. Use bounded fair scheduling rather than unlimited strict + priority if the existing output test exposes starvation. + +## 7. Required tests + +### Pure gate/state tests + +- [x] Multiple ordinary permits coexist. +- [x] Lifecycle becomes pending and waits for all earlier ordinary permits. +- [x] New ordinary admission is rejected while lifecycle is pending. +- [x] Lifecycle becomes active only after ordinary and internal counts reach + zero. +- [x] A second lifecycle request is rejected in both pending and active phases. +- [x] Cancelling a pending lifecycle operation reopens ordinary admission. +- [x] Dropping the active lifecycle permit reopens ordinary admission. +- [x] A stale lifecycle permit cannot reopen or mutate a later generation. +- [x] Internal settlement work can run during pending and is included in drain. +- [x] Internal work is not admitted while lifecycle is active. +- [x] Gate closure cancels/rejects admission and drains all permits. +- [x] Closing during pending and closing during active lifecycle work cannot + reopen admission when the lifecycle future or permit later drops. +- [x] Near-limit warnings and typed count-limit errors name the configuration + path. + +### Operation-table tests + +- [x] Ordinary request `42` followed by progress request `42` on the same + connection rejects the second request. +- [x] Progress request `42` followed by ordinary request `42` rejects the + second request. +- [x] Request `42` on two different connections remains valid. +- [x] Ordinary and progress class budgets are independent despite sharing the + ID table. +- [x] Cancellation versus natural completion publishes exactly one response. +- [x] Forced shutdown takeover publishes at most one response and releases all + accounting. +- [x] Closing a connection rejects later admission and cancels only matching + operations. +- [x] Scope drain reaches zero count and zero retained request bytes. + +### Real protocol-loop tests + +- [x] Same-VM read and write requests execute concurrently. +- [x] VM-A lifecycle drain does not delay VM-B ordinary work. +- [x] Configure/dispose waits for already-admitted same-VM work and rejects new + same-VM work while pending. +- [x] A sleeping ACP prompt does not retain the VM gate; same-VM file read/write + completes during the prompt. +- [x] Two ACP prompts on different routes run concurrently. +- [x] A second prompt on the same route receives ACP's typed `session_busy`. +- [x] ACP cancellation progresses while ordinary request admission is full. +- [x] ACP cancellation reaches adapter stdin while ordinary extension-service + capacity is saturated. +- [x] Disposal during an active prompt cancels, drains, and removes all route, + operation, gate, permission, and event-waiter state. +- [x] Repeated lifecycle/cancel/process-exit races produce no duplicate terminal + response, leaked permit, lost event, or hot spin. + +### Model/fuzz/load coverage + +- [x] Add a deterministic model test that compares randomized gate operations + against a small reference state machine. Cover admit/drop/cancel/close + and stale-generation sequences. +- [x] Extend protocol frame fuzzing with mixed ordinary/progress duplicate IDs, + lifecycle transitions, cancellation, and shutdown. +- [x] Add a bounded load test with multiple VMs, concurrent ordinary work, + periodic lifecycle requests, ACP progress, and deliberately delayed + completions. +- [x] The load test asserts a finite completion deadline, per-VM independence, + exactly-once responses, and zero final accounting rather than relying on + sleeps or log inspection. +- [x] Safeguard-firing saturation tests remain cheap and active in PR CI; tests + attempting to prove absence of a resource bound remain explicitly + ignored. + +## 8. Cleanup requirements + +- [x] Remove obsolete types, aliases, compatibility wrappers, error variants, + metrics, comments, and tests instead of leaving deprecated paths. +- [x] Do not retain both the old coordinator and new gate behind a feature flag. +- [x] Do not introduce a second request scheduler, runtime, or worker pool. +- [x] Do not add a general ordinary request FIFO. +- [x] Remove duplicated admission preflight where one atomic admission method + can reserve all required state safely. +- [x] Keep output reservations outside the VM gate; no gate permit may be held + while waiting for output capacity. +- [x] Keep public and Rust client behavior identical if any wire/config surface + changes. +- [x] Update `request-concurrency-fix-prompt.md` to describe the final simplified + ownership/gate model and remove the obsolete ordering-key description. +- [x] Update native-sidecar architecture comments and guards to enforce behavior + and forbidden dependencies, not incidental private symbol names. +- [x] `cargo fmt` and Clippy complete without new allows for dead code, complex + types, or too many arguments in the new subsystem. +- [x] The production ownership/request/gate implementation is net smaller than + the baseline, or the final validation record explains every net-new + production abstraction. + +## 9. Suggested implementation sequence + +1. [x] Add failing tests for cross-class duplicate IDs and the standalone VM + gate state machine. +2. [x] Introduce the single operation table and shared publication guard. +3. [x] Introduce `VmLifecycleGate` with ordinary, internal, and lifecycle RAII + permits. +4. [x] Rewire core VM requests and extension service commands to the gate. +5. [x] Move disposal cancellation/drain to the authoritative operation table. +6. [x] Remove per-entity operation maps and registration guards. +7. [x] Remove generic extension ordering hooks and conflict state. +8. [x] Run the real-loop ACP/filesystem/lifecycle regressions. +9. [x] Saturate the extension service path and close any progress-reservation + gap exposed by the test. +10. [x] Run model/fuzz/load coverage and the final source/architecture audits. +11. [x] Update the original concurrency design document and record final + validation below. + +The implementing agent may reorder mechanical steps, but must not temporarily +weaken the production invariants on the final revision. + +## 10. Validation commands + +Run focused checks first: + +```bash +cargo fmt --check +cargo test -p agentos-native-sidecar request_operations +cargo test -p agentos-native-sidecar ownership_coordinator +cargo test -p agentos-native-sidecar request_concurrency +cargo test -p agentos-native-sidecar --test architecture_guards +cargo test -p agentos-sidecar acp +pnpm --dir packages/core test:pr +``` + +Then run repository gates proportional to the final diff: + +```bash +cargo check --workspace +cargo clippy --workspace --all-targets -- -D warnings +pnpm check-types +pnpm build +``` + +If exact test filters change as obsolete modules are removed, replace them with +the new target names and record the actual commands below. Do not silently skip +a behavioral category because its old filter no longer matches tests. + +Required source audits: + +```bash +rg -n \ + 'ExtensionOrderingPolicy|request_ordering_key|request_ordering_policy|ConflictPolicy::Extension|extension_conflicts|ProgressRequestRegistry' \ + crates/native-sidecar/src crates/agentos-sidecar/src + +rg -n \ + 'ConnectionOperationRegistration|SessionOperationRegistration|VmOperationRegistration' \ + crates/native-sidecar/src +``` + +Both commands must return no production matches. Test-only reference-model +names should also be renamed so the audit remains unambiguous. + +## 11. Completion definition + +This work is complete only when: + +- [x] Every checkbox in sections 1 through 8 is complete. +- [x] The original prompt/filesystem/cancel reproduction passes. +- [x] All required focused validation passes; repository-wide gates either pass + or have a reproduced, unrelated baseline failure recorded below. +- [x] Source audits show no obsolete ownership or ordering implementation. +- [x] A reviewer can identify one authoritative operation table and one + lifecycle gate per VM without tracing duplicate shadow state. +- [x] The final diff contains no unrelated refactor or generated artifacts. +- [x] The working copy contains no accidental empty jj revision in the stack. +- [x] The implementation revision has a plain conventional-commit description + with no coding-agent attribution. + +## Final validation record + +Implementing agent fills this section in before handoff. + +- Implementation revision: `xulmxrrl` — + `refactor(native-sidecar): simplify VM lifecycle coordination` +- Final reviewer revision: `xulmxrrl` +- Focused Rust tests: `cargo fmt --check`; operation-table tests 20/20; + lifecycle-coordinator tests 11/11; request-concurrency tests 25/25; + architecture guards 41/41; ACP tests 52/52. +- ACP/public-client regression: the delayed response after 256 updates and the + two-prompts-plus-filesystem regression pass; all four native-sidecar migration + parity scenarios pass (6/6 tests across the two files). +- Model/fuzz/load tests: the deterministic randomized gate model, mixed + ordinary/progress protocol framing, duplicate-ID cases, lifecycle/cancel + races, bounded progress saturation, and multi-VM load deadline all pass in + the focused Rust suites. +- Workspace checks: `cargo check --workspace`, package-scoped Core and + runtime-core builds/typechecks, fixed-version verification, and publish + helper typechecks/tests pass. The repository-wide Clippy gate still reaches + unchanged `large_enum_variant` and disabled browser-target failures; root + `pnpm check-types` still reaches an unmaterialized example dependency; root + `pnpm build` still requires the generated Codex WASI release artifact; and + `packages/core test:pr` still has the unchanged top-level `pread` export + assertion (98/99 unit tests). These baseline failures are recorded in the + agentOS friction log; the focused changed-package and real-loop gates pass. +- Source audits: both required `rg` commands return no production matches. + `OperationTable` is the sole inbound request table and each `VmRecord` owns + one independent `VmLifecycleGate`. +- Production lines removed/added in the simplified subsystem: before test + modules, ownership coordination changed from 1,843 to 1,686 lines (-157) and + request operations from 1,484 to 1,524 (+40), for a net reduction of 117 + production lines. The request-operation increase is the shared progress + projection and cross-class identity/publication coverage replacing the + deleted second registry. +- Remaining known risks: no known change-specific correctness gap. The + unrelated repository-baseline gates above remain cleanup debt; the release + workflow remains the authoritative generated-artifact build.