sbx-mgr: inject the node-manager client transport and template-manager resolver

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<CapturedRootfs, BtrfsViaNodemgrError> {
     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<u64, 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()))?;
     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<u64, 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
@@ -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<u64, BtrfsViaNodemgrError> {
     // 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;

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"]);
+}

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<FirecrackerVmState, FirecrackerStateReadError> {
     // 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<FirecrackerVmState, FirecrackerStateReadError> {
     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

crates/libs/sbx/mgr/server/src/lib.rs

@@ -16,6 +16,7 @@ mod guest_filesystem_flush;
 mod inventory_budget;
 mod lifecycle_metrics;
 pub mod local_artifacts;
+pub mod node_manager_transport;
 mod node_state_sync;
 pub mod restore_operation;
 pub(crate) mod retained_source;
@@ -34,6 +35,7 @@ pub mod snapshot_scheduler;
 pub mod snapshot_upload;
 pub mod snapshot_upload_jobs;
 pub mod state;
+pub mod template_resolver;
 mod transfer_directory;
 pub mod watchdog;
 

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<NodeManagerGrpcClient, ConnectError>;
+
+    /// A connected client after a single dial, for callers that own their
+    /// retry loop.
+    async fn connect_once(&self) -> Result<NodeManagerGrpcClient, ConnectError>;
+
+    /// 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<NodeManagerGrpcClient, ConnectError> {
+        pplx_space_node_mgr_client::connect().await
+    }
+
+    async fn connect_once(&self) -> Result<NodeManagerGrpcClient, ConnectError> {
+        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<dyn NodeManagerTransport>,
+    /// 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<crate::btrfs_via_nodemgr::CapabilityCache>,
+}
+
+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<NodeManagerConnector> =
+            std::sync::LazyLock::new(|| NodeManagerConnector::new(Arc::new(HostSocket)));
+        HOST.clone()
+    }
+
+    pub fn new(transport: Arc<dyn NodeManagerTransport>) -> Self {
+        Self {
+            transport,
+            inspects_send_sources: Arc::new(crate::btrfs_via_nodemgr::CapabilityCache::new()),
+        }
+    }
+
+    pub async fn connect(&self) -> Result<NodeManagerGrpcClient, ConnectError> {
+        self.transport.connect().await
+    }
+
+    pub async fn connect_once(&self) -> Result<NodeManagerGrpcClient, ConnectError> {
+        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;