From e1e992711a833d87ed68e49f588246c8df1e05e8 Mon Sep 17 00:00:00 2001 From: Abhinav Chordia Date: Mon, 12 Oct 2026 05:26:26 +0000 Subject: [PATCH 1/2] sbx-mgr: inject the node-manager client transport Every node-manager client sandbox-manager builds now comes from one NodeManagerConnector, carried in SandboxManagerServiceOptions and handed to the subsystems the service constructs: the watchdog loops, the snapshot scheduler's checkpoint context, the snapshot-upload worker and its retained-source inspector, the snapshot reaper, the restore pipeline, btrfs-via-node-manager, the encryption helpers, and the startup takeover and recovery probes in run(). NodeManagerConnector::host(), the default, forwards connect(), connect_once() and lazy() to the pplx-space-node-mgr-client functions of the same name. Production therefore dials the same fixed socket with the same retry policy, keeps the same lazy or eager timing per call site, and shares the client crate's process-wide lazy channel. Error mapping and log lines at each call site are unchanged. Two pieces of process-wide state now follow the connector: - The retained-source inspection capability answer is cached per connector. host() always returns the same instance, so production keeps its single cache. - Each immediate-reclaim backlog entry carries the connector that submitted it. The worker that drains the backlog deletes every entry through that entry's own node-manager. Tests reach an in-process node-manager instead: FakeNodeManager serves scripted unary RPCs on a socket in a test-owned temporary directory and refuses every other call with Unimplemented. - fatal_kill_signal_failure_still_schedules_admitted_cleanup moves to service/tests/node_manager_transport_tests.rs. node-manager now refuses StopProcess, and the test asserts the refusal reaches the caller, the fenced StopProcess request, and that the admitted cleanup still stops the same execution with its filesystem kept and ends Stopped. A host node-manager can no longer satisfy it. - blocking_stop_follows_a_retryable_attempt_to_durable_completion drops the cleanup executor's test dispatch hook. It drives the real CleanupSandbox path against a fake that fails once with Unavailable, and asserts both requests and the Stopped row. - The immediate-reclaim staging test and the confirmed-pause test pass a fake connector, so their node-manager fallbacks cannot reach the host. - New tests pin both pieces of per-connector state. immediate_reclaim_backlog_uses_each_entrys_own_node_manager backlogs two entries through two fakes, and send_source_inspection_is_cached_per_connector gives two fakes opposing capability answers. - grpc_core's isolated_service, which lists every option, gives node-manager an unbound connector on an absent socket under the test's scratch directory, beside its absent host-control and credential-manager sockets. Any node-manager call fails closed. --- .../sbx/mgr/server/src/btrfs_via_nodemgr.rs | 46 +++-- .../capability_cache_tests.rs | 34 ++++ .../mgr/server/src/firecracker_via_nodemgr.rs | 8 +- crates/libs/sbx/mgr/server/src/lib.rs | 1 + .../mgr/server/src/node_manager_transport.rs | 93 +++++++++ .../server/src/node_manager_transport/fake.rs | 178 ++++++++++++++++++ .../sbx/mgr/server/src/retained_source.rs | 5 +- .../sbx/mgr/server/src/rootfs_encryption.rs | 6 +- crates/libs/sbx/mgr/server/src/run.rs | 41 ++-- .../mgr/server/src/runtime_reconciliation.rs | 4 +- .../server/src/service/activation_support.rs | 3 +- .../sbx/mgr/server/src/service/auth_lookup.rs | 4 +- .../mgr/server/src/service/cache_cleanup.rs | 5 +- .../sbx/mgr/server/src/service/cleanup.rs | 12 +- .../src/service/cleanup_executor/tests.rs | 51 +++-- .../libs/sbx/mgr/server/src/service/create.rs | 30 +-- .../sbx/mgr/server/src/service/grpc_core.rs | 2 +- .../mgr/server/src/service/grpc_core/tests.rs | 7 +- .../service/grpc_lifecycle/artifact_backup.rs | 7 +- .../grpc_lifecycle/backup_submission.rs | 4 +- .../server/src/service/grpc_lifecycle/kill.rs | 4 +- .../service/grpc_lifecycle/pause_resume.rs | 12 +- .../src/service/grpc_lifecycle/suspend.rs | 3 +- .../grpc_lifecycle/suspend_snapshot.rs | 7 +- .../src/service/grpc_lifecycle/tests.rs | 2 + .../server/src/service/grpc_net_gateway.rs | 2 +- .../src/service/grpc_pool_wait_clone.rs | 26 +-- .../server/src/service/grpc_spaced_files.rs | 4 +- .../service/grpc_spaced_process_snapshots.rs | 6 +- .../server/src/service/grpc_state_snapshot.rs | 17 +- crates/libs/sbx/mgr/server/src/service/mod.rs | 22 ++- .../sbx/mgr/server/src/service/recovery.rs | 9 +- .../src/service/recovery/staging_cleanup.rs | 28 ++- .../service/recovery/staging_cleanup/tests.rs | 76 +++++++- .../server/src/service/recovery_cross_node.rs | 14 +- .../src/service/route_identity_tests.rs | 1 + .../mgr/server/src/service/spaced_bridge.rs | 9 +- .../libs/sbx/mgr/server/src/service/start.rs | 39 ++-- .../server/src/service/startup_resident.rs | 14 +- .../mgr/server/src/service/tail_support.rs | 8 +- .../libs/sbx/mgr/server/src/service/tests.rs | 37 +--- .../tests/node_manager_transport_tests.rs | 127 +++++++++++++ .../sbx/mgr/server/src/snapshot_encryption.rs | 7 +- crates/libs/sbx/mgr/server/src/snapshot_gc.rs | 23 ++- .../sbx/mgr/server/src/snapshot_gc/tests.rs | 2 + .../sbx/mgr/server/src/snapshot_restore.rs | 30 ++- .../mgr/server/src/snapshot_restore/tests.rs | 8 +- .../sbx/mgr/server/src/snapshot_scheduler.rs | 9 +- .../server/src/snapshot_scheduler/tests.rs | 1 + .../sbx/mgr/server/src/snapshot_upload.rs | 2 + .../mgr/server/src/snapshot_upload_jobs.rs | 16 +- .../server/src/snapshot_upload_jobs/tests.rs | 11 +- crates/libs/sbx/mgr/server/src/watchdog.rs | 51 ++++- .../libs/sbx/mgr/server/src/watchdog/tests.rs | 1 + 54 files changed, 942 insertions(+), 227 deletions(-) create mode 100644 crates/libs/sbx/mgr/server/src/btrfs_via_nodemgr/capability_cache_tests.rs create mode 100644 crates/libs/sbx/mgr/server/src/node_manager_transport.rs create mode 100644 crates/libs/sbx/mgr/server/src/node_manager_transport/fake.rs create mode 100644 crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs 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..5b6e8b1f53 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; 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..3ac46671d9 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 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..b960533438 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( @@ -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/grpc_core.rs b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs index e9a386455e..559bace91a 100644 --- a/crates/libs/sbx/mgr/server/src/service/grpc_core.rs +++ b/crates/libs/sbx/mgr/server/src/service/grpc_core.rs @@ -1010,7 +1010,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 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..d7808ccfa3 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`, and the host-control, credential-manager and node-manager +/// sockets are absent paths under `scratch`. fn isolated_service( conn: rusqlite::Connection, scratch: &std::path::Path, @@ -57,6 +57,9 @@ 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"), + ), }, ) } 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..ccdc17b1ee 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 @@ -640,7 +640,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 } @@ -872,18 +872,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..ce3c26f0f5 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,7 @@ pub struct SandboxManagerService { credential_manager_client: Arc>, network_credentials_enabled: bool, + node_manager: NodeManagerConnector, 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 +972,8 @@ pub struct SandboxManagerServiceOptions { pub route_ed25519_verification_bundle: Option< Arc, >, + /// Every node-manager client the service and its subsystems use. + pub node_manager: NodeManagerConnector, } impl Default for SandboxManagerServiceOptions { @@ -988,6 +992,7 @@ impl Default for SandboxManagerServiceOptions { node_role: None, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: NodeManagerConnector::host(), } } } @@ -1621,6 +1626,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 +1646,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( @@ -1715,7 +1725,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 +1739,10 @@ 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, drain_state, feature_flags, route_ed25519_verification_bundle: Arc::new(ArcSwapOption::from( @@ -1809,6 +1822,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/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..90c182efbc 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,7 +1765,7 @@ impl SandboxManagerService { .ok_or_else(|| { RestoreError::Other("restore execution identity is missing".into()) })?; - let phase2_result = restore_pull_rootfs(&staged, execution_fence) + let phase2_result = restore_pull_rootfs(&staged, execution_fence, &self.node_manager) .instrument(phase2_span) .await; if let Err(ref e) = phase2_result { 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..53c72fffd1 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,7 @@ impl RouteTestService { node_role: None, node_managed_by: None, route_ed25519_verification_bundle: None, + node_manager: crate::node_manager_transport::NodeManagerConnector::host(), }, ); 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..fd82475fdf 100644 --- a/crates/libs/sbx/mgr/server/src/service/start.rs +++ b/crates/libs/sbx/mgr/server/src/service/start.rs @@ -513,6 +513,7 @@ async fn prepare_suspend_restore_dir( execution_fence, storage_uri, transfer_scope_id, + &service.node_manager, ) .await } @@ -521,6 +522,7 @@ async fn prepare_suspend_restore_dir( &sandbox_id, execution_fence, transfer_scope_id, + &service.node_manager, ) .await } @@ -750,7 +752,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?; @@ -1780,7 +1782,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 +2121,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 +2168,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..9270f1c5ae 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::*; @@ -10951,39 +10954,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 +12401,7 @@ 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(), } } 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..46a0de93ad --- /dev/null +++ b/crates/libs/sbx/mgr/server/src/service/tests/node_manager_transport_tests.rs @@ -0,0 +1,127 @@ +//! 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`. +pub(crate) fn service_with_node_manager( + conn: rusqlite::Connection, + node_manager: NodeManagerConnector, + 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, + }, + ) +} + +#[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..422d0c7ecc 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs @@ -1256,6 +1256,7 @@ 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, ) -> Result { let latest = latest_ready_snapshot(sandbox_id).await?; let staged = restore_pull_metadata_from_source( @@ -1264,7 +1265,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).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1278,6 +1279,7 @@ 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, ) -> Result { let config = SnapshotUploadConfig::from_uri(storage_uri) .map_err(|e| RestoreError::Other(format!("parse storage_uri: {e}")))?; @@ -1290,7 +1292,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).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1849,6 +1851,7 @@ 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, ) -> Result<(), RestoreError> { use pplx_space_node_mgr_client::execution_fence::InstanceTarget; let target = stage.filesystem_scope_id.map_or( @@ -1879,7 +1882,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 +1897,7 @@ pub async fn restore_pull_rootfs( crate::snapshot_upload::artifact_kind_from_meta(&stage.meta), stage.transfer_scope_id, target, + node_manager, ) .await } @@ -1919,7 +1923,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 +1932,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 +1949,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 +2001,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 +2012,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 +2504,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; node-manager is the injected transport" +)] async fn stage_btrfs_send_restore( identity: RestoreTransferIdentity<'_>, config: &SnapshotUploadConfig, @@ -2500,6 +2516,7 @@ 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, ) -> Result<(), RestoreError> { let sandbox_id = identity.sandbox_id; let sidecar_bytes = std::fs::read(sidecar_path).map_err(|e| { @@ -2564,6 +2581,7 @@ async fn stage_btrfs_send_restore( template_name, content_hash, target, + node_manager, ) .instrument(tracing::info_span!( "sbxmgr.start.stage", 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/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 From 17387b762114314f200156e7347ee7223d5ab3bf Mon Sep 17 00:00:00 2001 From: Abhinav Chordia Date: Mon, 12 Oct 2026 05:45:37 +0000 Subject: [PATCH 2/2] sbx-mgr: inject the template-manager resolver Every template-manager GetTemplate sandbox-manager issues now goes through one TemplateResolver, carried in SandboxManagerServiceOptions as template_resolver. That covers resolve_template_paths and resolve_template_paths_for_start, which serve create, start, wake, pool adoption and donor cloning, plus restore staging's parent-template lookup and the template network-policy cache. Both resolver methods take a tonic::Request and return the tonic::Response, so request and response metadata cross the seam; callers call into_inner() as before. The default, HostTemplateManager, makes the same calls the callers made directly. get_template dials pplx_space_template_mgr_client::connect() for each call, and get_template_on_shared_channel uses the client crate's process-wide lazy() channel. Connect and RPC failures stay distinct (TemplateLookupError), so every caller keeps its error codes and messages. TemplateNetworkPolicyCache no longer has a Default constructor that bypassed the injected resolver. Tests substitute FakeTemplateResolver. It records each lookup's channel, request metadata and message, and returns a scripted answer: a response, a refusal, or a connect failure. - create_template_lookup_tests covers create's auto-lifecycle authority checks. Forwarded peers' Local authority, Unspecified and unknown authorities are refused before any lookup. Cluster authority from a forwarded peer, and Local from a local peer, send one dialed lookup for the requested template and surface its refusal as "GetTemplate failed:". An unreachable template-manager surfaces as Internal "cannot connect to template-manager:". No row is created, and node-manager, a refusing fake, receives no call. - The latest-alias start test and the orphaned-sandbox prepare test assert that no template lookup is sent. Previously they relied on an absent template-manager. - The network-policy cache tests build the cache with a refusing fake. The hint-admission and custom-override tests assert that no owner lookup is sent. - grpc_core's isolated_service, which lists every option, gets a FakeTemplateResolver that refuses every lookup, so those tests can no longer reach the fixed template-manager socket. --- crates/libs/sbx/mgr/server/src/lib.rs | 1 + .../sbx/mgr/server/src/service/auth_lookup.rs | 10 +- .../libs/sbx/mgr/server/src/service/create.rs | 2 +- .../mgr/server/src/service/create_support.rs | 60 +++--- .../src/service/create_support/tests.rs | 22 ++- .../sbx/mgr/server/src/service/grpc_core.rs | 12 +- .../grpc_core/create_template_lookup_tests.rs | 178 ++++++++++++++++++ .../mgr/server/src/service/grpc_core/tests.rs | 9 +- .../src/service/grpc_pool_wait_clone.rs | 4 +- crates/libs/sbx/mgr/server/src/service/mod.rs | 6 + .../mgr/server/src/service/network_policy.rs | 41 ++-- .../src/service/network_policy/tests.rs | 36 +++- .../server/src/service/recovery_cross_node.rs | 11 +- .../src/service/route_identity_tests.rs | 1 + .../libs/sbx/mgr/server/src/service/start.rs | 11 +- .../libs/sbx/mgr/server/src/service/tests.rs | 7 +- .../tests/node_manager_transport_tests.rs | 25 ++- .../sbx/mgr/server/src/snapshot_restore.rs | 53 +++--- .../sbx/mgr/server/src/template_resolver.rs | 67 +++++++ .../mgr/server/src/template_resolver/fake.rs | 108 +++++++++++ 20 files changed, 561 insertions(+), 103 deletions(-) create mode 100644 crates/libs/sbx/mgr/server/src/service/grpc_core/create_template_lookup_tests.rs create mode 100644 crates/libs/sbx/mgr/server/src/template_resolver.rs create mode 100644 crates/libs/sbx/mgr/server/src/template_resolver/fake.rs diff --git a/crates/libs/sbx/mgr/server/src/lib.rs b/crates/libs/sbx/mgr/server/src/lib.rs index 5b6e8b1f53..54498db7a1 100644 --- a/crates/libs/sbx/mgr/server/src/lib.rs +++ b/crates/libs/sbx/mgr/server/src/lib.rs @@ -35,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/service/auth_lookup.rs b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs index 3ac46671d9..792ae03797 100644 --- a/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs +++ b/crates/libs/sbx/mgr/server/src/service/auth_lookup.rs @@ -1005,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/create.rs b/crates/libs/sbx/mgr/server/src/service/create.rs index b960533438..e46a361d44 100644 --- a/crates/libs/sbx/mgr/server/src/service/create.rs +++ b/crates/libs/sbx/mgr/server/src/service/create.rs @@ -628,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; 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 559bace91a..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, @@ -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 d7808ccfa3..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, credential-manager and node-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, @@ -60,6 +60,11 @@ fn isolated_service( 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_pool_wait_clone.rs b/crates/libs/sbx/mgr/server/src/service/grpc_pool_wait_clone.rs index ccdc17b1ee..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(), @@ -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 diff --git a/crates/libs/sbx/mgr/server/src/service/mod.rs b/crates/libs/sbx/mgr/server/src/service/mod.rs index ce3c26f0f5..f42685dc99 100644 --- a/crates/libs/sbx/mgr/server/src/service/mod.rs +++ b/crates/libs/sbx/mgr/server/src/service/mod.rs @@ -403,6 +403,7 @@ pub struct SandboxManagerService { 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. @@ -974,6 +975,8 @@ pub struct SandboxManagerServiceOptions { >, /// 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 { @@ -993,6 +996,7 @@ impl Default for SandboxManagerServiceOptions { node_managed_by: None, route_ed25519_verification_bundle: None, node_manager: NodeManagerConnector::host(), + template_resolver: Arc::new(crate::template_resolver::HostTemplateManager), } } } @@ -1678,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)), @@ -1743,6 +1748,7 @@ impl SandboxManagerService { 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( 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_cross_node.rs b/crates/libs/sbx/mgr/server/src/service/recovery_cross_node.rs index 90c182efbc..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 @@ -1765,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, &self.node_manager) - .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 53c72fffd1..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 @@ -86,6 +86,7 @@ impl RouteTestService { 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/start.rs b/crates/libs/sbx/mgr/server/src/service/start.rs index fd82475fdf..8dc53f2d9b 100644 --- a/crates/libs/sbx/mgr/server/src/service/start.rs +++ b/crates/libs/sbx/mgr/server/src/service/start.rs @@ -514,6 +514,7 @@ async fn prepare_suspend_restore_dir( storage_uri, transfer_scope_id, &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -523,6 +524,7 @@ async fn prepare_suspend_restore_dir( execution_fence, transfer_scope_id, &service.node_manager, + service.template_resolver.as_ref(), ) .await } @@ -1363,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(), @@ -1645,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 { @@ -1661,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; diff --git a/crates/libs/sbx/mgr/server/src/service/tests.rs b/crates/libs/sbx/mgr/server/src/service/tests.rs index 9270f1c5ae..cd8073a246 100644 --- a/crates/libs/sbx/mgr/server/src/service/tests.rs +++ b/crates/libs/sbx/mgr/server/src/service/tests.rs @@ -5409,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")); @@ -12402,6 +12406,7 @@ pub(crate) mod service_tests { 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 index 46a0de93ad..879bbdfc5d 100644 --- 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 @@ -14,10 +14,12 @@ use pplx_space_node_mgr_client::proto::{ CleanupSandboxRequest, CleanupSandboxResponse, StopProcessRequest, StopProcessResponse, }; -/// A service whose node-manager clients all come from `node_manager`. -pub(crate) fn service_with_node_manager( +/// 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( @@ -35,10 +37,29 @@ pub(crate) fn service_with_node_manager( 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())); diff --git a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs index 422d0c7ecc..83252de5b6 100644 --- a/crates/libs/sbx/mgr/server/src/snapshot_restore.rs +++ b/crates/libs/sbx/mgr/server/src/snapshot_restore.rs @@ -1257,6 +1257,7 @@ pub async fn prepare_restore_dir( 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( @@ -1265,7 +1266,7 @@ pub async fn prepare_restore_dir( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence, node_manager).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1280,6 +1281,7 @@ pub async fn prepare_restore_dir_from_uri( 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}")))?; @@ -1292,7 +1294,7 @@ pub async fn prepare_restore_dir_from_uri( RestoreMetadataDestination::InPlace(transfer_scope_id), ) .await?; - restore_pull_rootfs(&staged, execution_fence, node_manager).await?; + restore_pull_rootfs(&staged, execution_fence, node_manager, templates).await?; Ok(RestorePrepared { snapshot_dir: staged.snapshot_dir, storage_uri: staged.storage_uri, @@ -1852,6 +1854,7 @@ 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( @@ -1898,6 +1901,7 @@ pub async fn restore_pull_rootfs( stage.transfer_scope_id, target, node_manager, + templates, ) .await } @@ -2506,7 +2510,7 @@ async fn stage_encrypted_restore( /// responsible for cleanup of the partial staging directory. #[expect( clippy::too_many_arguments, - reason = "each argument is an independent restore input; node-manager is the injected transport" + 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<'_>, @@ -2517,6 +2521,7 @@ async fn stage_btrfs_send_restore( 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| { @@ -2568,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) }), @@ -2654,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/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()) + } + }) + } +}