diff --git a/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr.rs b/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr.rs index 0a81769ce6..7e4549a18d 100644 --- a/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr.rs +++ b/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr.rs @@ -21,6 +21,8 @@ use pplx_space_node_mgr_client::proto::{ use thiserror::Error; use tonic::Status; +use crate::node_manager_transport::NodeManagerConnector; + // `pub` only because `RestoreError::BtrfsReceive` carries it and `RestoreError` // is public; the module is `pub(crate)`, so the type is still not exported. #[derive(Debug, Error)] @@ -134,12 +136,14 @@ pub(crate) async fn encode_rootfs_send( content_hash: &str, snapshot_dir: &str, execution_fence: ExecutionFence, + node_manager: &NodeManagerConnector, ) -> Result { const RPC: &str = "BtrfsEncodeRootfsSend"; - let mut client = pplx_space_node_mgr_client::connect() + let mut client = node_manager + .connect() .await .map_err(|e| BtrfsViaNodemgrError::Connect(e.to_string()))?; - let defer = defer_send().await; + let defer = defer_send(node_manager).await; let response = client .btrfs_encode_rootfs_send(BtrfsEncodeRootfsSendRequest { sandbox_id: sandbox_id.to_string(), @@ -173,8 +177,10 @@ pub(crate) async fn encode_snapshot_rootfs_send( content_hash: &str, snapshot_name: &str, dest_dir: &str, + node_manager: &NodeManagerConnector, ) -> Result { - let mut client = pplx_space_node_mgr_client::connect() + let mut client = node_manager + .connect() .await .map_err(|e| BtrfsViaNodemgrError::Connect(e.to_string()))?; let response = client @@ -207,8 +213,10 @@ pub(crate) async fn receive_instance( template_name: &str, content_hash: &str, target: InstanceTarget, + node_manager: &NodeManagerConnector, ) -> Result<(), BtrfsViaNodemgrError> { - let mut client = pplx_space_node_mgr_client::connect() + let mut client = node_manager + .connect() .await .map_err(|e| BtrfsViaNodemgrError::Connect(e.to_string()))?; client @@ -232,8 +240,10 @@ pub(crate) async fn receive_instance( pub(crate) async fn delete_sandbox_instance( sandbox_id: &str, target: InstanceTarget, + node_manager: &NodeManagerConnector, ) -> Result<(), BtrfsViaNodemgrError> { - let mut client = pplx_space_node_mgr_client::connect() + let mut client = node_manager + .connect() .await .map_err(|e| BtrfsViaNodemgrError::Connect(e.to_string()))?; client @@ -262,8 +272,10 @@ pub(crate) async fn materialize_rootfs_send( sandbox_id: &str, artifact_dir: &std::path::Path, source_id: uuid::Uuid, + node_manager: &NodeManagerConnector, ) -> Result { - let mut client = pplx_space_node_mgr_client::connect() + let mut client = node_manager + .connect() .await .map_err(|e| BtrfsViaNodemgrError::Connect(e.to_string()))?; client @@ -289,11 +301,13 @@ pub(crate) async fn inspect_send_source( sandbox_id: &str, artifact_dir: &std::path::Path, source_id: uuid::Uuid, + node_manager: &NodeManagerConnector, ) -> Result { // Inspections are frequent and bounded; they share node-manager's lazy, // reconnecting channel instead of dialing the socket per call. A dead // socket surfaces as an RPC status, which proves nothing about the source. - pplx_space_node_mgr_client::lazy() + node_manager + .lazy() .btrfs_inspect_send_source(BtrfsInspectSendSourceRequest { sandbox_id: sandbox_id.to_owned(), artifact_dir: artifact_dir.to_string_lossy().into_owned(), @@ -361,7 +375,6 @@ impl CapabilityCache { } static VOLUME_MANAGER_STREAMS: CapabilityCache = CapabilityCache::new(); -static NODE_MANAGER_INSPECTS: CapabilityCache = CapabilityCache::new(); /// Whether volume-manager can stream an upload from a retained source. Gates /// existing retained artifacts too: without it they materialize first. @@ -395,15 +408,16 @@ async fn probe_volume_manager_streams_btrfs_send() -> bool { /// Whether node-manager can inspect retained sources and report their size. /// A deferred artifact is unverifiable without both. -async fn node_manager_inspects_send_sources() -> bool { - NODE_MANAGER_INSPECTS - .get(probe_node_manager_inspects_send_sources()) +async fn node_manager_inspects_send_sources(node_manager: &NodeManagerConnector) -> bool { + node_manager + .inspects_send_sources_cache() + .get(probe_node_manager_inspects_send_sources(node_manager)) .await } -async fn probe_node_manager_inspects_send_sources() -> bool { +async fn probe_node_manager_inspects_send_sources(node_manager: &NodeManagerConnector) -> bool { let probe = async { - let mut client = pplx_space_node_mgr_client::lazy(); + let mut client = node_manager.lazy(); client .get_btrfs_send_capabilities(GetBtrfsSendCapabilitiesRequest {}) .await @@ -421,16 +435,18 @@ async fn probe_node_manager_inspects_send_sources() -> bool { } } -async fn defer_send() -> bool { +async fn defer_send(node_manager: &NodeManagerConnector) -> bool { if !retain_send_sources() { return false; } let (streams, inspects) = tokio::join!( volume_manager_streams_btrfs_send(), - node_manager_inspects_send_sources() + node_manager_inspects_send_sources(node_manager) ); streams && inspects } +#[cfg(test)] +mod capability_cache_tests; #[cfg(test)] mod tests; diff --git a/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr/capability_cache_tests.rs b/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr/capability_cache_tests.rs new file mode 100644 index 0000000000..4557cf231e --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr/capability_cache_tests.rs @@ -0,0 +1,34 @@ +use super::*; +use crate::node_manager_transport::fake::{FakeNodeManager, ServedNodeManager}; +use pplx_space_node_mgr_client::proto::GetBtrfsSendCapabilitiesResponse; + +fn answering(inspect_send_source: bool) -> ServedNodeManager { + FakeNodeManager::default() + .unary( + "GetBtrfsSendCapabilities", + move |_: GetBtrfsSendCapabilitiesRequest| { + Ok(GetBtrfsSendCapabilitiesResponse { + supported: true, + inspect_send_source, + }) + }, + ) + .serve() +} + +/// A capability answer belongs to the node-manager that gave it: another +/// connector asks its own node-manager, and repeat questions within the cache +/// lifetime are answered without a second probe. +#[tokio::test] +async fn send_source_inspection_is_cached_per_connector() { + let inspecting = answering(true); + let refusing = answering(false); + + for _ in 0..2 { + assert!(node_manager_inspects_send_sources(&inspecting.connector()).await); + assert!(!node_manager_inspects_send_sources(&refusing.connector()).await); + } + + assert_eq!(inspecting.calls(), ["GetBtrfsSendCapabilities"]); + assert_eq!(refusing.calls(), ["GetBtrfsSendCapabilities"]); +} diff --git a/crates/libs/sbx/mgr/server/src/firecracker_via_nodemgr.rs b/crates/libs/sbx/mgr/server/src/firecracker_via_nodemgr.rs index 7b1ce9551f..a0369202c6 100644 --- a/crates/libs/sbx/mgr/server/src/firecracker_via_nodemgr.rs +++ b/crates/libs/sbx/mgr/server/src/firecracker_via_nodemgr.rs @@ -5,6 +5,8 @@ use pplx_space_node_mgr_client::proto::{ use pplx_space_node_runtime_core::fc_api::FirecrackerVmState; use std::time::Duration; +use crate::node_manager_transport::NodeManagerConnector; + const FIRECRACKER_VM_STATE_RPC_TIMEOUT: Duration = Duration::from_millis(500); #[derive(Debug, thiserror::Error)] @@ -20,10 +22,11 @@ pub(crate) enum FirecrackerStateReadError { async fn get_firecracker_vm_state_inner( sandbox_id: SandboxId, + node_manager: &NodeManagerConnector, ) -> Result { // Reuse the reconnecting channel; the outer read deadline also covers // channel readiness after a node-manager restart. - let mut client = pplx_space_node_mgr_client::lazy(); + let mut client = node_manager.lazy(); let response = client .get_firecracker_vm_state(GetFirecrackerVmStateRequest { sandbox_id: sandbox_id.to_string(), @@ -54,9 +57,10 @@ where pub(crate) async fn get_firecracker_vm_state( sandbox_id: SandboxId, + node_manager: &NodeManagerConnector, ) -> Result { bounded_firecracker_vm_state_read( - get_firecracker_vm_state_inner(sandbox_id), + get_firecracker_vm_state_inner(sandbox_id, node_manager), FIRECRACKER_VM_STATE_RPC_TIMEOUT, ) .await diff --git a/crates/libs/sbx/mgr/server/src/lib.rs b/crates/libs/sbx/mgr/server/src/lib.rs index e8715d0a1b..54498db7a1 100644 --- a/crates/libs/sbx/mgr/server/src/lib.rs +++ b/crates/libs/sbx/mgr/server/src/lib.rs @@ -16,6 +16,7 @@ mod guest_filesystem_flush; mod inventory_budget; mod lifecycle_metrics; pub mod local_artifacts; +pub mod node_manager_transport; mod node_state_sync; pub mod restore_operation; pub(crate) mod retained_source; @@ -34,6 +35,7 @@ pub mod snapshot_scheduler; pub mod snapshot_upload; pub mod snapshot_upload_jobs; pub mod state; +pub mod template_resolver; mod transfer_directory; pub mod watchdog; diff --git a/crates/libs/sbx/mgr/server/src/node_manager_transport.rs b/crates/libs/sbx/mgr/server/src/node_manager_transport.rs new file mode 100644 index 0000000000..710d78b538 --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/node_manager_transport.rs @@ -0,0 +1,93 @@ +//! How sandbox-manager obtains node-manager clients. +//! +//! Every node-manager client sandbox-manager uses comes from one +//! [`NodeManagerConnector`], which the service hands to the subsystems it +//! constructs. Production uses [`NodeManagerConnector::host`]: each method +//! calls the `pplx-space-node-mgr-client` constructor of the same name, so +//! every client dials node-manager's fixed Unix socket with that +//! constructor's retry policy and connection timing. Tests substitute a +//! transport serving an in-process node-manager. + +use std::sync::Arc; + +pub use pplx_space_node_mgr_client::{ConnectError, NodeManagerGrpcClient}; + +/// A source of node-manager clients. The methods mirror the client crate's +/// constructors; a caller picks the one whose dial timing it needs. +#[async_trait::async_trait] +pub trait NodeManagerTransport: Send + Sync { + /// A connected client, dialing under the default connect-retry policy. + async fn connect(&self) -> Result; + + /// A connected client after a single dial, for callers that own their + /// retry loop. + async fn connect_once(&self) -> Result; + + /// A client whose channel dials on first use and reconnects afterwards. + fn lazy(&self) -> NodeManagerGrpcClient; +} + +/// Node-manager's fixed host socket. +struct HostSocket; + +#[async_trait::async_trait] +impl NodeManagerTransport for HostSocket { + async fn connect(&self) -> Result { + pplx_space_node_mgr_client::connect().await + } + + async fn connect_once(&self) -> Result { + pplx_space_node_mgr_client::connect_once().await + } + + fn lazy(&self) -> NodeManagerGrpcClient { + pplx_space_node_mgr_client::lazy() + } +} + +/// Cloneable handle to the node-manager transport. +#[derive(Clone)] +pub struct NodeManagerConnector { + transport: Arc, + /// Whether this node-manager inspects retained btrfs-send sources. The + /// answer belongs to the node-manager the transport reaches, so it is + /// cached per connector and shared by its clones. + inspects_send_sources: Arc, +} + +impl NodeManagerConnector { + /// Node-manager at [`pplx_space_node_mgr_client::SOCKET_PATH`]. Every call + /// returns the same connector, so the process keeps one capability cache + /// for the host's node-manager however many places construct it. + pub fn host() -> Self { + static HOST: std::sync::LazyLock = + std::sync::LazyLock::new(|| NodeManagerConnector::new(Arc::new(HostSocket))); + HOST.clone() + } + + pub fn new(transport: Arc) -> Self { + Self { + transport, + inspects_send_sources: Arc::new(crate::btrfs_via_nodemgr::CapabilityCache::new()), + } + } + + pub async fn connect(&self) -> Result { + self.transport.connect().await + } + + pub async fn connect_once(&self) -> Result { + self.transport.connect_once().await + } + + pub fn lazy(&self) -> NodeManagerGrpcClient { + self.transport.lazy() + } + + pub(crate) fn inspects_send_sources_cache(&self) -> &crate::btrfs_via_nodemgr::CapabilityCache { + &self.inspects_send_sources + } +} + +#[cfg(test)] +pub(crate) mod fake; diff --git a/crates/libs/sbx/mgr/server/src/node_manager_transport/fake.rs b/crates/libs/sbx/mgr/server/src/node_manager_transport/fake.rs new file mode 100644 index 0000000000..0c562b8159 --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/node_manager_transport/fake.rs @@ -0,0 +1,178 @@ +//! An in-process node-manager for tests. +//! +//! The server listens on a socket in a test-owned temporary directory and +//! answers only the RPCs a test scripts. Every other call is refused with +//! `Unimplemented`, so code that regresses into an unexpected node-manager +//! call stops at the fake instead of reaching the host's node-manager. + +use std::collections::HashMap; +use std::marker::PhantomData; +use std::path::Path; +use std::sync::{Arc, Mutex}; + +use futures_util::future::BoxFuture; +use tokio_util::sync::CancellationToken; +use tonic::Status; + +use super::{ConnectError, NodeManagerConnector, NodeManagerGrpcClient, NodeManagerTransport}; + +const SERVICE_PATH: &str = "/space.node.v1.NodeManager/"; + +type Route = Arc< + dyn Fn( + http::Request, + ) -> BoxFuture<'static, http::Response> + + Send + + Sync, +>; + +/// The RPCs a [`FakeNodeManager`] answers, keyed by method name. +#[derive(Default)] +pub(crate) struct FakeNodeManager { + routes: HashMap<&'static str, Route>, +} + +impl FakeNodeManager { + /// Answer `method` (for example `"StopProcess"`) with `handler`. + pub(crate) fn unary(mut self, method: &'static str, handler: F) -> Self + where + Req: prost::Message + Default + Send + 'static, + Resp: prost::Message + Send + 'static, + F: Fn(Req) -> Result + Send + Sync + 'static, + { + let handler = Arc::new(handler); + let route: Route = Arc::new(move |request| { + let handler = Arc::clone(&handler); + Box::pin(async move { + tonic::server::Grpc::new(tonic_prost::ProstCodec::::default()) + .unary(Unary(handler, PhantomData), request) + .await + }) + }); + self.routes.insert(method, route); + self + } + + /// Start serving. The server stops when the returned handle drops. + pub(crate) fn serve(self) -> ServedNodeManager { + let sockets = tempfile::tempdir().unwrap(); + let socket = sockets.path().join("node-manager.sock"); + let listener = tokio::net::UnixListener::bind(&socket).unwrap(); + let routes = Arc::new(self.routes); + let calls = Arc::new(Mutex::new(Vec::new())); + let recorded = Arc::clone(&calls); + let rpc = tower::service_fn(move |request: http::Request| { + let method = request + .uri() + .path() + .strip_prefix(SERVICE_PATH) + .unwrap_or(request.uri().path()) + .to_owned(); + recorded.lock().unwrap().push(method.clone()); + let route = routes.get(method.as_str()).cloned(); + async move { + Ok::<_, std::convert::Infallible>(match route { + Some(route) => route(request).await, + None => { + Status::unimplemented(format!("fake node-manager does not answer {method}")) + .into_http() + } + }) + } + }); + let stop = CancellationToken::new(); + tokio::spawn( + tonic::transport::Server::builder().serve_with_incoming_shutdown( + rpc, + tokio_stream::wrappers::UnixListenerStream::new(listener), + stop.clone().cancelled_owned(), + ), + ); + ServedNodeManager { + connector: fixed_client(lazy_client(&socket)), + calls, + stop, + _sockets: sockets, + } + } +} + +/// A running [`FakeNodeManager`]. +pub(crate) struct ServedNodeManager { + connector: NodeManagerConnector, + calls: Arc>>, + stop: CancellationToken, + _sockets: tempfile::TempDir, +} + +impl ServedNodeManager { + /// A connector whose clients all reach this server. + pub(crate) fn connector(&self) -> NodeManagerConnector { + self.connector.clone() + } + + /// Method names of every call received, answered or refused, in arrival + /// order. + pub(crate) fn calls(&self) -> Vec { + self.calls.lock().unwrap().clone() + } +} + +impl Drop for ServedNodeManager { + fn drop(&mut self) { + self.stop.cancel(); + } +} + +/// A connector whose clients are all `client`, for tests that already serve a +/// node-manager of their own. +pub(crate) fn fixed_client(client: NodeManagerGrpcClient) -> NodeManagerConnector { + NodeManagerConnector::new(Arc::new(InProcess(client))) +} + +/// A connector whose clients all dial `socket`, a test-owned path nothing +/// binds, so every node-manager call fails before reaching any server. +pub(crate) fn unbound(socket: &Path) -> NodeManagerConnector { + fixed_client(lazy_client(socket)) +} + +fn lazy_client(socket: &Path) -> NodeManagerGrpcClient { + let channel = pplx_space_common_grpc::connect_uds_channel_lazy( + socket, + tonic::transport::Endpoint::from_static("http://[::]:50051"), + ); + NodeManagerGrpcClient::new(pplx_space_common_svc::bootstrap::wrap_grpc_client_channel( + channel, + )) +} + +struct InProcess(NodeManagerGrpcClient); + +#[async_trait::async_trait] +impl NodeManagerTransport for InProcess { + async fn connect(&self) -> Result { + Ok(self.0.clone()) + } + + async fn connect_once(&self) -> Result { + Ok(self.0.clone()) + } + + fn lazy(&self) -> NodeManagerGrpcClient { + self.0.clone() + } +} + +struct Unary(Arc, PhantomData Resp>); + +impl tonic::server::UnaryService for Unary +where + F: Fn(Req) -> Result, +{ + type Response = Resp; + type Future = std::future::Ready, Status>>; + + fn call(&mut self, request: tonic::Request) -> Self::Future { + std::future::ready((self.0)(request.into_inner()).map(tonic::Response::new)) + } +} diff --git a/crates/libs/sbx/mgr/server/src/retained_source.rs b/crates/libs/sbx/mgr/server/src/retained_source.rs index 4a08a9b779..a7f2876726 100644 --- a/crates/libs/sbx/mgr/server/src/retained_source.rs +++ b/crates/libs/sbx/mgr/server/src/retained_source.rs @@ -162,7 +162,9 @@ pub(crate) trait RetainedSourceInspector: Send + Sync { } /// One node-manager call with a bounded deadline. -pub(crate) struct NodeManagerSourceInspector; +pub(crate) struct NodeManagerSourceInspector( + pub(crate) crate::node_manager_transport::NodeManagerConnector, +); #[async_trait::async_trait] impl RetainedSourceInspector for NodeManagerSourceInspector { @@ -171,6 +173,7 @@ impl RetainedSourceInspector for NodeManagerSourceInspector { &source.sandbox_id, &source.artifact_dir, source.source_id, + &self.0, ); let (check, code) = match tokio::time::timeout(INSPECTION_DEADLINE, call).await { Ok(Ok(rootfs_bytes)) => (RetainedSourceCheck::Present { rootfs_bytes }, None), diff --git a/crates/libs/sbx/mgr/server/src/rootfs_encryption.rs b/crates/libs/sbx/mgr/server/src/rootfs_encryption.rs index bcf6c16eaa..f3ae4c57bc 100644 --- a/crates/libs/sbx/mgr/server/src/rootfs_encryption.rs +++ b/crates/libs/sbx/mgr/server/src/rootfs_encryption.rs @@ -52,8 +52,9 @@ pub(crate) async fn prepare_encrypted_rootfs( sandbox_id: &str, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, restore_cow_source: Option<&str>, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result { - let mut node = pplx_space_node_mgr_client::connect().await?; + let mut node = node_manager.connect().await?; let prepared = node .prepare_encrypted_rootfs( pplx_space_node_mgr_client::proto::PrepareEncryptedRootfsRequest { @@ -176,8 +177,9 @@ impl PoolRootfsPrepareError { pub(crate) async fn prepare_snapshot_rootfs( sandbox_id: &str, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result { - let mut node = pplx_space_node_mgr_client::connect().await?; + let mut node = node_manager.connect().await?; let prepared = node .prepare_snapshot_rootfs( pplx_space_node_mgr_client::proto::PrepareSnapshotRootfsRequest { diff --git a/crates/libs/sbx/mgr/server/src/run.rs b/crates/libs/sbx/mgr/server/src/run.rs index fad7a83446..2fb911c158 100644 --- a/crates/libs/sbx/mgr/server/src/run.rs +++ b/crates/libs/sbx/mgr/server/src/run.rs @@ -779,12 +779,9 @@ pub async fn run( .map_err(RunSbxMgrError::RouteVerificationBundle)?, )); - let health_server = start_takeover_health( - &config, - &shutdown_token, - pplx_space_node_mgr_client::connect_once, - ) - .await?; + let node_manager = crate::node_manager_transport::NodeManagerConnector::host(); + let health_server = + start_takeover_health(&config, &shutdown_token, || node_manager.connect_once()).await?; let health_socket = config .health_socket_path .as_ref() @@ -851,6 +848,7 @@ pub async fn run( node_role, node_managed_by, route_ed25519_verification_bundle, + node_manager: node_manager.clone(), ..Default::default() }, service_config, @@ -864,23 +862,20 @@ pub async fn run( ) .map_err(RunSbxMgrError::EnableDependencyRestartDrain) .inspect_err(report_active_failure)?; - let startup_recovery_lease = match admit_startup_recovery( - &service, - &shutdown_token, - pplx_space_node_mgr_client::connect_once, - ) - .await - .inspect_err(report_active_failure)? - { - Some(lease) => lease, - None => { - tracing::warn!("terminated while waiting for startup recovery"); - health_server.join().await; - active_ownership.exit_without_quiescence( - "sandbox-manager cancelled after construction without a drain proof", - ); - } - }; + let startup_recovery_lease = + match admit_startup_recovery(&service, &shutdown_token, || node_manager.connect_once()) + .await + .inspect_err(report_active_failure)? + { + Some(lease) => lease, + None => { + tracing::warn!("terminated while waiting for startup recovery"); + health_server.join().await; + active_ownership.exit_without_quiescence( + "sandbox-manager cancelled after construction without a drain proof", + ); + } + }; service .prepare_execution_authority() .await diff --git a/crates/libs/sbx/mgr/server/src/runtime_reconciliation.rs b/crates/libs/sbx/mgr/server/src/runtime_reconciliation.rs index 35ab22d98b..9a97668a1c 100644 --- a/crates/libs/sbx/mgr/server/src/runtime_reconciliation.rs +++ b/crates/libs/sbx/mgr/server/src/runtime_reconciliation.rs @@ -86,8 +86,10 @@ where pub(crate) async fn observe_firecracker_vm_state( sandbox_id: &SandboxId, probe_context: &'static str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> FirecrackerVmStateObservation { - match crate::firecracker_via_nodemgr::get_firecracker_vm_state(*sandbox_id).await { + match crate::firecracker_via_nodemgr::get_firecracker_vm_state(*sandbox_id, node_manager).await + { Ok(state) => FirecrackerVmStateObservation::Observed(state), Err(error) => { let reason = format!("Firecracker state probe failed: {error}"); diff --git a/crates/libs/sbx/mgr/server/src/service/activation_support.rs b/crates/libs/sbx/mgr/server/src/service/activation_support.rs index 7949bc843d..c82c18b7ae 100644 --- a/crates/libs/sbx/mgr/server/src/service/activation_support.rs +++ b/crates/libs/sbx/mgr/server/src/service/activation_support.rs @@ -1619,6 +1619,7 @@ pub(super) fn spawn_restore_state_transition( cleanup_id: SandboxCleanupEpisodeId, ) -> Result<(), Status> { let db = Arc::clone(&service.db); + let node_manager = service.node_manager.clone(); service.spawn_state_writer("RestoreFinalizer", async move { match complete_restore_and_mark_final( db.clone(), @@ -1631,7 +1632,7 @@ pub(super) fn spawn_restore_state_transition( { Ok(()) => { if let Some(snapshot_dir) = snapshot_dir { - cleanup_restored_snapshot_dir_async(&sandbox_id, &snapshot_dir); + cleanup_restored_snapshot_dir_async(&sandbox_id, &snapshot_dir, &node_manager); } cleanup_guard.complete_success(); } diff --git a/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs index b1f910ede5..792ae03797 100644 --- a/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs +++ b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs @@ -816,7 +816,7 @@ impl SandboxManagerService { } self.watchdog.record_lifecycle_edge(&sandbox.id).await; - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|error| resident_resume_connect_error(&self.db, &sandbox.id, error))?; @@ -869,6 +869,7 @@ impl SandboxManagerService { let observation = observe_firecracker_vm_state( &sandbox_id_typed, "resident resume response reconciliation", + &self.node_manager, ) .await; let reason = match observation.unobservable_reason() { @@ -926,6 +927,7 @@ impl SandboxManagerService { let observation = observe_firecracker_vm_state( &sandbox_id_typed, "resident resume guest setup rollback", + &self.node_manager, ) .await; let probe_detail = observation @@ -1003,13 +1005,15 @@ impl SandboxManagerService { } if inherit_template_default { - let resolved_template = - resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths( + self.template_resolver.as_ref(), + pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; policy = resolved_template.default_auto_lifecycle; } diff --git a/crates/libs/sbx/mgr/server/src/service/cache_cleanup.rs b/crates/libs/sbx/mgr/server/src/service/cache_cleanup.rs index 0ab2488d3c..57ba861f42 100644 --- a/crates/libs/sbx/mgr/server/src/service/cache_cleanup.rs +++ b/crates/libs/sbx/mgr/server/src/service/cache_cleanup.rs @@ -1387,10 +1387,7 @@ impl SandboxManagerService { fn orphan_execution_observer( &self, ) -> impl Fn(SandboxId) -> Result, Status> + use<> { - orphan_execution_observer( - pplx_space_node_mgr_client::lazy(), - self.shutdown_token.clone(), - ) + orphan_execution_observer(self.node_manager.lazy(), self.shutdown_token.clone()) } pub(super) fn reconcile_sandbox_inventory( diff --git a/crates/libs/sbx/mgr/server/src/service/cleanup.rs b/crates/libs/sbx/mgr/server/src/service/cleanup.rs index aca1fd7951..571919ddc5 100644 --- a/crates/libs/sbx/mgr/server/src/service/cleanup.rs +++ b/crates/libs/sbx/mgr/server/src/service/cleanup.rs @@ -563,6 +563,7 @@ impl SandboxManagerService { pplx_space_node_mgr_client::execution_fence::InstanceTarget::RestoreScope( transfer_scope_id, ), + &self.node_manager, ) .await .map_err(|error| { @@ -2204,7 +2205,7 @@ impl SandboxManagerService { admission, cleanup_id, || async move { - connect_node().await.map_err(|error| { + connect_node(&self.node_manager).await.map_err(|error| { Status::internal(format!( "{context}: cannot connect to node-manager: {error}" )) @@ -2526,7 +2527,7 @@ impl SandboxManagerService { || async move { self.flush_guest_before_termination(dispatch_sandbox_id, dispatch_guest_flush) .await; - let mut node_client = connect_node().await.map_err(|error| { + let mut node_client = connect_node(&self.node_manager).await.map_err(|error| { Status::internal(format!( "{dispatch_context}: cannot connect to node-manager: {error}" )) @@ -2772,7 +2773,8 @@ impl SandboxManagerService { /// [`BeforeStopFlush`]. async fn resume_paused_vm_before_flush(&self, sandbox_id: &str) -> Result<(), StopFlushError> { let id = SandboxId::from_str(sandbox_id).map_err(StopFlushError::SandboxId)?; - let state = crate::firecracker_via_nodemgr::get_firecracker_vm_state(id).await; + let state = + crate::firecracker_via_nodemgr::get_firecracker_vm_state(id, &self.node_manager).await; if let Err(error) = &state { tracing::debug!(%sandbox_id, %error, "VM state unreadable before the stop flush"); } @@ -2786,7 +2788,9 @@ impl SandboxManagerService { .map_err(StopFlushError::Fence)? .map(Into::into) }; - let mut node_client = connect_node().await.map_err(StopFlushError::NodeManager)?; + let mut node_client = connect_node(&self.node_manager) + .await + .map_err(StopFlushError::NodeManager)?; node_client .resume_sandbox(ResumeSandboxRequest { execution_fence, diff --git a/crates/libs/sbx/mgr/server/src/service/cleanup_executor/tests.rs b/crates/libs/sbx/mgr/server/src/service/cleanup_executor/tests.rs index 7e60e9a4c8..ae71d8450e 100644 --- a/crates/libs/sbx/mgr/server/src/service/cleanup_executor/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/cleanup_executor/tests.rs @@ -29,23 +29,36 @@ fn work(conn: &rusqlite::Connection, kill: bool) -> (JobId, CleanupWork) { #[tokio::test] async fn blocking_stop_follows_a_retryable_attempt_to_durable_completion() { + use pplx_space_node_mgr_client::proto::{CleanupSandboxRequest, CleanupSandboxResponse}; + + let cleanups = Arc::new(Mutex::new(Vec::new())); + let node_manager = crate::node_manager_transport::fake::FakeNodeManager::default() + .unary("CleanupSandbox", { + let cleanups = Arc::clone(&cleanups); + move |request: CleanupSandboxRequest| { + let mut cleanups = cleanups.lock().unwrap(); + cleanups.push(request); + if cleanups.len() == 1 { + Err(Status::unavailable("runtime still exiting")) + } else { + Ok(CleanupSandboxResponse::default()) + } + } + }) + .serve(); let conn = crate::db::test_conn(); let sandbox_id = SandboxId::new(); let subject = sandbox_id.to_string(); insert_test_sandbox(&conn, &subject, SandboxState::Stopping); - let service = service_for_teardown_test(conn); - let calls = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let attempts = Arc::clone(&calls); - *service.cleanup_executor.dispatch.lock().unwrap() = Some(Arc::new(move |_, _| { - let first = attempts.fetch_add(1, Ordering::SeqCst) == 0; - Box::pin(async move { - if first { - Err(Status::unavailable("runtime still exiting")) - } else { - Ok(()) - } - }) - })); + let sockets = tempfile::tempdir().unwrap(); + let service = crate::service::tests::node_manager_transport_tests::service_with_node_manager( + conn, + node_manager.connector(), + sockets.path(), + ); + let execution_fence = + crate::db::node_execution_fence(&service.db.lock().unwrap(), &subject).unwrap(); + assert!(execution_fence.is_some()); let label = SandboxRuntimeLabel::from_sandbox_id(sandbox_id); let request = SynchronousCleanupRequest { sandbox_id: &subject, @@ -64,7 +77,17 @@ async fn blocking_stop_follows_a_retryable_attempt_to_durable_completion() { .await .unwrap() .unwrap(); - assert_eq!(calls.load(Ordering::SeqCst), 2); + // Each attempt asks node-manager for the stop the caller requested. + let requested = CleanupSandboxRequest { + execution_fence: execution_fence.map(Into::into), + sandbox_id: subject.clone(), + force: true, + keep_filesystem: true, + return_after_slot_release: true, + immediate_kill: true, + }; + assert_eq!(*cleanups.lock().unwrap(), [requested.clone(), requested]); + assert_eq!(node_manager.calls(), ["CleanupSandbox", "CleanupSandbox"]); assert_eq!( crate::db::get_sandbox(&service.db.lock().unwrap(), &subject) .unwrap() diff --git a/crates/libs/sbx/mgr/server/src/service/create.rs b/crates/libs/sbx/mgr/server/src/service/create.rs index 6f97cecdc5..e46a361d44 100644 --- a/crates/libs/sbx/mgr/server/src/service/create.rs +++ b/crates/libs/sbx/mgr/server/src/service/create.rs @@ -71,6 +71,7 @@ async fn reconstruct_encrypted_restore_devices( sandbox_id: &str, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, snapshot_dir: &str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result<(String, String), Status> { let restore_source_if_present = |file: &str| -> Option { let path = format!("{snapshot_dir}/{file}"); @@ -83,6 +84,7 @@ async fn reconstruct_encrypted_restore_devices( sandbox_id, execution_fence, cow_source.as_deref(), + node_manager, ) .await .map_err(encrypted_restore_rootfs_error_status)?; @@ -91,6 +93,7 @@ async fn reconstruct_encrypted_restore_devices( Some(execution_fence.into()), false, snapshot_restore_source.as_deref(), + node_manager, ) .await .map_err(encrypted_restore_snapshot_error_status)? @@ -337,7 +340,7 @@ impl SandboxManagerService { }; if req.autostart { - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; self.reject_autostart_create_without_ready_slot(&mut node_client) @@ -451,7 +454,7 @@ impl SandboxManagerService { } let status = match base_snapshot_memory_file_mib(snapshot_dir_resolved) { - Ok(snapshot_memory_mib) => match connect_node().await { + Ok(snapshot_memory_mib) => match connect_node(&self.node_manager).await { Ok(mut node_client) => { match Self::node_sandbox_memory_policy(&mut node_client).await { Ok(node_memory_policy) => ensure_snapshot_memory_allowed( @@ -625,7 +628,7 @@ impl SandboxManagerService { "restore.template_prepare_source", START_TEMPLATE_RESOLUTION_TIMEOUT, id, - template_source_for_sandbox_prepare(sandbox), + template_source_for_sandbox_prepare(self.template_resolver.as_ref(), sandbox), ) .await?; let prepare_template = prepare_source.template; @@ -656,7 +659,7 @@ impl SandboxManagerService { START_NODE_METADATA_TIMEOUT, id, async { - connect_node() + connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}"))) }, @@ -931,12 +934,17 @@ impl SandboxManagerService { // `encrypted` marker, which equals it by invariant); plaintext // sandboxes take the plaintext path with no service round-trip. let (rootfs_block_device_path, snapshot_volume_mount_dir) = if sandbox.volume_encryption { - match reconstruct_encrypted_restore_devices(id, execution_fence, &snapshot_dir) - .instrument(tracing::info_span!( - "sbxmgr.restore.encrypted_devices_prepare", - sandbox_id = %id - )) - .await + match reconstruct_encrypted_restore_devices( + id, + execution_fence, + &snapshot_dir, + &self.node_manager, + ) + .instrument(tracing::info_span!( + "sbxmgr.restore.encrypted_devices_prepare", + sandbox_id = %id + )) + .await { Ok((device_path, mount_dir)) => (Some(device_path), Some(mount_dir)), Err(status) => { @@ -1168,7 +1176,7 @@ impl SandboxManagerService { resume_err = Some(e); if attempt < 2 { // Reconnect to node-manager in case the connection was reset - if let Ok(new_client) = connect_node().await { + if let Ok(new_client) = connect_node(&self.node_manager).await { node_client = new_client; } tokio::time::sleep(std::time::Duration::from_millis(200)).await; @@ -1320,7 +1328,7 @@ impl SandboxManagerService { return Err(activation_cleanup.failure(error)); } if sandbox.pool_template.is_none() { - cleanup_restored_snapshot_dir_async(id, &snapshot_dir); + cleanup_restored_snapshot_dir_async(id, &snapshot_dir, &self.node_manager); } activation_cleanup.complete_success(); } else { diff --git a/crates/libs/sbx/mgr/server/src/service/create_support.rs b/crates/libs/sbx/mgr/server/src/service/create_support.rs index d8365eaad1..1b00b5af80 100644 --- a/crates/libs/sbx/mgr/server/src/service/create_support.rs +++ b/crates/libs/sbx/mgr/server/src/service/create_support.rs @@ -564,21 +564,29 @@ pub(super) fn checkpoint_start_response_for_existing_started_sandbox( } pub(super) async fn resolve_template_paths( + templates: &dyn crate::template_resolver::TemplateResolver, template_ref: pplx_space_template_mgr_client::proto::TemplateRef, ) -> Result { - let mut tmpl_client = pplx_space_template_mgr_client::connect() + use crate::template_resolver::TemplateLookupError; + let inner = templates + .get_template(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(template_ref), + vcpu: None, + memory_mib: None, + }, + )) .await - .map_err(|e| tonic::Status::internal(format!("cannot connect to template-manager: {e}")))?; - let resp = tmpl_client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(template_ref), - vcpu: None, - memory_mib: None, - }) - .await - .map_err(|e| tonic::Status::new(e.code(), format!("GetTemplate failed: {e}")))?; - let inner = resp.into_inner(); + .map_err(|error| match error { + TemplateLookupError::Connect(e) => { + tonic::Status::internal(format!("cannot connect to template-manager: {e}")) + } + TemplateLookupError::Rpc(e) => { + tonic::Status::new(e.code(), format!("GetTemplate failed: {e}")) + } + })? + .into_inner(); let tmpl = inner.template.unwrap_or_default(); let arch = template_metadata_arch(&inner.template_dir).unwrap_or(inner.arch); let spaced_compatibility = template_metadata_spaced_compatibility(&inner.template_dir); @@ -670,6 +678,7 @@ pub(super) fn stage_base_snapshot_runtime_metadata( } pub(super) async fn resolve_template_paths_for_start( + templates: &dyn crate::template_resolver::TemplateResolver, template_ref: pplx_space_template_mgr_client::proto::TemplateRef, ) -> Result { // Start restores an existing sandbox's identity. A name-only lookup could @@ -683,7 +692,7 @@ pub(super) async fn resolve_template_paths_for_start( "cannot recover sandbox artifact without its saved template hash", )); } - resolve_template_paths(template_ref) + resolve_template_paths(templates, template_ref) .await .map_err(start_template_resolution_failed_status) } @@ -701,6 +710,7 @@ fn start_template_resolution_failed_status(error: Status) -> Status { } pub(super) async fn template_runtime_defaults_for_sandbox_start( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox: &crate::db::SandboxRow, prepared_template_dir: Option<&str>, ) -> TemplateRuntimeDefaults { @@ -708,11 +718,14 @@ pub(super) async fn template_runtime_defaults_for_sandbox_start( return template_runtime_defaults(template_dir); } - match resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { - name: sandbox.template.clone(), - content_hash: sandbox.template_hash.clone(), - tag: None, - }) + match resolve_template_paths( + templates, + pplx_space_template_mgr_client::proto::TemplateRef { + name: sandbox.template.clone(), + content_hash: sandbox.template_hash.clone(), + tag: None, + }, + ) .await { Ok(resolved) => resolved.runtime_defaults, @@ -742,6 +755,7 @@ pub(super) fn sandbox_rootfs_exists_in(cache_dir: &std::path::Path, sandbox_id: } pub(super) async fn template_source_for_sandbox_prepare( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox: &crate::db::SandboxRow, ) -> Result { if sandbox_rootfs_exists(&sandbox.id) { @@ -764,13 +778,15 @@ pub(super) async fn template_source_for_sandbox_prepare( ))); } - let resolved_template = - resolve_template_paths_for_start(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths_for_start( + templates, + pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; reject_incompatible_template_arch( &sandbox.template, resolved_template diff --git a/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs b/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs index 2790329b4e..425243bf02 100644 --- a/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/create_support/tests.rs @@ -476,17 +476,23 @@ fn restore_provenance_classifies_legacy_remote_metadata() { #[tokio::test] async fn existing_sandbox_cannot_resolve_a_missing_artifact_by_latest_alias() { - // With no template-manager available, this must reject locally rather than - // asking it to pick the newest build after the sandbox lost its identity. - let error = - resolve_template_paths_for_start(pplx_space_template_mgr_client::proto::TemplateRef { + // This must reject locally rather than asking template-manager to pick + // the newest build after the sandbox lost its identity. + let templates = crate::template_resolver::fake::FakeTemplateResolver::refusing( + tonic::Status::not_found("a latest-alias lookup must not be sent"), + ); + let error = resolve_template_paths_for_start( + &templates, + pplx_space_template_mgr_client::proto::TemplateRef { name: "app-runtime".into(), content_hash: None, tag: None, - }) - .await - .err() - .expect("missing identity cannot select a replacement artifact"); + }, + ) + .await + .err() + .expect("missing identity cannot select a replacement artifact"); assert_eq!(error.code(), tonic::Code::FailedPrecondition); assert!(error.message().contains("saved template hash")); + assert!(templates.requests().is_empty()); } diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs index e9a386455e..bb018178bf 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_core.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs @@ -197,13 +197,15 @@ impl SandboxManagerService { // Resolve template to content-addressed paths via template-manager. // If caller pinned a content_hash, use it; otherwise resolve newest. - let resolved_template = - resolve_template_paths(pplx_space_template_mgr_client::proto::TemplateRef { + let resolved_template = resolve_template_paths( + self.template_resolver.as_ref(), + pplx_space_template_mgr_client::proto::TemplateRef { name: template.to_string(), content_hash: req.content_hash.clone(), tag: None, - }) - .await?; + }, + ) + .await?; if let Some(policy) = template_network_policy_from_proto( req.template_network_policy.as_ref(), req.template_network_policy_generation, @@ -1010,7 +1012,7 @@ impl SandboxManagerService { let mut usage_bytes = std::collections::HashMap::new(); if !resolved_ids.is_empty() { - let mut client = connect_node() + let mut client = connect_node(&self.node_manager) .await .map_err(|e| Status::unavailable(format!("cannot connect to node-manager: {e}")))?; let response = client @@ -1176,5 +1178,7 @@ fn project_template_catalog( .collect() } +#[cfg(test)] +mod create_template_lookup_tests; #[cfg(test)] mod tests; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs new file mode 100644 index 0000000000..9e28442e0e --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs @@ -0,0 +1,178 @@ +//! Create's auto-lifecycle authority checks, which run just before template +//! lookup, against a scripted template resolver. Template-manager and +//! node-manager are fakes, so a request that gets past a check stops at the +//! resolver's scripted answer. + +use super::*; +use crate::node_manager_transport::fake::FakeNodeManager; +use crate::template_resolver::TemplateLookupError; +use crate::template_resolver::fake::{FakeTemplateResolver, LookupChannel}; +use pplx_space_common_core::host::linux::identities::users; +use pplx_space_common_svc::PeerIdentity; +use pplx_space_net_common::proto::PolicyAuthority; +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, TemplateRef}; + +/// The resolver's refusal, distinct from every refusal a check produces before +/// lookup. +const LOOKUP_REFUSAL: &str = "scripted template lookup refusal"; + +fn refusing_lookup() -> FakeTemplateResolver { + FakeTemplateResolver::refusing(Status::not_found(LOOKUP_REFUSAL)) +} + +fn unreachable_template_manager() -> FakeTemplateResolver { + FakeTemplateResolver::new(|| { + let invalid = "http://[".parse::().unwrap_err(); + Err(TemplateLookupError::Connect( + pplx_space_template_mgr_client::ConnectError::InvalidUri(invalid), + )) + }) +} + +#[tokio::test] +async fn create_checks_auto_lifecycle_authority_before_template_lookup() { + if !crate::service::tests::enter_route_test_process(&[ + ("EPPO_SDK_KEY", None), + ("SPACE_FLAGD_ENABLED", None), + ]) { + return; + } + let forwarded = PeerIdentity::new( + users::SPACE_NODEAPI.uid, + users::SPACE_NODEAPI.primary_group.gid, + None, + ); + let local = PeerIdentity::new(0, 0, None); + struct Case { + name: &'static str, + peer: PeerIdentity, + authority: i32, + templates: fn() -> FakeTemplateResolver, + code: tonic::Code, + message_prefix: &'static str, + message_contains: &'static str, + looks_up_template: bool, + } + let cases = [ + Case { + name: "local authority from a forwarded peer", + peer: forwarded.clone(), + authority: PolicyAuthority::Local as i32, + templates: refusing_lookup, + code: tonic::Code::PermissionDenied, + message_prefix: "authority=local requires a local operator peer", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "unspecified authority", + peer: forwarded.clone(), + authority: PolicyAuthority::Unspecified as i32, + templates: refusing_lookup, + code: tonic::Code::InvalidArgument, + message_prefix: "policy authority proto value is unspecified", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "unknown authority", + peer: forwarded.clone(), + authority: 99, + templates: refusing_lookup, + code: tonic::Code::InvalidArgument, + message_prefix: "policy authority proto value 99 is unknown", + message_contains: "", + looks_up_template: false, + }, + Case { + name: "cluster authority from a forwarded peer", + peer: forwarded.clone(), + authority: PolicyAuthority::Cluster as i32, + templates: refusing_lookup, + code: tonic::Code::NotFound, + message_prefix: "GetTemplate failed: ", + message_contains: LOOKUP_REFUSAL, + looks_up_template: true, + }, + Case { + name: "local authority from a local peer", + peer: local, + authority: PolicyAuthority::Local as i32, + templates: refusing_lookup, + code: tonic::Code::NotFound, + message_prefix: "GetTemplate failed: ", + message_contains: LOOKUP_REFUSAL, + looks_up_template: true, + }, + Case { + name: "accepted authority with template-manager unreachable", + peer: forwarded, + authority: PolicyAuthority::Cluster as i32, + templates: unreachable_template_manager, + code: tonic::Code::Internal, + message_prefix: "cannot connect to template-manager: ", + message_contains: "invalid gRPC endpoint URI", + looks_up_template: true, + }, + ]; + for case in cases { + let templates = Arc::new((case.templates)()); + let node_manager = FakeNodeManager::default().serve(); + let sockets = tempfile::tempdir().unwrap(); + let service = crate::service::tests::node_manager_transport_tests::isolated_service( + crate::db::test_conn(), + node_manager.connector(), + Arc::clone(&templates) as Arc, + sockets.path(), + ); + let sandbox_id = SandboxId::new().to_string(); + let mut request = Request::new(CreateRequest { + template: "seed".into(), + owner: "root@local".into(), + runtime: pplx_space_common_core::types::RUNTIME_FIRECRACKER_ALIAS.into(), + boot_mode: BootMode::Cold as i32, + sandbox_id: Some(sandbox_id.clone()), + auto_lifecycle: Some(Default::default()), + auto_lifecycle_authority: case.authority, + ..Default::default() + }); + request.extensions_mut().insert(case.peer); + + let error = service.create(request).await.expect_err(case.name); + + assert_eq!(error.code(), case.code, "{}: {error}", case.name); + assert!( + error.message().starts_with(case.message_prefix) + && error.message().contains(case.message_contains), + "{}: {error}", + case.name + ); + let lookups = templates.lookups(); + let expected_lookups = if case.looks_up_template { + vec![GetTemplateRequest { + caller_org_id: None, + template: Some(TemplateRef { + name: "seed".into(), + content_hash: None, + tag: None, + }), + vcpu: None, + memory_mib: None, + }] + } else { + Vec::new() + }; + assert_eq!(templates.requests(), expected_lookups, "{}", case.name); + // Create dials template-manager for its lookup and sends no metadata. + for lookup in &lookups { + assert_eq!(lookup.channel, LookupChannel::Dialed, "{}", case.name); + assert!(lookup.metadata.is_empty(), "{}", case.name); + } + assert!(node_manager.calls().is_empty(), "{}", case.name); + assert!( + !crate::db::sandbox_exists(&service.db_conn().unwrap(), &sandbox_id).unwrap(), + "{}", + case.name + ); + } +} diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs index b3b7760dd3..04c5a61c8d 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core/tests.rs @@ -33,8 +33,8 @@ fn enter_isolated_test_process() -> bool { /// A service built from explicit options. `Default` options read the node /// identity from a fixed host file and probe runtimes under the fixed run /// directory; here the node identity is fixed, runtime liveness reads -/// `scratch`, and the host-control and credential-manager sockets are absent -/// paths under `scratch`. +/// `scratch`, the host-control, credential-manager and node-manager sockets +/// are absent paths under `scratch`, and every template lookup is refused. fn isolated_service( conn: rusqlite::Connection, scratch: &std::path::Path, @@ -57,6 +57,14 @@ fn isolated_service( node_role, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: crate::node_manager_transport::fake::unbound( + &scratch.join("absent-node-manager.sock"), + ), + template_resolver: Arc::new( + crate::template_resolver::fake::FakeTemplateResolver::refusing( + Status::unavailable("template-manager is outside this test"), + ), + ), }, ) } diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/artifact_backup.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/artifact_backup.rs index c255036780..b2c59a5ee1 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/artifact_backup.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/artifact_backup.rs @@ -738,9 +738,10 @@ impl SandboxManagerService { (SandboxState::Stopped, None) => { let id = LocalArtifactId::new(); let snapshot_name = format!("d-stop-{id}"); - let mut node_client = connect_node().await.map_err(|error| { - Status::internal(format!("connect to node-manager: {error}")) - })?; + let mut node_client = + connect_node(&self.node_manager).await.map_err(|error| { + Status::internal(format!("connect to node-manager: {error}")) + })?; node_client .create_btrfs_snapshot( pplx_space_node_mgr_client::proto::CreateBtrfsSnapshotRequest { diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/backup_submission.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/backup_submission.rs index cc1103d11b..0053efe075 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/backup_submission.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/backup_submission.rs @@ -410,7 +410,9 @@ impl SandboxManagerService { .map_err(ArtifactBackupSubmissionError::Begin)?; } if payload.job_ids.is_empty() { - match crate::snapshot_upload::submit_prepared_snapshot(&payload).await { + match crate::snapshot_upload::submit_prepared_snapshot(&payload, &self.node_manager) + .await + { Ok(submitted) => { match crate::snapshot_upload_jobs::persist_submitted_transfer_ids( &self.db, job_id, &submitted, diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/kill.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/kill.rs index 39c7e275b5..51c2662c2e 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/kill.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/kill.rs @@ -36,7 +36,7 @@ impl SandboxManagerService { // Non-fatal signals (e.g. SIGHUP, SIGUSR1) are just forwarded to the process. let is_fatal = signal == 9 || signal == 15; if !is_fatal { - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; node_client @@ -73,7 +73,7 @@ impl SandboxManagerService { let execution_fence = sandbox.node_execution_fence(); let dispatch_error = match self .dispatch_node_action_using_admission(sandbox_id, &admission, || async move { - let mut node_client = connect_node().await.map_err(|error| { + let mut node_client = connect_node(&self.node_manager).await.map_err(|error| { Status::internal(format!("cannot connect to node-manager: {error}")) })?; node_client diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/pause_resume.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/pause_resume.rs index 4758c8223b..1ab59abdda 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/pause_resume.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/pause_resume.rs @@ -35,9 +35,11 @@ pub(super) fn pause_reconciliation_outcome( async fn reconcile_state_after_unconfirmed_pause( sandbox_id: &SandboxId, reason: &str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> PauseReconciliationOutcome { let observation = - observe_firecracker_vm_state(sandbox_id, "pause response reconciliation").await; + observe_firecracker_vm_state(sandbox_id, "pause response reconciliation", node_manager) + .await; let outcome = pause_reconciliation_outcome(&observation); tracing::warn!( sandbox_id = %sandbox_id, @@ -58,6 +60,7 @@ pub(super) async fn pause_with_reconciliation( sandbox_id: &SandboxId, reason_context: &str, guest_preparation_budget: GuestSnapshotPreparationBudget, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> SandboxState { let sandbox_id_string = sandbox_id.to_string(); let mut last_reason = None; @@ -73,7 +76,9 @@ pub(super) async fn pause_with_reconciliation( Err(error) => { let reason = format!("{reason_context}: {error}"); last_reason = Some(reason.clone()); - let outcome = reconcile_state_after_unconfirmed_pause(sandbox_id, &reason).await; + let outcome = + reconcile_state_after_unconfirmed_pause(sandbox_id, &reason, node_manager) + .await; let retry_requires_prepare = matches!(&outcome, PauseReconciliationOutcome::Running); match outcome { @@ -181,7 +186,7 @@ impl SandboxManagerService { ))); } - let mut node_client = match connect_node().await { + let mut node_client = match connect_node(&self.node_manager).await { Ok(node_client) => node_client, Err(e) => { update_state_best_effort(&self.db, &sandbox.id, SandboxState::Started); @@ -202,6 +207,7 @@ impl SandboxManagerService { &sandbox_id, "pause failed", GuestSnapshotPreparationBudget::Standard, + &self.node_manager, ) .await; if pause_state != SandboxState::Paused { diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend.rs index 8a4d3ee80e..c1616cda41 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend.rs @@ -331,7 +331,7 @@ impl SandboxManagerService { // Auto-pause if started if sandbox.state == SandboxState::Started { let pause_started_at = std::time::Instant::now(); - let mut node_client = connect_node().await.map_err(|e| { + let mut node_client = connect_node(&self.node_manager).await.map_err(|e| { pplx_space_metrics::snapshot::create_failed( create_kind, context.create_codec, @@ -401,6 +401,7 @@ impl SandboxManagerService { &sandbox_id, "auto-pause failed", guest_preparation_budget, + &self.node_manager, ) .await; if pause_state != SandboxState::Paused { diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend_snapshot.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend_snapshot.rs index 04d5cff0c4..1460d63154 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend_snapshot.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/suspend_snapshot.rs @@ -407,7 +407,7 @@ impl SandboxManagerService { // the success path. Harmless no-op on paths that never unsubscribed. let mut monitoring_guard = SuspendMonitoringGuard::new(self, &sandbox); let create_kind = "suspend"; - let mut node_client = connect_node().await.map_err(|e| { + let mut node_client = connect_node(&self.node_manager).await.map_err(|e| { pplx_space_metrics::snapshot::create_failed( create_kind, create_codec, @@ -533,6 +533,7 @@ impl SandboxManagerService { sandbox.node_execution_fence(), fresh, None, + &self.node_manager, ) .await { @@ -595,6 +596,7 @@ impl SandboxManagerService { &sandbox.id, sandbox.node_execution_fence(), "", + &self.node_manager, ) .await { @@ -649,6 +651,7 @@ impl SandboxManagerService { &sandbox.id, sandbox.node_execution_fence(), &staging_dir_string, + &self.node_manager, ) .await { @@ -857,6 +860,7 @@ impl SandboxManagerService { drop(crate::service::recovery::staging_cleanup::reclaim_garbage_now( sandbox.id.clone(), path, + self.node_manager.clone(), )); } None @@ -1167,6 +1171,7 @@ impl SandboxManagerService { .as_ref() .ok_or_else(|| Status::failed_precondition("execution identity is missing"))? .try_into()?, + &self.node_manager, ) .await .map_err(|e| e.into_status())?; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/tests.rs b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/tests.rs index cba1ee61f9..82ba947137 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_lifecycle/tests.rs @@ -1908,6 +1908,7 @@ async fn confirmed_node_pause_reports_paused_and_carries_the_execution_fence() { previous_id: None, }; + let connector = crate::node_manager_transport::fake::fixed_client(client.clone()); let state = super::pause_with_reconciliation( &mut client, Some(fence.clone()), @@ -1915,6 +1916,7 @@ async fn confirmed_node_pause_reports_paused_and_carries_the_execution_fence() { &sandbox_id, "pause failed", GuestSnapshotPreparationBudget::Standard, + &connector, ) .await; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_net_gateway.rs b/crates/libs/sbx/mgr/server/src/service/grpc_net_gateway.rs index e61a9894a3..e8cd6d7c35 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_net_gateway.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_net_gateway.rs @@ -129,7 +129,7 @@ impl SandboxManagerService { // check and the signal. That is not a failure: skip further signaling // and fall through to the relaunch, which either forks a fresh // gateway or recovers the watchdog's already-live replacement. - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|error| fail(format!("cannot connect to node-manager: {error}")))?; let term_signalled = node_client diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs b/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs index f3b72c4431..7f10447b82 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs @@ -258,6 +258,7 @@ impl SandboxManagerService { let resolved_template = if req.auto_lifecycle.is_none() { Some( resolve_template_paths( + self.template_resolver.as_ref(), pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), @@ -640,7 +641,7 @@ impl SandboxManagerService { &self, request: Request, ) -> Result, Status> { - self.clone_from_donor_with_node(request, connect_clone_node()) + self.clone_from_donor_with_node(request, connect_clone_node(&self.node_manager)) .await } @@ -693,7 +694,8 @@ impl SandboxManagerService { let uuid8 = SandboxRuntimeLabel::from_sandbox_id(new_sandbox_id); // Resolve template and compute hash - let resolved_template = resolve_template_paths(template_ref).await?; + let resolved_template = + resolve_template_paths(self.template_resolver.as_ref(), template_ref).await?; reject_incompatible_template_arch( template, resolved_template @@ -872,18 +874,18 @@ impl SandboxManagerService { } } -async fn connect_clone_node() -> Result { - pplx_space_node_mgr_client::connect() - .await - .map_err(|error| match error { - pplx_space_node_mgr_client::ConnectError::InvalidUri(_) => { - Status::internal("node-manager endpoint is invalid") - } - pplx_space_node_mgr_client::ConnectError::Transport(_) - | pplx_space_node_mgr_client::ConnectError::Connect(_) => { - Status::unavailable("cannot connect to node-manager") - } - }) +async fn connect_clone_node( + node_manager: &crate::node_manager_transport::NodeManagerConnector, +) -> Result { + node_manager.connect().await.map_err(|error| match error { + pplx_space_node_mgr_client::ConnectError::InvalidUri(_) => { + Status::internal("node-manager endpoint is invalid") + } + pplx_space_node_mgr_client::ConnectError::Transport(_) + | pplx_space_node_mgr_client::ConnectError::Connect(_) => { + Status::unavailable("cannot connect to node-manager") + } + }) } pub(super) fn pool_adoption_rollback_failure(original: Status, rollback: Status) -> Status { diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_spaced_files.rs b/crates/libs/sbx/mgr/server/src/service/grpc_spaced_files.rs index ded8eb5f90..c8568cf1b8 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_spaced_files.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_spaced_files.rs @@ -40,7 +40,7 @@ impl SandboxManagerService { .slot .ok_or_else(|| Status::failed_precondition("sandbox has no allocated network slot"))?; - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; @@ -204,7 +204,7 @@ impl SandboxManagerService { } if updates_resources { - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::unavailable(format!("cannot connect to node-manager: {e}")))?; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs b/crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs index 46cd332240..2aa3731989 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_spaced_process_snapshots.rs @@ -851,7 +851,7 @@ impl SandboxManagerService { let mut snapshots = Vec::new(); let mut response_budget = InventoryBudget::new(MAX_ITEMS, MAX_ENCODED_BYTES); - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; @@ -1240,7 +1240,7 @@ impl SandboxManagerService { .map(|artifact| artifact.local_path) }); - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; @@ -1486,7 +1486,7 @@ impl SandboxManagerService { }) .await?; - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; diff --git a/crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs b/crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs index 76ad958720..aaf8a034ef 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_state_snapshot.rs @@ -1174,6 +1174,7 @@ impl SandboxManagerService { db: self.db.clone(), auto_lifecycle: self.auto_lifecycle.clone(), health_lifecycle: self.watchdog.lifecycle_epochs(), + node_manager: self.node_manager.clone(), }; let completion = crate::snapshot_scheduler::do_checkpoint(&ctx, id, &snapshot_name) @@ -1189,7 +1190,9 @@ impl SandboxManagerService { Some(completion) } else { // Disk snapshot: btrfs subvolume snapshot via node-manager - let mut node_client = pplx_space_node_mgr_client::connect() + let mut node_client = self + .node_manager + .connect() .await .map_err(|e| Status::internal(format!("connect to node-manager: {e}")))?; node_client @@ -1249,7 +1252,7 @@ impl SandboxManagerService { ) }; if let Err(error) = publish_result { - match pplx_space_node_mgr_client::connect().await { + match self.node_manager.connect().await { Ok(mut node_client) => { if let Err(cleanup_error) = node_client .delete_btrfs_snapshot( @@ -1457,9 +1460,10 @@ impl SandboxManagerService { })?; } - let mut node_client = pplx_space_node_mgr_client::connect() - .await - .map_err(|error| Status::unavailable(format!("connect to node-manager: {error}")))?; + let mut node_client = + self.node_manager.connect().await.map_err(|error| { + Status::unavailable(format!("connect to node-manager: {error}")) + })?; node_client .sync_guest(pplx_space_node_mgr_client::proto::SyncGuestRequest { sandbox_id: sandbox_id.clone(), @@ -1567,6 +1571,7 @@ impl SandboxManagerService { content_hash, snapshot_name.as_str(), &paths.artifact, + &self.node_manager, ) .await .map_err(crate::btrfs_via_nodemgr::BtrfsViaNodemgrError::into_status), @@ -1809,7 +1814,7 @@ impl SandboxManagerService { Status::internal(format!("create template snapshot staging dir: {e}")) })?; - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; let phase = Instant::now(); diff --git a/crates/libs/sbx/mgr/server/src/service/mod.rs b/crates/libs/sbx/mgr/server/src/service/mod.rs index 231fb56850..f42685dc99 100644 --- a/crates/libs/sbx/mgr/server/src/service/mod.rs +++ b/crates/libs/sbx/mgr/server/src/service/mod.rs @@ -96,6 +96,7 @@ pub use self::restart_admission::{ SandboxManagerRestartAdmissionLayer, handoff_cancellable_response_stream_name, }; use self::start_support::ActivationAdmissionLimiter; +use crate::node_manager_transport::NodeManagerConnector; use crate::node_state_sync::{mark_sandbox_dirty, mark_sandboxes_dirty, should_skip_node_publish}; use crate::proto::sandbox_manager_server::SandboxManager; use crate::proto::{ @@ -401,6 +402,8 @@ pub struct SandboxManagerService { credential_manager_client: Arc>, network_credentials_enabled: bool, + node_manager: NodeManagerConnector, + template_resolver: Arc, gateway_caller_client: pplx_space_node_mgr_client::NodeManagerGrpcClient, gateway_authority: gateway_authority::GatewayAuthority, /// S3 upload config for suspend snapshots. `None` means uploads are disabled. @@ -970,6 +973,10 @@ pub struct SandboxManagerServiceOptions { pub route_ed25519_verification_bundle: Option< Arc, >, + /// Every node-manager client the service and its subsystems use. + pub node_manager: NodeManagerConnector, + /// Every template-manager `GetTemplate` the service and its subsystems ask. + pub template_resolver: Arc, } impl Default for SandboxManagerServiceOptions { @@ -988,6 +995,8 @@ impl Default for SandboxManagerServiceOptions { node_role: None, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), } } } @@ -1621,6 +1630,7 @@ impl SandboxManagerService { db: db.clone(), auto_lifecycle: auto_lifecycle.clone(), health_lifecycle: health_lifecycle.clone(), + node_manager: options.node_manager.clone(), }); let snapshot_scheduler = Arc::new( crate::snapshot_scheduler::SnapshotScheduler::new_with_ownership( @@ -1640,12 +1650,16 @@ impl SandboxManagerService { options.shutdown_token.clone(), Arc::clone(&drain_state), &options.task_tracker, + options.node_manager.clone(), )); if let Some(generation) = expected_route_verification_key_generation { watchdog.set_expected_route_verification_key_generation(generation); } - let snapshot_upload_worker = - crate::snapshot_upload_jobs::build_worker(Arc::clone(&db), Arc::clone(&drain_state)); + let snapshot_upload_worker = crate::snapshot_upload_jobs::build_worker( + Arc::clone(&db), + Arc::clone(&drain_state), + options.node_manager.clone(), + ); let local_artifacts = crate::local_artifacts::LocalArtifacts::new(Arc::clone(&db)); let teardown_suspend_driver_concurrency = TEARDOWN_SUSPEND_DRIVER_CONCURRENCY; pplx_space_metrics::sandbox::worker::teardown_suspend_driver_capacity( @@ -1668,6 +1682,7 @@ impl SandboxManagerService { template_network_policies: TemplateNetworkPolicyCache::new( options.shutdown_token.clone(), options.task_tracker.clone(), + Arc::clone(&options.template_resolver), ), network_policy_reconcile_queue, network_policy_reconciler_started: Arc::new(AtomicBool::new(false)), @@ -1715,7 +1730,7 @@ impl SandboxManagerService { .as_deref(), false, ), - gateway_caller_client: pplx_space_node_mgr_client::lazy(), + gateway_caller_client: options.node_manager.lazy(), gateway_authority, snapshot_upload_config: { let cfg = crate::snapshot_upload::SnapshotUploadConfig::from_env(); @@ -1729,7 +1744,11 @@ impl SandboxManagerService { }, snapshot_upload_worker, sandbox_runtime_live: options.sandbox_runtime_live, - retained_source_inspector: Arc::new(crate::retained_source::NodeManagerSourceInspector), + retained_source_inspector: Arc::new( + crate::retained_source::NodeManagerSourceInspector(options.node_manager.clone()), + ), + node_manager: options.node_manager, + template_resolver: options.template_resolver, drain_state, feature_flags, route_ed25519_verification_bundle: Arc::new(ArcSwapOption::from( @@ -1809,6 +1828,7 @@ impl SandboxManagerService { coordinator: self.auto_lifecycle.clone(), drain: Arc::clone(&self.drain_state), shutdown: shutdown_token, + node_manager: self.node_manager.clone(), }; task_tracker.spawn(async move { crate::snapshot_gc::run_loop( diff --git a/crates/libs/sbx/mgr/server/src/service/network_policy.rs b/crates/libs/sbx/mgr/server/src/service/network_policy.rs index 1c7b614cee..9111e83f4c 100644 --- a/crates/libs/sbx/mgr/server/src/service/network_policy.rs +++ b/crates/libs/sbx/mgr/server/src/service/network_policy.rs @@ -130,20 +130,22 @@ pub(crate) fn network_policy_override_from_set_request( } pub(crate) async fn fetch_template_default_network_policy( + templates: &dyn crate::template_resolver::TemplateResolver, template: &str, ) -> Result { - let mut client = pplx_space_template_mgr_client::lazy(); - let response = match client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(pplx_space_template_mgr_client::proto::TemplateRef { - name: template.to_string(), - content_hash: None, - tag: None, - }), - vcpu: None, - memory_mib: None, - }) + let response = match templates + .get_template_on_shared_channel(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(pplx_space_template_mgr_client::proto::TemplateRef { + name: template.to_string(), + content_hash: None, + tag: None, + }), + vcpu: None, + memory_mib: None, + }, + )) .await { Ok(response) => response.into_inner(), @@ -175,15 +177,7 @@ pub(crate) struct TemplateNetworkPolicyCache { fetch_slots: Arc, shutdown_token: tokio_util::sync::CancellationToken, task_tracker: tokio_util::task::TaskTracker, -} - -impl Default for TemplateNetworkPolicyCache { - fn default() -> Self { - Self::new( - tokio_util::sync::CancellationToken::new(), - tokio_util::task::TaskTracker::new(), - ) - } + templates: Arc, } #[derive(Default)] @@ -216,6 +210,7 @@ impl TemplateNetworkPolicyCache { pub(crate) fn new( shutdown_token: tokio_util::sync::CancellationToken, task_tracker: tokio_util::task::TaskTracker, + templates: Arc, ) -> Self { Self { inner: Arc::default(), @@ -224,6 +219,7 @@ impl TemplateNetworkPolicyCache { )), shutdown_token, task_tracker, + templates, } } @@ -233,8 +229,9 @@ impl TemplateNetworkPolicyCache { &self, template: &str, ) -> Result { + let templates = Arc::clone(&self.templates); self.get_or_fetch_with(template, |template| async move { - fetch_template_default_network_policy(&template).await + fetch_template_default_network_policy(templates.as_ref(), &template).await }) .await } diff --git a/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs b/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs index 9063e9f87f..4c10af9fd4 100644 --- a/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/network_policy/tests.rs @@ -1,4 +1,19 @@ use super::*; +use crate::template_resolver::fake::FakeTemplateResolver; + +/// A cache whose owner lookups are refused, with the resolver that refuses +/// them. Tests that supply their own fetch never reach it. +fn refusing_cache() -> (TemplateNetworkPolicyCache, Arc) { + let templates = Arc::new(FakeTemplateResolver::refusing(Status::unavailable( + "template-manager is outside this test", + ))); + let cache = TemplateNetworkPolicyCache::new( + tokio_util::sync::CancellationToken::new(), + tokio_util::task::TaskTracker::new(), + Arc::clone(&templates) as Arc, + ); + (cache, templates) +} #[test] fn legacy_policy_hash_is_stable() { @@ -27,7 +42,7 @@ async fn template_policy_refresh_uses_current_owner_across_local_changes_and_rec (Local, 0, Local, 0), // Old template-manager local updates reuse zero. (Local, 9, Cluster, 1), // Delete/recreate legitimately changes authority. ] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; cache .get_or_fetch_with("seed", move |_| async move { Ok(EffectiveNetworkPolicy::new( @@ -111,7 +126,7 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() Some(tonic::Code::FailedPrecondition), ), ] { - let cache = TemplateNetworkPolicyCache::default(); + let (cache, templates) = refusing_cache(); cache .get_or_fetch_with("seed", move |_| async move { Ok(EffectiveNetworkPolicy::new( @@ -134,6 +149,8 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() true, ); let admission_error = cache.require_hint("seed", &hint, true).await.err(); + // The hint is checked against the cached owner answer, not a new lookup. + assert!(templates.lookups().is_empty()); assert_eq!(admission_error.as_ref().map(|error| error.code()), expected); if let Some(error) = &admission_error { assert!(pplx_space_sbx_mgr_client::is_template_policy_dependency( @@ -173,7 +190,7 @@ async fn template_policy_hints_are_admission_requirements_not_cache_authority() #[tokio::test(start_paused = true)] async fn custom_policy_override_does_not_require_template_owner_progress() { - let cache = TemplateNetworkPolicyCache::default(); + let (cache, templates) = refusing_cache(); let hint = EffectiveNetworkPolicy::new( NetworkPolicyV1::deny_all(), 50, @@ -190,11 +207,12 @@ async fn custom_policy_override_does_not_require_template_owner_progress() { inherits_template.then_some(tonic::Code::DeadlineExceeded) ); } + assert!(templates.lookups().is_empty()); } #[tokio::test(start_paused = true)] async fn template_policy_invalidation_fences_fetch_callers_and_successor_claims() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let mut old = Box::pin(cache.get_or_fetch_with("seed", |_| std::future::pending())); assert!(futures_util::poll!(&mut old).is_pending()); let old_claim = Arc::clone(&cache.inner.lock().await.pending_fetches["seed"].claim); @@ -258,7 +276,7 @@ async fn template_policy_invalidation_fences_fetch_callers_and_successor_claims( #[tokio::test(start_paused = true)] async fn template_policy_wait_deadline_includes_cache_admission() { for release_midway in [false, true] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let guard = cache.inner.lock().await; let mut waiting = Box::pin(cache.get_or_fetch_with("seed", |_| std::future::pending())); assert!(futures_util::poll!(&mut waiting).is_pending()); @@ -280,7 +298,7 @@ async fn template_policy_wait_deadline_includes_cache_admission() { #[tokio::test(start_paused = true)] async fn template_policy_fetch_progress_survives_foreground_timeouts_and_cancellation() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let (alive, cancelled) = tokio::sync::oneshot::channel::<()>(); let (release, released) = tokio::sync::oneshot::channel::<()>(); let mut first = Box::pin(cache.get_or_fetch_with("seed", |_| async move { @@ -324,7 +342,7 @@ async fn template_policy_fetch_progress_survives_foreground_timeouts_and_cancell #[tokio::test(start_paused = true)] async fn template_policy_fetch_capacity_and_deadline_release_claims() { for shutdown in [false, true] { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let names: Vec<_> = (0..TEMPLATE_NETWORK_POLICY_FETCH_CAPACITY) .map(|i| format!("template-{i}")) .collect(); @@ -391,7 +409,7 @@ async fn template_policy_fetch_capacity_and_deadline_release_claims() { // A ready Shared can be completed by its caller before its tracked owner // is scheduled. Those queued owner tasks still consume the same hard cap. - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; for i in 0..TEMPLATE_NETWORK_POLICY_FETCH_CAPACITY { let name = format!("ready-{i}"); let mut caller = Box::pin(tokio::task::unconstrained(cache.get_or_fetch_with( @@ -426,7 +444,7 @@ async fn template_policy_fetch_capacity_and_deadline_release_claims() { #[tokio::test(start_paused = true)] async fn template_policy_cache_bounds_churn_without_retiring_pending_claims() { - let cache = TemplateNetworkPolicyCache::default(); + let cache = refusing_cache().0; let mut waiting = Box::pin(cache.get_or_fetch_with("pending", |_| std::future::pending())); assert!(futures_util::poll!(&mut waiting).is_pending()); let claim = Arc::clone(&cache.inner.lock().await.pending_fetches["pending"].claim); diff --git a/crates/libs/sbx/mgr/server/src/service/recovery.rs b/crates/libs/sbx/mgr/server/src/service/recovery.rs index 870e6b2335..9c41a274c9 100644 --- a/crates/libs/sbx/mgr/server/src/service/recovery.rs +++ b/crates/libs/sbx/mgr/server/src/service/recovery.rs @@ -411,6 +411,7 @@ impl SandboxManagerService { crate::runtime_reconciliation::observe_firecracker_vm_state( &sandbox_id, "startup pause recovery", + &self.node_manager, ) .await }) @@ -883,7 +884,11 @@ impl SandboxManagerService { let error = match published { Ok(garbage) => { for path in garbage { - drop(staging_cleanup::reclaim_garbage_now(row.id.clone(), path)); + drop(staging_cleanup::reclaim_garbage_now( + row.id.clone(), + path, + self.node_manager.clone(), + )); } None } @@ -2233,7 +2238,7 @@ impl SandboxManagerService { let _guard = guard; let mut recovery_cursor = crate::db::PeriodicLifecycleRecoveryScan::default(); let mut staging_scan = staging_cleanup::StagingCleanupScan::default(); - let mut node_client = pplx_space_node_mgr_client::lazy(); + let mut node_client = service.node_manager.lazy(); let mut inventory_scan = super::cache_cleanup::PeriodicInventoryScan::default(); let mut quarantine_scan = crate::db::PeriodicQuarantineRecoveryScan::default(); let mut transfer_scan = crate::db::inventory::RestoreTransferScopeScan::default(); diff --git a/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup.rs b/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup.rs index 3e240241b1..a5baf20f6d 100644 --- a/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup.rs +++ b/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup.rs @@ -314,6 +314,7 @@ impl SandboxManagerService { pub(in crate::service) fn reclaim_garbage_now( sandbox_id: String, garbage: PathBuf, + node_manager: crate::node_manager_transport::NodeManagerConnector, ) -> Option> { // Admission and enqueue happen under the backlog lock, and a worker gives // its slot back under the same lock only after seeing the backlog empty, @@ -326,7 +327,7 @@ pub(in crate::service) fn reclaim_garbage_now( Ok(permit) => permit, Err(_) => { if backlog.len() < IMMEDIATE_RECLAIM_BACKLOG { - backlog.push_back((sandbox_id, garbage)); + backlog.push_back((sandbox_id, garbage, node_manager)); } else { tracing::debug!(sandbox_id, path = %garbage.display(), "immediate reclaim backlog full; the staging sweep collects the entry"); } @@ -336,9 +337,10 @@ pub(in crate::service) fn reclaim_garbage_now( }; Some(tokio::spawn(async move { let mut permit = Some(permit); - let mut next = Some((sandbox_id, garbage)); - while let Some((sandbox_id, garbage)) = next { - reclaim_one(&sandbox_id, &garbage).await; + let mut next = Some((sandbox_id, garbage, node_manager)); + // Each entry is deleted through the connector that submitted it. + while let Some((sandbox_id, garbage, node_manager)) = next { + reclaim_one(&sandbox_id, &garbage, &node_manager).await; let mut backlog = immediate_reclaim_backlog(); next = backlog.pop_front(); if next.is_none() { @@ -349,15 +351,25 @@ pub(in crate::service) fn reclaim_garbage_now( })) } -fn immediate_reclaim_backlog() -> std::sync::MutexGuard<'static, VecDeque<(String, PathBuf)>> { - static BACKLOG: std::sync::Mutex> = +type ReclaimEntry = ( + String, + PathBuf, + crate::node_manager_transport::NodeManagerConnector, +); + +fn immediate_reclaim_backlog() -> std::sync::MutexGuard<'static, VecDeque> { + static BACKLOG: std::sync::Mutex> = std::sync::Mutex::new(VecDeque::new()); BACKLOG .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) } -async fn reclaim_one(sandbox_id: &str, garbage: &Path) { +async fn reclaim_one( + sandbox_id: &str, + garbage: &Path, + node_manager: &crate::node_manager_transport::NodeManagerConnector, +) { let local = { let garbage = garbage.to_path_buf(); tokio::task::spawn_blocking(move || std::fs::remove_dir_all(&garbage)).await @@ -374,7 +386,7 @@ async fn reclaim_one(sandbox_id: &str, garbage: &Path) { return; } }; - let mut client = pplx_space_node_mgr_client::lazy(); + let mut client = node_manager.lazy(); let request = pplx_space_node_mgr_client::proto::DeleteCacheArtifactRequest { local_path: garbage.to_string_lossy().into_owned(), }; diff --git a/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup/tests.rs b/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup/tests.rs index f3a3da7f47..b164bcfd72 100644 --- a/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/recovery/staging_cleanup/tests.rs @@ -563,6 +563,8 @@ async fn immediate_reclaim_removes_ordinary_entries_and_leaves_the_rest_in_the_l // Read-only directory: unlinking its entries fails for a normal user, as // a retained subvolume's do. std::fs::set_permissions(&bundle, std::fs::Permissions::from_mode(0o555)).unwrap(); + // Node-manager refuses the delete the stubborn entry falls back to. + let node_manager = crate::node_manager_transport::fake::FakeNodeManager::default().serve(); // At capacity nothing is spawned; the entries join the backlog. The // slots here are held by the test rather than by workers, so no worker @@ -581,19 +583,21 @@ async fn immediate_reclaim_removes_ordinary_entries_and_leaves_the_rest_in_the_l .map(|_| IMMEDIATE_RECLAIM_SLOTS.try_acquire().unwrap()) .collect(); for dir in &backlog { - assert!(reclaim_garbage_now("sandbox".into(), dir.clone()).is_none()); + assert!( + reclaim_garbage_now("sandbox".into(), dir.clone(), node_manager.connector()).is_none() + ); assert!(dir.exists()); } drop(held); - reclaim_garbage_now("sandbox".into(), ordinary.clone()) + reclaim_garbage_now("sandbox".into(), ordinary.clone(), node_manager.connector()) .unwrap() .await .unwrap(); for dir in &backlog { assert!(!dir.exists(), "backlog drained once a slot was free"); } - reclaim_garbage_now("sandbox".into(), stubborn.clone()) + reclaim_garbage_now("sandbox".into(), stubborn.clone(), node_manager.connector()) .unwrap() .await .unwrap(); @@ -612,3 +616,69 @@ async fn immediate_reclaim_removes_ordinary_entries_and_leaves_the_rest_in_the_l } assert!(remaining.len() <= 1); } + +/// A backlogged entry falls back to node-manager through the connector that +/// submitted it, not through the connector of whichever worker drains it. +#[tokio::test] +async fn immediate_reclaim_backlog_uses_each_entrys_own_node_manager() { + use pplx_space_node_mgr_client::proto::{ + DeleteCacheArtifactRequest, DeleteCacheArtifactResponse, + }; + use std::sync::{Arc, Mutex}; + + fn recording_node_manager( + deletes: &Arc>>, + ) -> crate::node_manager_transport::fake::ServedNodeManager { + let deletes = Arc::clone(deletes); + crate::node_manager_transport::fake::FakeNodeManager::default() + .unary( + "DeleteCacheArtifact", + move |request: DeleteCacheArtifactRequest| { + deletes.lock().unwrap().push(request.local_path); + Ok(DeleteCacheArtifactResponse::default()) + }, + ) + .serve() + } + + let cache = tempfile::tempdir().unwrap(); + let lane = suspend_snapshot_publication::suspend_snapshot_staging_parent_dir( + &cache.path().join("sandbox"), + ); + std::fs::create_dir_all(&lane).unwrap(); + // A regular file is not a directory, so the local removal fails even for + // root and every entry falls back to node-manager. + let entry = |name: &str| { + let path = lane.join(format!("{RETIRED_PREFIX}{name}")); + std::fs::write(&path, b"retired").unwrap(); + path + }; + let (first, second, drainer) = (entry("first"), entry("second"), entry("drainer")); + let deletes = [(); 3].map(|()| Arc::new(Mutex::new(Vec::new()))); + let node_managers = deletes.each_ref().map(recording_node_manager); + + let held: Vec<_> = (0..IMMEDIATE_RECLAIM_CONCURRENCY) + .map(|_| IMMEDIATE_RECLAIM_SLOTS.try_acquire().unwrap()) + .collect(); + for (path, node_manager) in [(&first, &node_managers[0]), (&second, &node_managers[1])] { + assert!( + reclaim_garbage_now("sandbox".into(), path.clone(), node_manager.connector()).is_none(), + "both slots are held, so the entry is backlogged" + ); + } + drop(held); + reclaim_garbage_now( + "sandbox".into(), + drainer.clone(), + node_managers[2].connector(), + ) + .unwrap() + .await + .unwrap(); + + let deleted = deletes.map(|deletes| deletes.lock().unwrap().clone()); + assert_eq!( + deleted, + [first, second, drainer].map(|path| vec![path.to_string_lossy().into_owned()]) + ); +} diff --git a/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs b/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs index f97af25616..371aa3173b 100644 --- a/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs +++ b/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs @@ -1303,13 +1303,11 @@ impl SandboxManagerService { // This is intentionally whole-sandbox disposal, not partial // restore cleanup. Keep teardown and recursive removal with // the node-manager's lifecycle owner through completion. - let mut node = pplx_space_node_mgr_client::connect().await.map_err( - |error| { - RestoreError::Other(format!( - "connect node-manager for orphan cleanup {canonical_id}: {error}" - )) - }, - )?; + let mut node = self.node_manager.connect().await.map_err(|error| { + RestoreError::Other(format!( + "connect node-manager for orphan cleanup {canonical_id}: {error}" + )) + })?; super::cleanup_support::cleanup_sandbox_via_node( &self.db, &mut node, @@ -1767,9 +1765,14 @@ impl SandboxManagerService { .ok_or_else(|| { RestoreError::Other("restore execution identity is missing".into()) })?; - let phase2_result = restore_pull_rootfs(&staged, execution_fence) - .instrument(phase2_span) - .await; + let phase2_result = restore_pull_rootfs( + &staged, + execution_fence, + &self.node_manager, + self.template_resolver.as_ref(), + ) + .instrument(phase2_span) + .await; if let Err(ref e) = phase2_result { emit_sandbox_event( &canonical_id, diff --git a/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs b/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs index f10242546a..f9c9078a4f 100644 --- a/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/route_identity_tests.rs @@ -85,6 +85,8 @@ impl RouteTestService { node_role: None, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), }, ); Self { diff --git a/crates/libs/sbx/mgr/server/src/service/spaced_bridge.rs b/crates/libs/sbx/mgr/server/src/service/spaced_bridge.rs index 1662b6d344..48496a6ac2 100644 --- a/crates/libs/sbx/mgr/server/src/service/spaced_bridge.rs +++ b/crates/libs/sbx/mgr/server/src/service/spaced_bridge.rs @@ -77,7 +77,11 @@ pub(super) async fn launch_spaced_host_via_node( } } -pub(super) fn cleanup_restored_snapshot_dir_async(sandbox_id: &str, snapshot_dir: &str) { +pub(super) fn cleanup_restored_snapshot_dir_async( + sandbox_id: &str, + snapshot_dir: &str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, +) { // The consumed snapshot moves into the staging sweep's lane first, so a // deferred capture's read-only subvolumes, which this service cannot // delete, have an owner even if nothing below runs; then it is reclaimed @@ -91,6 +95,7 @@ pub(super) fn cleanup_restored_snapshot_dir_async(sandbox_id: &str, snapshot_dir drop(super::recovery::staging_cleanup::reclaim_garbage_now( sandbox_id.to_string(), garbage, + node_manager.clone(), )); } Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} @@ -137,7 +142,7 @@ impl SandboxManagerService { } let runtime = spaced_endpoint_runtime_identity(sandbox)?; - let mut node_client = pplx_space_node_mgr_client::lazy(); + let mut node_client = self.node_manager.lazy(); let proof = match sandbox.spaced_endpoint_proof.clone() { Some(proof) => proof, None => { diff --git a/crates/libs/sbx/mgr/server/src/service/start.rs b/crates/libs/sbx/mgr/server/src/service/start.rs index 9ff9a723e7..8dc53f2d9b 100644 --- a/crates/libs/sbx/mgr/server/src/service/start.rs +++ b/crates/libs/sbx/mgr/server/src/service/start.rs @@ -513,6 +513,8 @@ async fn prepare_suspend_restore_dir( execution_fence, storage_uri, transfer_scope_id, + &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -521,6 +523,8 @@ async fn prepare_suspend_restore_dir( &sandbox_id, execution_fence, transfer_scope_id, + &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -750,7 +754,7 @@ impl SandboxManagerService { &self, requested_memory_mib: u32, ) -> Result { - let mut node_client = connect_node() + let mut node_client = connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}")))?; let node_memory_policy = Self::node_sandbox_memory_policy(&mut node_client).await?; @@ -1361,6 +1365,7 @@ impl SandboxManagerService { if !vmstate_path.exists() { let resolved_template = resolve_template_paths_for_start( + self.template_resolver.as_ref(), pplx_space_template_mgr_client::proto::TemplateRef { name: sandbox.template.clone(), content_hash: sandbox.template_hash.clone(), @@ -1643,7 +1648,7 @@ impl SandboxManagerService { "fresh.template_prepare_source", START_TEMPLATE_RESOLUTION_TIMEOUT, &sandbox.id, - template_source_for_sandbox_prepare(&sandbox), + template_source_for_sandbox_prepare(self.template_resolver.as_ref(), &sandbox), ) .await?; let selection_mode = if prepare_source.existing_artifact { @@ -1659,7 +1664,11 @@ impl SandboxManagerService { START_TEMPLATE_RESOLUTION_TIMEOUT, &sandbox.id, TemplateRuntimeDefaults::default(), - template_runtime_defaults_for_sandbox_start(&sandbox, prepare_template_dir.as_deref()), + template_runtime_defaults_for_sandbox_start( + self.template_resolver.as_ref(), + &sandbox, + prepare_template_dir.as_deref(), + ), ) .await; @@ -1780,7 +1789,7 @@ impl SandboxManagerService { START_NODE_METADATA_TIMEOUT, id, async { - connect_node() + connect_node(&self.node_manager) .await .map_err(|e| Status::internal(format!("cannot connect to node-manager: {e}"))) }, @@ -2119,12 +2128,17 @@ impl SandboxManagerService { // one node-local round-trip for the warm-pool composite *shape* below // — a device-layout choice, not an encryption decision. let rootfs_block_device_path = if sandbox.volume_encryption { - match crate::rootfs_encryption::prepare_encrypted_rootfs(id, execution_fence, None) - .instrument(tracing::info_span!( - "sbxmgr.start.encrypted_rootfs_prepare", - sandbox_id = %id - )) - .await + match crate::rootfs_encryption::prepare_encrypted_rootfs( + id, + execution_fence, + None, + &self.node_manager, + ) + .instrument(tracing::info_span!( + "sbxmgr.start.encrypted_rootfs_prepare", + sandbox_id = %id + )) + .await { Ok(path) => Some(path), Err(e) => { @@ -2161,12 +2175,16 @@ impl SandboxManagerService { // it; booting it past an existing COW (cow.plain, or cow.luks // from an unstamped re-home) silently rolls back every write the // guest made through it. - match crate::rootfs_encryption::prepare_snapshot_rootfs(id, execution_fence) - .instrument(tracing::info_span!( - "sbxmgr.start.pool_snapshot_rootfs_prepare", - sandbox_id = %id - )) - .await + match crate::rootfs_encryption::prepare_snapshot_rootfs( + id, + execution_fence, + &self.node_manager, + ) + .instrument(tracing::info_span!( + "sbxmgr.start.pool_snapshot_rootfs_prepare", + sandbox_id = %id + )) + .await { Ok(path) => Some(path), Err(e) diff --git a/crates/libs/sbx/mgr/server/src/service/startup_resident.rs b/crates/libs/sbx/mgr/server/src/service/startup_resident.rs index 57f174ad88..fcfb79e693 100644 --- a/crates/libs/sbx/mgr/server/src/service/startup_resident.rs +++ b/crates/libs/sbx/mgr/server/src/service/startup_resident.rs @@ -37,7 +37,12 @@ impl SandboxManagerService { row, Path::new(RUN_DIR), |id| async move { - observe_firecracker_vm_state(&id, "startup resident recovery").await + observe_firecracker_vm_state( + &id, + "startup resident recovery", + &self.node_manager, + ) + .await }, |target| self.prepare_startup_resident(&prepare_row, target), ) @@ -89,7 +94,12 @@ impl SandboxManagerService { row, Path::new(RUN_DIR), |id| async move { - observe_firecracker_vm_state(&id, "periodic resident recovery").await + observe_firecracker_vm_state( + &id, + "periodic resident recovery", + &self.node_manager, + ) + .await }, |target| self.prepare_startup_resident(&prepare_row, target), ) diff --git a/crates/libs/sbx/mgr/server/src/service/tail_support.rs b/crates/libs/sbx/mgr/server/src/service/tail_support.rs index 952d7ea0af..356197cadf 100644 --- a/crates/libs/sbx/mgr/server/src/service/tail_support.rs +++ b/crates/libs/sbx/mgr/server/src/service/tail_support.rs @@ -189,9 +189,11 @@ pub(super) enum NodeManagerConnectError { Client(String), } -pub(super) async fn connect_node() --> Result { - pplx_space_node_mgr_client::connect() +pub(super) async fn connect_node( + node_manager: &crate::node_manager_transport::NodeManagerConnector, +) -> Result { + node_manager + .connect() .await .map_err(|e| NodeManagerConnectError::Client(e.to_string())) } diff --git a/crates/libs/sbx/mgr/server/src/service/tests.rs b/crates/libs/sbx/mgr/server/src/service/tests.rs index 7ee38f4d7b..cd8073a246 100644 --- a/crates/libs/sbx/mgr/server/src/service/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/tests.rs @@ -239,6 +239,9 @@ mod cancelled_cleanup_rpc_tests; #[cfg(test)] mod local_backup_inventory_tests; +#[cfg(test)] +pub(crate) mod node_manager_transport_tests; + #[cfg(test)] pub(crate) mod service_tests { use super::start_support::*; @@ -5406,10 +5409,14 @@ pub(crate) mod service_tests { let sandbox_id = SandboxId::new().to_string(); insert_orphaned_test_sandbox(&conn, &sandbox_id, SandboxState::Failed); let row = crate::db::get_sandbox(&conn, &sandbox_id).unwrap().unwrap(); + let templates = crate::template_resolver::fake::FakeTemplateResolver::refusing( + Status::not_found("orphaned sandboxes must not look up templates"), + ); - let error = template_source_for_sandbox_prepare(&row) + let error = template_source_for_sandbox_prepare(&templates, &row) .await .expect_err("orphaned sandbox without rootfs is not startable"); + assert!(templates.requests().is_empty()); assert_eq!(error.code(), tonic::Code::FailedPrecondition); assert!(error.message().contains("cannot start orphaned sandbox")); @@ -10951,39 +10958,6 @@ pub(crate) mod service_tests { ); } - #[tokio::test(start_paused = true)] - async fn fatal_kill_signal_failure_still_schedules_admitted_cleanup() { - let conn = crate::db::test_conn(); - let sandbox_id = SandboxId::new(); - let sandbox_id_string = sandbox_id.to_string(); - insert_test_sandbox(&conn, &sandbox_id_string, SandboxState::Started); - let service = service_for_teardown_test(conn); - service - .update_runtime_fields( - sandbox_id, - runtime_fields(ExecutionGeneration::new(1).unwrap(), 7), - ) - .unwrap(); - - let error = service - .kill(internal_request(KillRequest { - sandbox_id: sandbox_id_string.clone(), - signal: 9, - blocking: false, - })) - .await - .expect_err("the unavailable node-manager must fail fatal signal dispatch"); - assert_eq!(error.code(), tonic::Code::Internal); - - tokio::time::timeout(Duration::from_secs(1), async { - while service.cleanup_executor.active.is_empty() { - tokio::time::sleep(Duration::from_millis(1)).await; - } - }) - .await - .expect("fatal signal failure must still hand admitted cleanup to a runner"); - } - #[tokio::test] async fn unpersisted_delete_clears_activity_without_replacing_live_lock_identity() { let service = SandboxManagerService::new(rusqlite::Connection::open_in_memory().unwrap()); @@ -12431,6 +12405,8 @@ pub(crate) mod service_tests { node_role: None, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), } } diff --git a/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs b/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs new file mode 100644 index 0000000000..879bbdfc5d --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs @@ -0,0 +1,148 @@ +//! Lifecycle handlers driven against an in-process node-manager. +//! +//! The service is built with every host dependency redirected: the +//! node-manager connector reaches a [`FakeNodeManager`], the host-control and +//! credential-manager sockets are test-owned paths nothing binds, and no +//! runtime is ever live. + +use super::service_tests::{insert_test_sandbox, runtime_fields}; +use super::*; +use crate::node_manager_transport::fake::FakeNodeManager; +use crate::service::activation_support::internal_request; +use pplx_space_common_core::types::ExecutionGeneration; +use pplx_space_node_mgr_client::proto::{ + CleanupSandboxRequest, CleanupSandboxResponse, StopProcessRequest, StopProcessResponse, +}; + +/// A service whose node-manager clients all come from `node_manager` and whose +/// template lookups all go to `template_resolver`. +pub(crate) fn isolated_service( + conn: rusqlite::Connection, + node_manager: NodeManagerConnector, + template_resolver: Arc, + sockets: &Path, +) -> SandboxManagerService { + SandboxManagerService::new_with_options( + conn, + SandboxManagerServiceOptions { + host_control_socket_path: sockets.join("absent-host-control.sock"), + credential_manager_socket_path: sockets.join("absent-credential-manager.sock"), + sandbox_runtime_live: Arc::new(|_| false), + shutdown_token: CancellationToken::new(), + task_tracker: TaskTracker::new(), + local_node_id: Some(NodeId::from_uuid( + uuid::Uuid::parse_str("019f8000-0000-7000-8000-0000000000fe").unwrap(), + )), + node_role: None, + node_managed_by: None, + route_ed25519_verification_bundle: None, + node_manager, + template_resolver, + }, + ) +} + +/// [`isolated_service`] whose template lookups are all refused. +pub(crate) fn service_with_node_manager( + conn: rusqlite::Connection, + node_manager: NodeManagerConnector, + sockets: &Path, +) -> SandboxManagerService { + isolated_service( + conn, + node_manager, + Arc::new( + crate::template_resolver::fake::FakeTemplateResolver::refusing(Status::unavailable( + "template-manager is outside this test", + )), + ), + sockets, + ) +} + +#[tokio::test] +async fn fatal_kill_signal_failure_still_schedules_admitted_cleanup() { + let stops = Arc::new(Mutex::new(Vec::new())); + let cleanups = Arc::new(Mutex::new(Vec::new())); + let node_manager = FakeNodeManager::default() + .unary("StopProcess", { + let stops = Arc::clone(&stops); + move |request: StopProcessRequest| { + stops.lock().unwrap().push(request); + Err::(Status::unavailable("injected StopProcess refusal")) + } + }) + .unary("CleanupSandbox", { + let cleanups = Arc::clone(&cleanups); + move |request: CleanupSandboxRequest| { + cleanups.lock().unwrap().push(request); + Ok(CleanupSandboxResponse::default()) + } + }) + .serve(); + let conn = crate::db::test_conn(); + let sandbox_id = SandboxId::new(); + let sandbox_id_string = sandbox_id.to_string(); + insert_test_sandbox(&conn, &sandbox_id_string, SandboxState::Started); + let sockets = tempfile::tempdir().unwrap(); + let service = service_with_node_manager(conn, node_manager.connector(), sockets.path()); + service + .update_runtime_fields( + sandbox_id, + runtime_fields(ExecutionGeneration::new(1).unwrap(), 7), + ) + .unwrap(); + let execution_fence = service + .lookup_sandbox(&sandbox_id_string) + .unwrap() + .node_execution_fence(); + assert!(execution_fence.is_some()); + + let error = service + .kill(internal_request(KillRequest { + sandbox_id: sandbox_id_string.clone(), + signal: 9, + blocking: false, + })) + .await + .expect_err("node-manager's refusal must fail fatal signal dispatch"); + assert_eq!(error.code(), tonic::Code::Internal); + assert!( + error.message().contains("send fatal signal failed") + && error.message().contains("injected StopProcess refusal"), + "the caller must see node-manager's refusal: {}", + error.message() + ); + // The signal reached node-manager once, fenced to the observed execution. + assert_eq!( + *stops.lock().unwrap(), + [StopProcessRequest { + execution_fence: execution_fence.clone(), + sandbox_id: sandbox_id_string.clone(), + signal: 9, + pid: None, + }] + ); + + // The cleanup admitted before the signal still runs: node-manager stops + // the same execution, keeping its filesystem, and the row ends stopped. + tokio::time::timeout(Duration::from_secs(5), async { + while service.lookup_sandbox(&sandbox_id_string).unwrap().state != SandboxState::Stopped { + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await + .expect("fatal signal failure must still run the admitted cleanup to completion"); + assert_eq!( + *cleanups.lock().unwrap(), + [CleanupSandboxRequest { + execution_fence, + sandbox_id: sandbox_id_string, + force: true, + keep_filesystem: true, + return_after_slot_release: true, + immediate_kill: true, + }] + ); + assert_eq!(node_manager.calls(), ["StopProcess", "CleanupSandbox"]); +} diff --git a/crates/libs/sbx/mgr/server/src/snapshot_encryption.rs b/crates/libs/sbx/mgr/server/src/snapshot_encryption.rs index 9de0c7a02f..933b41a7c4 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_encryption.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_encryption.rs @@ -93,8 +93,9 @@ pub(crate) async fn prepare_encrypted_snapshot_volume( execution_fence: Option, fresh: bool, restore_source: Option<&str>, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result { - let mut node = pplx_space_node_mgr_client::connect().await?; + let mut node = node_manager.connect().await?; let prepared = node .prepare_encrypted_volume( pplx_space_node_mgr_client::proto::PrepareEncryptedVolumeRequest { @@ -123,8 +124,10 @@ pub(crate) async fn release_encrypted_snapshot_volume( sandbox_id: &str, execution_fence: Option, stage_dir: &str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result<(), String> { - let mut node = pplx_space_node_mgr_client::connect() + let mut node = node_manager + .connect() .await .map_err(|e| format!("node-manager connect for snapshot-volume release: {e}"))?; // The two sealed ciphertext artifacts a suspend publishes: the volume's own diff --git a/crates/libs/sbx/mgr/server/src/snapshot_gc.rs b/crates/libs/sbx/mgr/server/src/snapshot_gc.rs index a24d405484..3755f10a73 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_gc.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_gc.rs @@ -78,6 +78,7 @@ pub(crate) struct GcOwnership { pub coordinator: crate::auto_lifecycle::AutoLifecycleCoordinator, pub drain: Arc, pub shutdown: CancellationToken, + pub node_manager: crate::node_manager_transport::NodeManagerConnector, } #[derive(Debug, thiserror::Error)] @@ -527,7 +528,12 @@ async fn gc_tick( // until the next sbx-mgr startup sweep, which now // also goes through this RPC. if let Some(staging) = staging_to_clean { - delete_cache_artifact_via_node_mgr(&staging, &candidate.key).await; + delete_cache_artifact_via_node_mgr( + &staging, + &candidate.key, + &ownership.node_manager, + ) + .await; } } Ok((skip, _)) => { @@ -587,8 +593,12 @@ const FORK_COPY_ORPHAN_MIN_AGE_SECS: u64 = 60 * 60; /// failure rather than abort, because the parent operation has /// already committed (the row is gone, or this was orphan cleanup /// to begin with). -async fn delete_cache_artifact_via_node_mgr(staging_path: &str, candidate_label: &str) { - let mut client = match pplx_space_node_mgr_client::connect().await { +async fn delete_cache_artifact_via_node_mgr( + staging_path: &str, + candidate_label: &str, + node_manager: &crate::node_manager_transport::NodeManagerConnector, +) { + let mut client = match node_manager.connect().await { Ok(c) => c, Err(err) => { warn!( @@ -702,7 +712,12 @@ async fn sweep_orphans( }; // Await the admitted mutation; caller cancellation stops the next // admission instead of dropping this mutation and its reservation. - delete_cache_artifact_via_node_mgr(&path.to_string_lossy(), "orphan-sweep").await; + delete_cache_artifact_via_node_mgr( + &path.to_string_lossy(), + "orphan-sweep", + &ownership.node_manager, + ) + .await; } } diff --git a/crates/libs/sbx/mgr/server/src/snapshot_gc/tests.rs b/crates/libs/sbx/mgr/server/src/snapshot_gc/tests.rs index db553d12a9..da09a3ef8c 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_gc/tests.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_gc/tests.rs @@ -13,6 +13,7 @@ fn test_ownership() -> GcOwnership { coordinator: crate::auto_lifecycle::AutoLifecycleCoordinator::default(), drain: Arc::new(crate::service::drain::DrainState::default()), shutdown: CancellationToken::new(), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), } } @@ -645,6 +646,7 @@ async fn gc_defers_to_publication_staging_drain_and_shutdown_ownership() { coordinator: ownership.coordinator.clone(), drain: crate::service::tests::closed_drain_for_test(), shutdown: CancellationToken::new(), + node_manager: ownership.node_manager.clone(), }; let report = gc_tick( artifacts.clone(), diff --git a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs index a0876e6670..83252de5b6 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs @@ -1256,6 +1256,8 @@ pub async fn prepare_restore_dir( sandbox_id: &str, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, transfer_scope_id: pplx_space_common_core::types::TransferScopeId, + node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result { let latest = latest_ready_snapshot(sandbox_id).await?; let staged = restore_pull_metadata_from_source( @@ -1264,7 +1266,7 @@ pub async fn prepare_restore_dir( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1278,6 +1280,8 @@ pub async fn prepare_restore_dir_from_uri( execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, storage_uri: &str, transfer_scope_id: pplx_space_common_core::types::TransferScopeId, + node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result { let config = SnapshotUploadConfig::from_uri(storage_uri) .map_err(|e| RestoreError::Other(format!("parse storage_uri: {e}")))?; @@ -1290,7 +1294,7 @@ pub async fn prepare_restore_dir_from_uri( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1849,6 +1853,8 @@ fn preserve_attributed_metadata_retry( pub async fn restore_pull_rootfs( stage: &RestoreMetadataStaged, execution_fence: pplx_space_node_mgr_client::execution_fence::ExecutionFence, + node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result<(), RestoreError> { use pplx_space_node_mgr_client::execution_fence::InstanceTarget; let target = stage.filesystem_scope_id.map_or( @@ -1879,7 +1885,7 @@ pub async fn restore_pull_rootfs( ) .await; if let Err(ref e) = stage_result { - handle_restore_stage_failure(stage, target, e).await; + handle_restore_stage_failure(stage, target, e, node_manager).await; } return stage_result; } @@ -1894,6 +1900,8 @@ pub async fn restore_pull_rootfs( crate::snapshot_upload::artifact_kind_from_meta(&stage.meta), stage.transfer_scope_id, target, + node_manager, + templates, ) .await } @@ -1919,7 +1927,7 @@ pub async fn restore_pull_rootfs( }; if let Err(ref e) = stage_result { - handle_restore_stage_failure(stage, target, e).await; + handle_restore_stage_failure(stage, target, e, node_manager).await; } stage_result } @@ -1928,6 +1936,7 @@ async fn handle_restore_stage_failure( stage: &RestoreMetadataStaged, target: pplx_space_node_mgr_client::execution_fence::InstanceTarget, error: &RestoreError, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) { if matches!(error, RestoreError::AdmissionRejected(_)) { preserve_completed_restore_artifacts(stage); @@ -1944,7 +1953,13 @@ async fn handle_restore_stage_failure( snapshot_dir = %stage.snapshot_dir.display(), "cross-node restore staging failed; removing partial artifacts" ); - cleanup_partial_restore_artifacts(&stage.sandbox_dir, &stage.meta.sandbox_id, target).await; + cleanup_partial_restore_artifacts( + &stage.sandbox_dir, + &stage.meta.sandbox_id, + target, + node_manager, + ) + .await; } fn preserve_completed_restore_artifacts(stage: &RestoreMetadataStaged) { @@ -1990,6 +2005,7 @@ async fn cleanup_partial_restore_artifacts( sandbox_dir: &Path, sandbox_id: &str, target: pplx_space_node_mgr_client::execution_fence::InstanceTarget, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) { use pplx_space_node_mgr_client::execution_fence::InstanceTarget; let scope = match target { @@ -2000,7 +2016,7 @@ async fn cleanup_partial_restore_artifacts( sandbox_dir, sandbox_id, scope, - crate::btrfs_via_nodemgr::delete_sandbox_instance(sandbox_id, target), + crate::btrfs_via_nodemgr::delete_sandbox_instance(sandbox_id, target, node_manager), ) .await; } @@ -2492,6 +2508,10 @@ async fn stage_encrypted_restore( /// deterministic `received_uuid`. Failure on any of these surfaces as /// `RestoreError::Other` with a precise diagnostic — the caller is /// responsible for cleanup of the partial staging directory. +#[expect( + clippy::too_many_arguments, + reason = "each argument is an independent restore input; the node-manager and template-manager clients are injected" +)] async fn stage_btrfs_send_restore( identity: RestoreTransferIdentity<'_>, config: &SnapshotUploadConfig, @@ -2500,6 +2520,8 @@ async fn stage_btrfs_send_restore( artifact_kind: ArtifactKind, transfer_scope_id: Option, target: pplx_space_node_mgr_client::execution_fence::InstanceTarget, + node_manager: &crate::node_manager_transport::NodeManagerConnector, + templates: &dyn crate::template_resolver::TemplateResolver, ) -> Result<(), RestoreError> { let sandbox_id = identity.sandbox_id; let sidecar_bytes = std::fs::read(sidecar_path).map_err(|e| { @@ -2551,7 +2573,7 @@ async fn stage_btrfs_send_restore( } run_btrfs_staging( Box::pin(async { - ensure_local_template(sandbox_id, template_name, content_hash) + ensure_local_template(templates, sandbox_id, template_name, content_hash) .await .map(drop) }), @@ -2564,6 +2586,7 @@ async fn stage_btrfs_send_restore( template_name, content_hash, target, + node_manager, ) .instrument(tracing::info_span!( "sbxmgr.start.stage", @@ -2636,31 +2659,35 @@ async fn run_btrfs_staging<'a>( /// tmplmgr — same RPC the sandbox-create path uses. Returns the local /// directory the template lives at. async fn ensure_local_template( + templates: &dyn crate::template_resolver::TemplateResolver, sandbox_id: &str, template_name: &str, content_hash: &str, ) -> Result { + use crate::template_resolver::TemplateLookupError; async { - let mut client = pplx_space_template_mgr_client::connect() + let resp = templates + .get_template(tonic::Request::new( + pplx_space_template_mgr_client::proto::GetTemplateRequest { + caller_org_id: None, + template: Some(pplx_space_template_mgr_client::proto::TemplateRef { + name: template_name.to_string(), + content_hash: Some(content_hash.to_string()), + tag: None, + }), + vcpu: None, + memory_mib: None, + }, + )) .await - .map_err(|e| { - RestoreError::Other(format!( + .map_err(|error| match error { + TemplateLookupError::Connect(e) => RestoreError::Other(format!( "tmplmgr connect for GetTemplate({template_name}): {e}" - )) - })?; - let resp = client - .get_template(pplx_space_template_mgr_client::proto::GetTemplateRequest { - caller_org_id: None, - template: Some(pplx_space_template_mgr_client::proto::TemplateRef { - name: template_name.to_string(), - content_hash: Some(content_hash.to_string()), - tag: None, - }), - vcpu: None, - memory_mib: None, - }) - .await - .map_err(|error| template_pull_error(template_name, content_hash, error))? + )), + TemplateLookupError::Rpc(error) => { + template_pull_error(template_name, content_hash, error) + } + })? .into_inner(); if resp.template_dir.is_empty() { return Err(RestoreError::Other(format!( diff --git a/crates/libs/sbx/mgr/server/src/snapshot_restore/tests.rs b/crates/libs/sbx/mgr/server/src/snapshot_restore/tests.rs index 7e6c165ded..56b73d4fdf 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_restore/tests.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_restore/tests.rs @@ -1009,7 +1009,13 @@ async fn admission_rejection_preserves_completed_restore_work_for_retry() { None, ), ); - handle_restore_stage_failure(&stage, target, &error).await; + handle_restore_stage_failure( + &stage, + target, + &error, + &crate::node_manager_transport::NodeManagerConnector::host(), + ) + .await; assert!( sandbox_dir.exists(), diff --git a/crates/libs/sbx/mgr/server/src/snapshot_scheduler.rs b/crates/libs/sbx/mgr/server/src/snapshot_scheduler.rs index 7bf9f54e99..0b0e58c5ff 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_scheduler.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_scheduler.rs @@ -230,6 +230,7 @@ pub struct CheckpointContext { pub db: Arc>, pub auto_lifecycle: crate::auto_lifecycle::AutoLifecycleCoordinator, pub health_lifecycle: crate::watchdog::HealthLifecycleEpochs, + pub node_manager: crate::node_manager_transport::NodeManagerConnector, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -421,7 +422,9 @@ async fn process_tick( checkpoint_tier: Option, ctx: &CheckpointContext, ) -> Result { - let mut node_client = pplx_space_node_mgr_client::connect() + let mut node_client = ctx + .node_manager + .connect() .await .map_err(|e| e.to_string())?; @@ -687,6 +690,7 @@ async fn rollback_checkpoint_failure( let observation = observe_firecracker_vm_state( &sandbox_id_typed, "checkpoint rollback resume reconciliation", + &ctx.node_manager, ) .await; if let Some(probe_error) = observation.unobservable_reason() @@ -840,7 +844,7 @@ pub async fn do_checkpoint( } // 2. Connect to node-manager - let mut node_client = match pplx_space_node_mgr_client::connect().await { + let mut node_client = match ctx.node_manager.connect().await { Ok(node_client) => node_client, Err(error) => { restore_checkpoint_state_best_effort(ctx, sandbox_id, previous_state); @@ -976,6 +980,7 @@ pub async fn do_checkpoint( let observation = observe_firecracker_vm_state( &sandbox_id_typed, "completed checkpoint resume reconciliation", + &ctx.node_manager, ) .await; let reason = match observation.unobservable_reason() { diff --git a/crates/libs/sbx/mgr/server/src/snapshot_scheduler/tests.rs b/crates/libs/sbx/mgr/server/src/snapshot_scheduler/tests.rs index 9e28b213d5..a20f7911d1 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_scheduler/tests.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_scheduler/tests.rs @@ -143,6 +143,7 @@ fn checkpoint_context(state: SandboxState) -> (CheckpointContext, String) { db: Arc::new(std::sync::Mutex::new(conn)), auto_lifecycle: crate::auto_lifecycle::AutoLifecycleCoordinator::default(), health_lifecycle: crate::watchdog::HealthLifecycleEpochs::default(), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), }, sandbox_id, ) diff --git a/crates/libs/sbx/mgr/server/src/snapshot_upload.rs b/crates/libs/sbx/mgr/server/src/snapshot_upload.rs index 057c2113e0..f88c2582a1 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_upload.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_upload.rs @@ -950,6 +950,7 @@ pub(crate) fn prepare_suspend_snapshot( /// replay after a worker restart rejoins vol-mgr's persisted transfer jobs. pub async fn submit_prepared_snapshot( prepared: &SnapshotUploadJobs, + node_manager: &crate::node_manager_transport::NodeManagerConnector, ) -> Result { let mut prepared = prepared.clone(); if !prepared.deferred_metadata @@ -963,6 +964,7 @@ pub async fn submit_prepared_snapshot( &prepared.meta.sandbox_id, &prepared.snapshot_dir, source_id, + node_manager, ) .await .map_err(|error| error.to_string())?; diff --git a/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs.rs b/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs.rs index bf385c1c4d..31141ff26f 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs.rs @@ -72,6 +72,7 @@ const TRANSFER_PROBE_TIMEOUT: Duration = Duration::from_secs(10); pub(crate) fn build_worker( db: std::sync::Arc>, drain_state: std::sync::Arc, + node_manager: crate::node_manager_transport::NodeManagerConnector, ) -> BackgroundWorker { let config = WorkerConfig::new(DEFAULT_CONCURRENCY, JOB_TIMEOUT, MAX_ATTEMPTS) .with_kind_concurrency(SNAPSHOT_UPLOAD_JOB_KIND, DEFAULT_CONCURRENCY) @@ -80,8 +81,9 @@ pub(crate) fn build_worker( db: std::sync::Arc::clone(&db), drain_state, retained_source_inspector: std::sync::Arc::new( - crate::retained_source::NodeManagerSourceInspector, + crate::retained_source::NodeManagerSourceInspector(node_manager.clone()), ), + node_manager, }; BackgroundWorker::new(db, executor, config) } @@ -109,6 +111,7 @@ pub struct SnapshotUploadExecutor { /// Sizes a retained source before its stream is reclaimed, so the local /// sidecar can switch to counting the image the artifact keeps. retained_source_inspector: std::sync::Arc, + node_manager: crate::node_manager_transport::NodeManagerConnector, } #[async_trait::async_trait] @@ -170,11 +173,12 @@ impl JobExecutor for SnapshotUploadExecutor { // jobs survive it. Persist the durable transfer handles before allowing // a teardown-gated caller to release the runtime. if payload.job_ids.is_empty() { - payload = crate::snapshot_upload::submit_prepared_snapshot(&payload) - .await - .map_err(|error| { - JobError::retriable(format!("vol-mgr submit prepared snapshot: {error}")) - })?; + payload = + crate::snapshot_upload::submit_prepared_snapshot(&payload, &self.node_manager) + .await + .map_err(|error| { + JobError::retriable(format!("vol-mgr submit prepared snapshot: {error}")) + })?; persist_submitted_transfer_ids(&self.db, &job.job_id, &payload)?; } diff --git a/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs/tests.rs b/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs/tests.rs index 16a2e73939..c164544270 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs/tests.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_upload_jobs/tests.rs @@ -37,6 +37,7 @@ async fn upload_executor_converges_unstaged_metadata_failure() { let executor = SnapshotUploadExecutor { db: std::sync::Arc::new(std::sync::Mutex::new(open_backfill_test_db())), drain_state: crate::service::tests::closed_drain_for_test(), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), retained_source_inspector: std::sync::Arc::new(crate::retained_source::FakeInspector::new( crate::retained_source::RetainedSourceCheck::Absent, )), @@ -184,6 +185,7 @@ async fn executor_rejects_non_snapshot_upload_job_kind() { rusqlite::Connection::open_in_memory().unwrap(), )), drain_state: std::sync::Arc::new(crate::service::drain::DrainState::default()), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), retained_source_inspector: std::sync::Arc::new(crate::retained_source::FakeInspector::new( crate::retained_source::RetainedSourceCheck::Absent, )), @@ -216,6 +218,7 @@ async fn executor_requires_begin_before_volume_staging() { rusqlite::Connection::open_in_memory().unwrap(), )), drain_state: std::sync::Arc::new(crate::service::drain::DrainState::default()), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), retained_source_inspector: std::sync::Arc::new(crate::retained_source::FakeInspector::new( crate::retained_source::RetainedSourceCheck::Absent, )), @@ -361,6 +364,7 @@ fn ready_backup_keeps_completed_evidence_when_bookkeeping_misses_execution_deadl executor: SnapshotUploadExecutor { db: std::sync::Arc::clone(&db), drain_state: std::sync::Arc::new(crate::service::drain::DrainState::default()), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), retained_source_inspector: std::sync::Arc::new( crate::retained_source::FakeInspector::new( crate::retained_source::RetainedSourceCheck::Absent, @@ -449,7 +453,11 @@ async fn maintenance_retries_failed_pruning_on_the_next_tick() { .unwrap(); let db = std::sync::Arc::new(std::sync::Mutex::new(conn)); let drain = std::sync::Arc::new(crate::service::drain::DrainState::default()); - let worker = build_worker(std::sync::Arc::clone(&db), std::sync::Arc::clone(&drain)); + let worker = build_worker( + std::sync::Arc::clone(&db), + std::sync::Arc::clone(&drain), + crate::node_manager_transport::NodeManagerConnector::host(), + ); let shutdown = tokio_util::sync::CancellationToken::new(); let maintenance = tokio::spawn(run_maintenance( std::sync::Arc::clone(&db), @@ -1129,6 +1137,7 @@ fn ready_backup_returns_before_the_retained_source_is_inspected() { let executor = SnapshotUploadExecutor { db: std::sync::Arc::clone(&db), drain_state: std::sync::Arc::new(crate::service::drain::DrainState::default()), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), retained_source_inspector: std::sync::Arc::new(GatedInspector { gate: std::sync::Arc::clone(&gate), }), diff --git a/crates/libs/sbx/mgr/server/src/template_resolver.rs b/crates/libs/sbx/mgr/server/src/template_resolver.rs new file mode 100644 index 0000000000..36d1ed34c9 --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/template_resolver.rs @@ -0,0 +1,67 @@ +//! How sandbox-manager looks up template builds. +//! +//! Create, start, wake, pool cloning, restore staging and the template +//! network-policy cache ask template-manager's `GetTemplate` through one +//! [`TemplateResolver`], carried in the service options. Production uses +//! [`HostTemplateManager`], which reaches template-manager's fixed Unix socket +//! through the `pplx-space-template-mgr-client` constructors exactly as each +//! caller did directly. Callers keep their own error mapping. Tests substitute +//! a resolver that records requests and returns scripted answers. + +use pplx_space_template_mgr_client::ConnectError; +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, GetTemplateResponse}; + +/// Why a [`TemplateResolver::get_template`] call failed. The connect and RPC +/// failures stay distinct because callers report them differently. +#[derive(Debug, thiserror::Error)] +pub enum TemplateLookupError { + #[error(transparent)] + Connect(#[from] ConnectError), + #[error(transparent)] + Rpc(#[from] tonic::Status), +} + +/// Requests and responses cross the resolver whole, so callers can attach and +/// read gRPC metadata as they could on the client directly. +#[async_trait::async_trait] +pub trait TemplateResolver: Send + Sync { + /// `GetTemplate` on a client dialed for this call under the default + /// connect-retry policy. + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError>; + + /// `GetTemplate` on the process-wide reconnecting channel, which dials on + /// first use. Connection failures surface as the RPC status. + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status>; +} + +/// Template-manager at [`pplx_space_template_mgr_client::SOCKET_PATH`]. +pub struct HostTemplateManager; + +#[async_trait::async_trait] +impl TemplateResolver for HostTemplateManager { + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + let mut client = pplx_space_template_mgr_client::connect().await?; + Ok(client.get_template(request).await?) + } + + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + pplx_space_template_mgr_client::lazy() + .get_template(request) + .await + } +} + +#[cfg(test)] +pub(crate) mod fake; diff --git a/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs b/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs new file mode 100644 index 0000000000..1fed44511c --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/template_resolver/fake.rs @@ -0,0 +1,108 @@ +//! A scripted template resolver for tests. + +use std::sync::{Mutex, PoisonError}; + +use pplx_space_template_mgr_client::proto::{GetTemplateRequest, GetTemplateResponse}; + +use super::{TemplateLookupError, TemplateResolver}; + +type Answer = + dyn Fn() -> Result, TemplateLookupError> + Send + Sync; + +/// Which [`TemplateResolver`] method a lookup used. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum LookupChannel { + /// [`TemplateResolver::get_template`]: a client dialed for the call. + Dialed, + /// [`TemplateResolver::get_template_on_shared_channel`]. + Shared, +} + +/// One lookup as the resolver received it. +#[derive(Clone, Debug)] +pub(crate) struct Lookup { + pub(crate) channel: LookupChannel, + pub(crate) metadata: tonic::metadata::MetadataMap, + pub(crate) request: GetTemplateRequest, +} + +/// Answers every lookup with a scripted result and records each lookup in +/// arrival order. A scripted connect failure reaches a shared-channel lookup +/// as `Unavailable`, as a lazy channel reports it. +pub(crate) struct FakeTemplateResolver { + answer: Box, + lookups: Mutex>, +} + +impl FakeTemplateResolver { + pub(crate) fn new( + answer: impl Fn() -> Result, TemplateLookupError> + + Send + + Sync + + 'static, + ) -> Self { + Self { + answer: Box::new(answer), + lookups: Mutex::new(Vec::new()), + } + } + + /// Refuse every lookup with `status`. + pub(crate) fn refusing(status: tonic::Status) -> Self { + Self::new(move || Err(TemplateLookupError::Rpc(status.clone()))) + } + + pub(crate) fn lookups(&self) -> Vec { + self.lookups + .lock() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + pub(crate) fn requests(&self) -> Vec { + self.lookups() + .into_iter() + .map(|lookup| lookup.request) + .collect() + } + + fn answer( + &self, + channel: LookupChannel, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + let (metadata, _, request) = request.into_parts(); + self.lookups + .lock() + .unwrap_or_else(PoisonError::into_inner) + .push(Lookup { + channel, + metadata, + request, + }); + (self.answer)() + } +} + +#[async_trait::async_trait] +impl TemplateResolver for FakeTemplateResolver { + async fn get_template( + &self, + request: tonic::Request, + ) -> Result, TemplateLookupError> { + self.answer(LookupChannel::Dialed, request) + } + + async fn get_template_on_shared_channel( + &self, + request: tonic::Request, + ) -> Result, tonic::Status> { + self.answer(LookupChannel::Shared, request) + .map_err(|error| match error { + TemplateLookupError::Rpc(status) => status, + TemplateLookupError::Connect(error) => { + tonic::Status::unavailable(error.to_string()) + } + }) + } +} diff --git a/crates/libs/sbx/mgr/server/src/watchdog.rs b/crates/libs/sbx/mgr/server/src/watchdog.rs index a32de568a5..cc0ba174cf 100644 --- a/crates/libs/sbx/mgr/server/src/watchdog.rs +++ b/crates/libs/sbx/mgr/server/src/watchdog.rs @@ -30,6 +30,7 @@ use pplx_space_metrics::sandbox::worker::{ use pplx_space_node_runtime_core::fc_api::FirecrackerVmState; use crate::auto_lifecycle::AutoLifecycleCoordinator; +use crate::node_manager_transport::NodeManagerConnector; use crate::node_state_sync::mark_sandbox_dirty; use crate::service::network_policy::NetworkPolicyReconcileQueue; use crate::sidecar::{sidecar_alive, sidecar_live_pid}; @@ -689,6 +690,7 @@ pub struct SandboxWatchdog { network_policy_reconcile_queue: NetworkPolicyReconcileQueue, expected_route_verification_key_generation: Arc, feature_flags: pplx_space_feature_flags::FeatureFlagClient, + node_manager: NodeManagerConnector, } #[derive(Clone, Default)] @@ -730,6 +732,7 @@ impl SandboxWatchdog { shutdown_token: CancellationToken, drain_state: Arc, task_tracker: &TaskTracker, + node_manager: NodeManagerConnector, ) -> Self { let state = lifecycle_epochs.state.clone(); let sidecars: Arc>> = @@ -743,6 +746,7 @@ impl SandboxWatchdog { Arc::clone(&expected_route_verification_key_generation); let task_feature_flags = feature_flags.clone(); let task_snapshot_scheduler = snapshot_scheduler.clone(); + let task_node_manager = node_manager.clone(); let sidecar_shutdown = shutdown_token.child_token(); let sidecar_guard = crate::service::drain::CriticalTaskGuard::new( "sidecar watchdog", @@ -760,6 +764,7 @@ impl SandboxWatchdog { task_expected_route_verification_key_generation, task_feature_flags, sidecar_shutdown, + task_node_manager, ) .await; }); @@ -777,6 +782,7 @@ impl SandboxWatchdog { auto_lifecycle, lifecycle_epochs, shutdown_token.child_token(), + node_manager.clone(), ); task_tracker.spawn(async move { let _guard = health_guard; @@ -789,6 +795,7 @@ impl SandboxWatchdog { network_policy_reconcile_queue, expected_route_verification_key_generation, feature_flags, + node_manager, } } @@ -806,6 +813,7 @@ impl SandboxWatchdog { db: db.clone(), auto_lifecycle: auto_lifecycle.clone(), health_lifecycle: lifecycle_epochs.clone(), + node_manager: NodeManagerConnector::host(), }, ))); Self::new( @@ -818,6 +826,7 @@ impl SandboxWatchdog { shutdown_token, Arc::new(crate::service::drain::DrainState::default()), task_tracker, + NodeManagerConnector::host(), ) } @@ -1008,6 +1017,7 @@ impl SandboxWatchdog { &self.network_policy_reconcile_queue, &self.state, &self.feature_flags, + &self.node_manager, ) .await } @@ -1228,6 +1238,10 @@ fn system_time_to_timestamp(value: SystemTime) -> Option }) } +#[expect( + clippy::too_many_arguments, + reason = "the loop receives each long-lived watchdog dependency from the constructor" +)] async fn sidecar_loop( state: Arc>>, sidecars: Arc>>, @@ -1236,6 +1250,7 @@ async fn sidecar_loop( expected_route_verification_key_generation: Arc, feature_flags: pplx_space_feature_flags::FeatureFlagClient, shutdown_token: CancellationToken, + node_manager: NodeManagerConnector, ) { let mut interval = tokio::time::interval(SIDECAR_CHECK_INTERVAL); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); @@ -1260,6 +1275,7 @@ async fn sidecar_loop( &network_policy_reconcile_queue, &state, &feature_flags, + &node_manager, ) .await; false @@ -1383,6 +1399,7 @@ async fn sidecar_loop( &network_policy_reconcile_queue, &state, &feature_flags, + &node_manager, ) .await; @@ -1410,6 +1427,10 @@ async fn sidecar_loop( } } +#[expect( + clippy::too_many_arguments, + reason = "each restart task needs the sidecar loop's long-lived dependencies" +)] async fn fill_sidecar_restart_slots( queue: &mut SidecarRestartQueue, futures: &mut FuturesUnordered, @@ -1418,6 +1439,7 @@ async fn fill_sidecar_restart_slots( network_policy_reconcile_queue: &NetworkPolicyReconcileQueue, state: &Arc>>, feature_flags: &pplx_space_feature_flags::FeatureFlagClient, + node_manager: &NodeManagerConnector, ) { while futures.len() < SIDECAR_RESTART_CONCURRENCY { let Some(request) = queue.pop_front() else { @@ -1446,12 +1468,13 @@ async fn fill_sidecar_restart_slots( let network_policy_reconcile_queue = network_policy_reconcile_queue.clone(); let state = Arc::clone(state); let feature_flags = feature_flags.clone(); + let node_manager = node_manager.clone(); let task = tokio_util::task::AbortOnDropHandle::new(tokio::spawn(async move { match request { SidecarRestartRequest::SpacedHost(info) => watchdog_timeout( SIDECAR_RESTART_TIMEOUT, operation, - restart_spaced_host(&info, db, sidecars), + restart_spaced_host(&info, db, sidecars, &node_manager), ) .await .map(|_| ()) @@ -1466,6 +1489,7 @@ async fn fill_sidecar_restart_slots( network_policy_reconcile_queue, &state, &feature_flags, + &node_manager, ), ) .await @@ -2107,6 +2131,10 @@ async fn probe_net_gateway_stats( } } +#[expect( + clippy::too_many_arguments, + reason = "the loop receives each long-lived watchdog dependency from the constructor" +)] async fn composed_health_loop( state: Arc>>, sidecars: Arc>>, @@ -2115,6 +2143,7 @@ async fn composed_health_loop( auto_lifecycle: AutoLifecycleCoordinator, lifecycle_epochs: HealthLifecycleEpochs, shutdown_token: CancellationToken, + node_manager: NodeManagerConnector, ) { let mut interval = tokio::time::interval(HEALTH_SCHEDULER_TICK); interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); @@ -2201,8 +2230,10 @@ async fn composed_health_loop( { diagnostics::report_transition(&evidence_queue, &transition); } + let node_manager = node_manager.clone(); pending.spawn(async move { - let result = probe_composed_health(&sandbox_id, &info, lifecycle).await; + let result = + probe_composed_health(&sandbox_id, &info, lifecycle, &node_manager).await; (sandbox_id, lifecycle, lifecycle_epoch, result) }); } @@ -2256,8 +2287,9 @@ async fn probe_composed_health( sandbox_id: &str, info: &SidecarInfo, lifecycle: Option, + node_manager: &NodeManagerConnector, ) -> ComposedProbeResult { - let firecracker = probe_firecracker_health(sandbox_id, lifecycle); + let firecracker = probe_firecracker_health(sandbox_id, lifecycle, node_manager); let spaced_host = probe_spaced_host_health(sandbox_id); let spaced_guest = probe_spaced_guest_health(info, lifecycle); let net_gateway = probe_net_gateway_health(sandbox_id); @@ -2293,6 +2325,7 @@ async fn read_guest_status_via_spaced_host( async fn probe_firecracker_health( sandbox_id: &str, lifecycle: Option, + node_manager: &NodeManagerConnector, ) -> HealthObservation { let started = Instant::now(); if !firecracker_alive(sandbox_id) { @@ -2308,7 +2341,7 @@ async fn probe_firecracker_health( return health_observation(started, HealthProbeOutcome::Unavailable, error.to_string()); } }; - match crate::firecracker_via_nodemgr::get_firecracker_vm_state(sandbox_id).await { + match crate::firecracker_via_nodemgr::get_firecracker_vm_state(sandbox_id, node_manager).await { Err(crate::firecracker_via_nodemgr::FirecrackerStateReadError::Timeout { .. }) => { health_observation( started, @@ -2814,8 +2847,9 @@ async fn restart_spaced_host( info: &SidecarInfo, db: Arc>, sidecars: Arc>>, + node_manager: &NodeManagerConnector, ) -> Result<(), SpacedHostRestartError> { - let mut node_client = pplx_space_node_mgr_client::connect().await?; + let mut node_client = node_manager.connect().await?; let id = &info.sandbox_id; let uid = info.uid.ok_or(SpacedHostRestartError::MissingUid)?; @@ -2914,6 +2948,7 @@ async fn restart_netgtw( network_policy_reconcile_queue: NetworkPolicyReconcileQueue, state: &Arc>>, feature_flags: &pplx_space_feature_flags::FeatureFlagClient, + node_manager: &NodeManagerConnector, ) -> Result<(), String> { relaunch_netgtw( info, @@ -2922,6 +2957,7 @@ async fn restart_netgtw( &network_policy_reconcile_queue, state, feature_flags, + node_manager, ) .await?; Ok(()) @@ -2943,14 +2979,13 @@ async fn relaunch_netgtw( network_policy_reconcile_queue: &NetworkPolicyReconcileQueue, state: &Arc>>, feature_flags: &pplx_space_feature_flags::FeatureFlagClient, + node_manager: &NodeManagerConnector, ) -> Result { let sandbox_id = &info.sandbox_id; let uid = info .uid .ok_or_else(|| "missing sandbox uid for net-gateway restart".to_string())?; - let mut node_client = pplx_space_node_mgr_client::connect() - .await - .map_err(|e| e.to_string())?; + let mut node_client = node_manager.connect().await.map_err(|e| e.to_string())?; let args = net_gateway_relaunch_args(db, sandbox_id, feature_flags) .map_err(|error| error.to_string())?; diff --git a/crates/libs/sbx/mgr/server/src/watchdog/tests.rs b/crates/libs/sbx/mgr/server/src/watchdog/tests.rs index 13913b979b..1f0f70b921 100644 --- a/crates/libs/sbx/mgr/server/src/watchdog/tests.rs +++ b/crates/libs/sbx/mgr/server/src/watchdog/tests.rs @@ -1103,6 +1103,7 @@ async fn pause_failed_recovery_fixture( db: db.clone(), auto_lifecycle: auto_lifecycle.clone(), health_lifecycle: lifecycle_epochs.clone(), + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), }, ))); snapshot_scheduler