From 251e4c75b204ad3e118d2c51ba776d7a4eaed7a5 Mon Sep 17 00:00:00 2001 From: Abhinav Chordia Date: Tue, 13 Oct 2026 04:43:39 +0000 Subject: [PATCH] api: ask the substrate for capabilities instead of matching on it Shared api code outside the substrate module asked "which substrate is this?" (a `Substrate::Direct`/`Substrate::Kubernetes` match, or a comparison of `managed_by()` with "k8s"/"space") where it meant "does this substrate support X?". Name those questions once. `SubstrateAdapter` gains a required `capabilities()` that returns a small `Copy` `SubstrateCapabilities`, one constant per substrate (`SubstrateCapabilities::DIRECT`, `SubstrateCapabilities::KUBERNETES`), forwarded by `impl SubstrateAdapter for Substrate`. Its fields, each read by a site converted here: - `caller_completes_claimed_restore` (Direct false, Kubernetes true): every restore-shaped wake claims `restoring` before its node call and leaves `restoring -> started` to the caller's completion CAS. The wake completion's published-suspend fence and the wake-cancellation shield rely on exactly that. Direct does not guarantee it: its worker-local wake leaves its pre-wake lifecycle (`suspended`, or `restoring` kept by an earlier lease restore) for the completion CAS to advance, and its lease restore claims `restoring` but writes `started` itself. The shield doc and two `lifecycle.rs` comments that said Direct restores never leave `suspended` are corrected. - `node_addressing` (`NodeAddressing::DirectEndpoint` / `SubstrateNodeId`): whether the fork artifact copy can reuse the source's active node target. Replaces the executor's private `CopyNodeEndpointRequirement`; it must match what the substrate's `resolve_node_api_endpoint` requires. - `registers_nodes_via_admin_api` (Direct true, Kubernetes false): the admin node registry routes (create, update, delete, template pull and delete). `require_direct_substrate` becomes `require_admin_node_registry`; the 501 response is unchanged. - `resolves_cold_create_shape` (Direct false, Kubernetes true): the control plane resolves a cold Firecracker create's shape, from a pool for a fallthrough create and otherwise from the template's build shape for values the request left unset. - `supports_node_pool_placement` (Direct false, Kubernetes true): folds in the trait method of that name, whose only caller was the admin create route. Behavior-preserving: every field has the value the replaced branch had for each substrate. Row-level checks of a row's `managed_by` stay where they are, as do `managed_by()`'s return type, the Kubernetes-only redrive hydration in `wake_completion`, and the console `/me` substrate label. Existing tests change only where a signature did: the cold-create shape test passes capabilities instead of a `managed_by` string, the fork copy target tests name `NodeAddressing`, and the two hand-written `SubstrateAdapter` test doubles implement `capabilities()`. New behavioural tests, one per field without existing coverage except `node_addressing`, each run against a real Direct substrate and (with the `kubernetes` feature) the store-backed Kubernetes test substrate, both wrapped in `Substrate`: - the wake completion fence chosen from the running substrate keeps a `suspended` row on Kubernetes and advances it to `started` on Direct; - deleting an unknown node through the admin route reaches the registry lookup (404) on Direct and is refused with the unchanged 501 on Kubernetes; - with placement enabled, a node-pool-pinned admin create is refused (400) on Direct and accepted on Kubernetes with the pin persisted in the queued create. The two admin tests run in a child test process that pins `SPACE_K8S_CONTROL_PLANE` off, so `AppState::from_pool` builds a Direct substrate whatever the host environment sets. --- .../api/src/controllers/sandbox/lifecycle.rs | 44 +++++--- .../wake_completion_postgres_tests.rs | 44 ++++++++ .../libs/api/src/controllers/sandbox/mod.rs | 4 +- .../api/src/controllers/sandbox/queries.rs | 10 +- .../sandbox/queries/create_postgres_tests.rs | 26 +++-- .../controllers/sandbox/wake_completion.rs | 28 ++--- .../api/src/controllers/substrate/direct.rs | 7 +- .../kubernetes/adapter/substrate_adapter.rs | 4 +- .../controllers/substrate/kubernetes/mod.rs | 3 +- .../libs/api/src/controllers/substrate/mod.rs | 77 +++++++++++-- .../libs/api/src/fork_operation_executor.rs | 28 ++--- .../api/src/fork_operation_executor/tests.rs | 6 +- .../api/src/routes/admin/node_maintenance.rs | 20 ++-- .../routes/admin/node_maintenance/tests.rs | 68 ++++++++++++ crates/libs/api/src/routes/admin/sandboxes.rs | 5 +- .../api/src/routes/admin/sandboxes/tests.rs | 105 ++++++++++++++++++ 16 files changed, 392 insertions(+), 87 deletions(-) diff --git a/crates/libs/api/src/controllers/sandbox/lifecycle.rs b/crates/libs/api/src/controllers/sandbox/lifecycle.rs index 31fd691209..5f8d68bd86 100644 --- a/crates/libs/api/src/controllers/sandbox/lifecycle.rs +++ b/crates/libs/api/src/controllers/sandbox/lifecycle.rs @@ -361,11 +361,13 @@ enum PublishedSuspendFence { } impl PublishedSuspendFence { - fn for_substrate(substrate: &crate::controllers::substrate::Substrate) -> Self { - match substrate { - crate::controllers::substrate::Substrate::Direct(_) => Self::DirectRestoreBaseline, - #[cfg(feature = "kubernetes")] - crate::controllers::substrate::Substrate::Kubernetes(_) => Self::Enforced, + fn for_capabilities( + capabilities: crate::controllers::substrate::SubstrateCapabilities, + ) -> Self { + if capabilities.caller_completes_claimed_restore { + Self::Enforced + } else { + Self::DirectRestoreBaseline } } } @@ -375,10 +377,13 @@ impl PublishedSuspendFence { /// A suspend claim can land after the node finishes waking but before this /// control-plane write. In that case the node-side wake succeeded, so callers /// may continue, but this stale completion must not replace `suspending` with -/// `started`. Kubernetes also fences a finalized `suspended` state. Direct -/// restores retain `suspended` as their owned baseline because that substrate -/// does not claim `restoring` before its node call. Other CAS losses retain the -/// existing error behavior. +/// `started`. A substrate whose restore-shaped wakes claim `restoring` before +/// the node call and leave `restoring -> started` to this CAS (Kubernetes) +/// also fences a finalized `suspended` state. Direct does not: its +/// worker-local wake leaves its pre-wake lifecycle (`suspended`, or +/// `restoring` kept by an earlier lease restore) for this CAS to advance, and +/// its lease restore writes `started` before returning. +/// Other CAS losses retain the existing error behavior. async fn complete_wake_lifecycle( pool: &PgPool, sandbox_id: SandboxId, @@ -1659,7 +1664,9 @@ impl SandboxController { if sandbox .lifecycle .parse::() - .is_ok_and(|lc| wake_completion::shield_applies(&self.substrate, lc)) + .is_ok_and(|lc| { + wake_completion::shield_applies(self.substrate.capabilities(), lc) + }) { completion_guard.arm(wake_completion::wake_completion_redrive( pool.clone(), @@ -1735,7 +1742,7 @@ impl SandboxController { // completion must never overwrite if a stop, delete, or // newer suspend wins while substrate work is in flight. let published_suspend_fence = - PublishedSuspendFence::for_substrate(&self.substrate); + PublishedSuspendFence::for_capabilities(self.substrate.capabilities()); let cas_result = complete_wake_lifecycle(pool, id, published_suspend_fence).await; // Explicit exit either way from here (Ok, cas-lost, or DB @@ -1826,7 +1833,9 @@ impl SandboxController { if sandbox .lifecycle .parse::() - .is_ok_and(|lc| wake_completion::shield_applies(&self.substrate, lc)) + .is_ok_and(|lc| { + wake_completion::shield_applies(self.substrate.capabilities(), lc) + }) { completion_guard.arm(wake_completion::wake_completion_redrive( pool.clone(), @@ -1875,20 +1884,21 @@ impl SandboxController { // Refuse to overwrite a terminal, teardown, or newer // suspend lifecycle. // We can't guard on `lifecycle = 'restoring'` alone - // because the two substrates enter restore from different + // because the two substrates reach this CAS from different // lifecycles: the Kubernetes substrate's // `begin_restore_lifecycle` flips // `suspended -> restoring` before the substrate call, so // the row sits in `restoring` during the slow part; the - // direct substrate's `restore_execution` doesn't touch - // lifecycle at all, so the row sits in `suspended` for the - // duration. The terminal-state exclusion is the only + // direct substrate's `restore_execution` claims + // `restoring` under its own lease and writes `started` + // itself before returning, so the row is already + // `started` here. The terminal-state exclusion is the only // predicate that catches "sweeper landed" uniformly across // substrates. On race-loss we surface a 5xx so the caller // re-issues rather than silently observing a `started` row // pointing at a deleted execution. let published_suspend_fence = - PublishedSuspendFence::for_substrate(&self.substrate); + PublishedSuspendFence::for_capabilities(self.substrate.capabilities()); let cas_result = complete_wake_lifecycle(pool, id, published_suspend_fence).await; // Explicit exit either way from here — see the Wake arm. diff --git a/crates/libs/api/src/controllers/sandbox/lifecycle/wake_completion_postgres_tests.rs b/crates/libs/api/src/controllers/sandbox/lifecycle/wake_completion_postgres_tests.rs index cb27ea8521..6155c2e2a1 100644 --- a/crates/libs/api/src/controllers/sandbox/lifecycle/wake_completion_postgres_tests.rs +++ b/crates/libs/api/src/controllers/sandbox/lifecycle/wake_completion_postgres_tests.rs @@ -126,3 +126,47 @@ async fn wake_completion_cas_outcome_for_every_prior_lifecycle() -> TestResult { } Ok(()) } + +/// The wake and restore arms pick the fence from the running substrate. On a +/// row still `suspended` at completion, Kubernetes, whose restore leaves +/// `restoring -> started` to this CAS, defers to it as a newer suspend; Direct, +/// whose worker-local wake keeps `suspended` as its baseline, advances it. +#[tokio::test(flavor = "current_thread")] +async fn wake_completion_fence_follows_the_running_substrate() -> TestResult { + use crate::controllers::substrate::{Substrate, direct::DirectSubstrate}; + + let db = pg_test_db().await?; + let pool = &db.pool; + + #[cfg_attr(not(feature = "kubernetes"), allow(unused_mut))] + let mut cases = vec![( + Substrate::Direct(Box::new(DirectSubstrate::new(pool.clone()))), + "started", + )]; + #[cfg(feature = "kubernetes")] + cases.push(( + Substrate::Kubernetes(Box::new( + crate::controllers::substrate::kubernetes::tests::substrate_with_store(pool.clone()), + )), + "suspended", + )); + + for (substrate, expected) in cases { + let sandbox_id = SandboxId::new(); + insert_sandbox(pool, sandbox_id, "suspended").await?; + + let fence = PublishedSuspendFence::for_capabilities(substrate.capabilities()); + complete_wake_lifecycle(pool, sandbox_id, fence).await?; + + let lifecycle = lifecycle_and_placement(pool, sandbox_id) + .await? + .map(|(lifecycle, _)| lifecycle); + assert_eq!( + lifecycle.as_deref(), + Some(expected), + "{} substrate", + substrate.managed_by() + ); + } + Ok(()) +} diff --git a/crates/libs/api/src/controllers/sandbox/mod.rs b/crates/libs/api/src/controllers/sandbox/mod.rs index df819456b1..1e1731bbfe 100644 --- a/crates/libs/api/src/controllers/sandbox/mod.rs +++ b/crates/libs/api/src/controllers/sandbox/mod.rs @@ -1122,8 +1122,8 @@ impl SandboxController { self.substrate.managed_by() } - pub fn supports_node_pool_placement(&self) -> bool { - self.substrate.supports_node_pool_placement() + pub fn substrate_capabilities(&self) -> super::substrate::SubstrateCapabilities { + self.substrate.capabilities() } pub async fn set_node_schedulability( diff --git a/crates/libs/api/src/controllers/sandbox/queries.rs b/crates/libs/api/src/controllers/sandbox/queries.rs index c9d5ed906c..a772b1d385 100644 --- a/crates/libs/api/src/controllers/sandbox/queries.rs +++ b/crates/libs/api/src/controllers/sandbox/queries.rs @@ -4,7 +4,7 @@ use super::*; use crate::controllers::substrate::{ AdmissionPoolClaim, AdmittedPoolClaim, AdmittedPoolClaimLookup, CreateBootMode, PooledClaimRefusal, PooledExecutionClaim, PooledExecutionClaimRequest, - PooledExecutionPoolShape, PooledPrepareFailureCleanup, + PooledExecutionPoolShape, PooledPrepareFailureCleanup, SubstrateCapabilities, }; use crate::controllers::{db_contention_error, internal_error, pool_store_error}; use crate::reconcile_queue::finalize_sandbox_db_delete_tx; @@ -566,7 +566,7 @@ struct KubernetesPoolAssignment { } struct ColdCreateShapeInput<'a> { - managed_by: &'static str, + capabilities: SubstrateCapabilities, runtime: RuntimeKind, boot_mode: CreateBootMode, template: &'a TemplateName, @@ -594,7 +594,7 @@ async fn resolve_cold_create_resources( input: ColdCreateShapeInput<'_>, ) -> Result<(SandboxResources, DiskIoLimits), ControllerError> { let ColdCreateShapeInput { - managed_by, + capabilities, runtime, boot_mode, template, @@ -603,7 +603,7 @@ async fn resolve_cold_create_resources( requested, restoring, } = input; - if managed_by != "k8s" || runtime != RuntimeKind::Firecracker { + if !capabilities.resolves_cold_create_shape || runtime != RuntimeKind::Firecracker { return Ok((requested, DiskIoLimits::default())); } if boot_mode == CreateBootMode::Fallthrough @@ -1318,7 +1318,7 @@ impl SandboxController { let (resources, disk_io_limits) = resolve_cold_create_resources( pool, ColdCreateShapeInput { - managed_by: self.substrate.managed_by(), + capabilities: self.substrate.capabilities(), runtime, boot_mode, template: &template_name, diff --git a/crates/libs/api/src/controllers/sandbox/queries/create_postgres_tests.rs b/crates/libs/api/src/controllers/sandbox/queries/create_postgres_tests.rs index 6364aa25f9..8c214556d3 100644 --- a/crates/libs/api/src/controllers/sandbox/queries/create_postgres_tests.rs +++ b/crates/libs/api/src/controllers/sandbox/queries/create_postgres_tests.rs @@ -716,6 +716,10 @@ impl SubstrateAdapter for RetryingPoolSubstrate { "k8s" } + fn capabilities(&self) -> SubstrateCapabilities { + SubstrateCapabilities::KUBERNETES + } + async fn create_execution( &self, _request: CreateExecutionRequest<'_>, @@ -1048,6 +1052,10 @@ impl SubstrateAdapter for NodeRpcStubbedKubernetes { self.inner.managed_by() } + fn capabilities(&self) -> SubstrateCapabilities { + self.inner.capabilities() + } + async fn create_execution( &self, _request: CreateExecutionRequest<'_>, @@ -1756,11 +1764,13 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap let release_id = insert_template_release(&db.pool, "app-runtime", "sha256:shape").await?; insert_inventory_shape(&db.pool, "builder", "sha256:shape", (4, 4096), 0).await?; let template = TemplateName::parse("app-runtime").unwrap(); - let resolve_with = |managed_by, boot_mode, requested, restoring| { + let k8s = SubstrateCapabilities::KUBERNETES; + let space = SubstrateCapabilities::DIRECT; + let resolve_with = |substrate, boot_mode, requested, restoring| { resolve_cold_create_resources( &db.pool, ColdCreateShapeInput { - managed_by, + capabilities: substrate, runtime: RuntimeKind::Firecracker, boot_mode, template: &template, @@ -1772,7 +1782,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap ) }; let resolve = - |managed_by, boot_mode, requested| resolve_with(managed_by, boot_mode, requested, false); + |substrate, boot_mode, requested| resolve_with(substrate, boot_mode, requested, false); let unset = || SandboxResources { vcpu: None, memory_mib: None, @@ -1781,7 +1791,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap let shape = |resources: &SandboxResources| (resources.vcpu, resources.memory_mib, resources.disk_mib); - let (resources, _) = resolve("k8s", CreateBootMode::Fallthrough, unset()).await?; + let (resources, _) = resolve(k8s, CreateBootMode::Fallthrough, unset()).await?; assert_eq!( shape(&resources), (Some(4), Some(4096), None), @@ -1789,7 +1799,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap ); let (resources, _) = resolve( - "k8s", + k8s, CreateBootMode::Cold, SandboxResources { vcpu: None, @@ -1804,14 +1814,14 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap "requested values win; only unset ones come from the template" ); - let (resources, _) = resolve("space", CreateBootMode::Fallthrough, unset()).await?; + let (resources, _) = resolve(space, CreateBootMode::Fallthrough, unset()).await?; assert_eq!( shape(&resources), (None, None, None), "direct creates keep the node's snapshot or metadata shape" ); - let (resources, _) = resolve_with("k8s", CreateBootMode::Cold, unset(), true).await?; + let (resources, _) = resolve_with(k8s, CreateBootMode::Cold, unset(), true).await?; assert_eq!( shape(&resources), (None, None, None), @@ -1829,7 +1839,7 @@ async fn postgres_kubernetes_create_fills_unset_resources_from_the_template_shap .bind(release_id) .execute(&db.pool) .await?; - let (resources, _) = resolve("k8s", CreateBootMode::Fallthrough, unset()).await?; + let (resources, _) = resolve(k8s, CreateBootMode::Fallthrough, unset()).await?; assert_eq!( shape(&resources), (Some(2), Some(8192), Some(32768)), diff --git a/crates/libs/api/src/controllers/sandbox/wake_completion.rs b/crates/libs/api/src/controllers/sandbox/wake_completion.rs index eb90aac551..2ecb1a82ac 100644 --- a/crates/libs/api/src/controllers/sandbox/wake_completion.rs +++ b/crates/libs/api/src/controllers/sandbox/wake_completion.rs @@ -41,7 +41,7 @@ use pplx_space_sandbox_store as sandbox_store; use sqlx::PgPool; use tokio::sync::Semaphore; -use crate::controllers::substrate::Substrate; +use crate::controllers::substrate::{Substrate, SubstrateCapabilities}; /// How often the redrive re-checks the row and re-attempts hydration. The /// node-side restore this waits on takes seconds to low tens of seconds; @@ -78,25 +78,27 @@ pub(super) fn redrive_budget() -> Duration { } /// Whether the cancellation shield applies to this wake: only restore-shaped -/// wakes (`suspended`/`restoring` pre-wake lifecycle) on the Kubernetes -/// substrate. Everything else keeps existing semantics: -/// * the direct substrate's restore never flips the row out of `suspended` -/// mid-flight, so a cancelled direct restore leaves an already-recoverable -/// state; +/// wakes (`suspended`/`restoring` pre-wake lifecycle) on a substrate that +/// leaves `restoring -> started` to the caller's completion CAS (Kubernetes). +/// Everything else keeps existing semantics: +/// * the direct substrate does not guarantee this: its worker-local +/// wake leaves the pre-wake lifecycle (`suspended`, or `restoring` kept by +/// an earlier lease restore) for the completion CAS, and its lease restore +/// writes `started` itself, so a cancelled direct wake leaves `suspended`, +/// `started`, or a `restoring` row that a later restore reclaims once any +/// lease expires; /// * pause/resume and stop/start wakes don't take the restore claim and have /// no execution to hydrate — their (much rarer) cancellation class is a /// separate piece of work. -pub(super) fn shield_applies(substrate: &Substrate, pre_wake_lifecycle: SandboxLifecycle) -> bool { +pub(super) fn shield_applies( + capabilities: SubstrateCapabilities, + pre_wake_lifecycle: SandboxLifecycle, +) -> bool { let restore_shaped = matches!( pre_wake_lifecycle, SandboxLifecycle::Suspended | SandboxLifecycle::Restoring ); - let kubernetes = match substrate { - Substrate::Direct(_) => false, - #[cfg(feature = "kubernetes")] - Substrate::Kubernetes(_) => true, - }; - restore_shaped && kubernetes + restore_shaped && capabilities.caller_completes_claimed_restore } /// Drop-guard that owns the boxed redrive future while the wake handler is diff --git a/crates/libs/api/src/controllers/substrate/direct.rs b/crates/libs/api/src/controllers/substrate/direct.rs index 9d1c7af18b..cd7aeba25a 100644 --- a/crates/libs/api/src/controllers/substrate/direct.rs +++ b/crates/libs/api/src/controllers/substrate/direct.rs @@ -27,7 +27,8 @@ use super::{ ExecutionBinding, NodeApiTarget, NodeInventoryEntry, OwnedNodeApiTarget, PooledExecutionClaim, PooledExecutionClaimAttempt, PooledExecutionClaimRequest, PooledExecutionPoolShape, RetireExecutionOutcome, RetireExecutionRequest, RuntimeClient, SubstrateAdapter, - SubstrateError, connect_builder, connect_worker, node_api_endpoint, substrate_db_error, + SubstrateCapabilities, SubstrateError, connect_builder, connect_worker, node_api_endpoint, + substrate_db_error, }; use crate::controllers::auto_lifecycle::policy_to_proto; use crate::controllers::{ControllerError, db_contention_error}; @@ -3087,6 +3088,10 @@ impl SubstrateAdapter for DirectSubstrate { "space" } + fn capabilities(&self) -> SubstrateCapabilities { + SubstrateCapabilities::DIRECT + } + fn pool_eligible_for_request( &self, boot_mode: CreateBootMode, diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/adapter/substrate_adapter.rs b/crates/libs/api/src/controllers/substrate/kubernetes/adapter/substrate_adapter.rs index 8b328c47be..ef036d4ad7 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/adapter/substrate_adapter.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/adapter/substrate_adapter.rs @@ -7,8 +7,8 @@ impl SubstrateAdapter for KubernetesSubstrate { "k8s" } - fn supports_node_pool_placement(&self) -> bool { - true + fn capabilities(&self) -> SubstrateCapabilities { + SubstrateCapabilities::KUBERNETES } async fn create_execution<'a>( diff --git a/crates/libs/api/src/controllers/substrate/kubernetes/mod.rs b/crates/libs/api/src/controllers/substrate/kubernetes/mod.rs index 1a6d84fe7c..825e08d21a 100644 --- a/crates/libs/api/src/controllers/substrate/kubernetes/mod.rs +++ b/crates/libs/api/src/controllers/substrate/kubernetes/mod.rs @@ -83,7 +83,8 @@ use super::{ ExecutionIdentity, NODE_RPC_TIMEOUT, NodeApiTarget, NodeInventoryEntry, ObservedPod, ObservedPodIdentity, OwnedNodeApiTarget, PodFenceMismatch, PodIdentityFence, RecordedExecution, RetireExecutionOutcome, RetireExecutionRequest, RetireFence, RuntimeClient, SubstrateAdapter, - SubstrateError, TeardownProgress, WaitReadyTimeoutPhase, connect_worker, node_api_endpoint, + SubstrateCapabilities, SubstrateError, TeardownProgress, WaitReadyTimeoutPhase, connect_worker, + node_api_endpoint, }; const START_WAIT_FAST_POLL_INTERVAL: Duration = Duration::from_millis(250); diff --git a/crates/libs/api/src/controllers/substrate/mod.rs b/crates/libs/api/src/controllers/substrate/mod.rs index d711259ccd..b3cdeb4132 100644 --- a/crates/libs/api/src/controllers/substrate/mod.rs +++ b/crates/libs/api/src/controllers/substrate/mod.rs @@ -2132,6 +2132,70 @@ impl RecoveryCharge { } } +/// How the substrate addresses a node for node-api calls: which field of a +/// [`NodeApiTarget`] must be present. Must match the field the substrate's +/// [`SubstrateAdapter::resolve_node_api_endpoint`] requires. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NodeAddressing { + /// The node row's stored node-api endpoint (`direct_endpoint`). + DirectEndpoint, + /// An opaque substrate node identifier (`substrate_node_id`) that the + /// substrate resolves to an endpoint. + SubstrateNodeId, +} + +/// Static facts about what a substrate supports. Shared code reads these +/// instead of matching on which substrate is active. Each substrate returns +/// one of the constants below; behavior that needs the substrate to act stays +/// a [`SubstrateAdapter`] method. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SubstrateCapabilities { + /// Every restore-shaped wake (from `suspended` or `restoring`) claims + /// `restoring` before its node call and leaves `restoring -> started` to + /// the caller's completion CAS. The caller can then treat a `suspended` + /// row at completion as a newer, finalized suspend and fence it, and must + /// shield that CAS against request cancellation. + /// + /// `false` does not mean the substrate's rows are never `restoring`. + /// Direct's worker-local wake leaves its pre-wake lifecycle (`suspended`, + /// or `restoring` kept by an earlier lease restore) for the completion CAS + /// to advance; its lease restore claims `restoring` and writes `started` + /// itself. + pub caller_completes_claimed_restore: bool, + /// How node-api targets are addressed. + pub node_addressing: NodeAddressing, + /// Nodes live in the admin API's node registry: they are registered, + /// updated and deleted there, and their templates are pulled and deleted + /// there, rather than discovered by the substrate. + pub registers_nodes_via_admin_api: bool, + /// The control plane resolves a cold Firecracker create's shape: a + /// fallthrough create takes a configured or policy pool shape, and + /// otherwise vCPU and memory the request left unset come from the + /// template's build shape. Without it, the requested shape goes to the + /// node unchanged. + pub resolves_cold_create_shape: bool, + /// Create requests may carry a Kubernetes node-pool placement pin. + pub supports_node_pool_placement: bool, +} + +impl SubstrateCapabilities { + pub const DIRECT: Self = Self { + caller_completes_claimed_restore: false, + node_addressing: NodeAddressing::DirectEndpoint, + registers_nodes_via_admin_api: true, + resolves_cold_create_shape: false, + supports_node_pool_placement: false, + }; + + pub const KUBERNETES: Self = Self { + caller_completes_claimed_restore: true, + node_addressing: NodeAddressing::SubstrateNodeId, + registers_nodes_via_admin_api: false, + resolves_cold_create_shape: true, + supports_node_pool_placement: true, + }; +} + #[allow(async_fn_in_trait)] pub trait SubstrateAdapter: Send + Sync { /// Stable label written to control-plane metadata for this substrate. @@ -2140,10 +2204,9 @@ pub trait SubstrateAdapter: Send + Sync { /// execution binding, not this label. fn managed_by(&self) -> &'static str; - /// Whether create requests may carry a Kubernetes node-pool placement pin. - fn supports_node_pool_placement(&self) -> bool { - false - } + /// What this substrate supports. Required, so every substrate (and the + /// `Substrate` forwarding impl) states its answer. + fn capabilities(&self) -> SubstrateCapabilities; /// Create the active execution object for a newly-created sandbox. /// @@ -3235,11 +3298,11 @@ impl SubstrateAdapter for Substrate { } } - fn supports_node_pool_placement(&self) -> bool { + fn capabilities(&self) -> SubstrateCapabilities { match self { - Self::Direct(substrate) => substrate.supports_node_pool_placement(), + Self::Direct(substrate) => substrate.capabilities(), #[cfg(feature = "kubernetes")] - Self::Kubernetes(substrate) => substrate.supports_node_pool_placement(), + Self::Kubernetes(substrate) => substrate.capabilities(), } } diff --git a/crates/libs/api/src/fork_operation_executor.rs b/crates/libs/api/src/fork_operation_executor.rs index 09eed77948..fea4b66a61 100644 --- a/crates/libs/api/src/fork_operation_executor.rs +++ b/crates/libs/api/src/fork_operation_executor.rs @@ -18,7 +18,9 @@ use pplx_space_public_api::{ }; use sqlx::PgPool; -use crate::controllers::substrate::{OwnedNodeApiTarget, SnapshotRef, Substrate, SubstrateAdapter}; +use crate::controllers::substrate::{ + NodeAddressing, OwnedNodeApiTarget, SnapshotRef, Substrate, SubstrateAdapter, +}; use crate::controllers::{ ControllerError, sandbox::{ @@ -44,28 +46,14 @@ const MAX_BACKUP_ATTEMPTS: u8 = 3; const CLAIM_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(10); const EXTERNAL_WORK_DEADLINE: Duration = Duration::from_secs(600); -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum CopyNodeEndpointRequirement { - DirectEndpoint, - SubstrateNodeId, -} - struct SuccessfulForkTargetRecord { target_id: SandboxId, result: ForkTargetOutcome, snapshot_id: Option, } -fn copy_node_endpoint_requirement(substrate: &Substrate) -> CopyNodeEndpointRequirement { - match substrate { - Substrate::Direct(_) => CopyNodeEndpointRequirement::DirectEndpoint, - #[cfg(feature = "kubernetes")] - Substrate::Kubernetes(_) => CopyNodeEndpointRequirement::SubstrateNodeId, - } -} - fn active_copy_node_target_is_complete( - endpoint_requirement: CopyNodeEndpointRequirement, + node_addressing: NodeAddressing, managed_by_present: bool, role_present: bool, direct_endpoint_present: bool, @@ -73,9 +61,9 @@ fn active_copy_node_target_is_complete( ) -> bool { managed_by_present && role_present - && match endpoint_requirement { - CopyNodeEndpointRequirement::DirectEndpoint => direct_endpoint_present, - CopyNodeEndpointRequirement::SubstrateNodeId => substrate_node_id_present, + && match node_addressing { + NodeAddressing::DirectEndpoint => direct_endpoint_present, + NodeAddressing::SubstrateNodeId => substrate_node_id_present, } } @@ -800,7 +788,7 @@ impl ForkSandboxOperationExecutor { .map_err(ForkExecutionError::retry_error)? .ok_or_else(|| ForkExecutionError::not_found("source sandbox disappeared"))?; let active_target_is_complete = active_copy_node_target_is_complete( - copy_node_endpoint_requirement(&self.substrate), + self.substrate.capabilities().node_addressing, source.active_node_managed_by.is_some(), source.active_node_role.is_some(), source.active_node_address.is_some(), diff --git a/crates/libs/api/src/fork_operation_executor/tests.rs b/crates/libs/api/src/fork_operation_executor/tests.rs index 0c37973211..5fe753e2bb 100644 --- a/crates/libs/api/src/fork_operation_executor/tests.rs +++ b/crates/libs/api/src/fork_operation_executor/tests.rs @@ -649,14 +649,14 @@ mod cold_copy_tests { #[test] fn kubernetes_copy_target_requires_a_substrate_node_id() { assert!(!active_copy_node_target_is_complete( - CopyNodeEndpointRequirement::SubstrateNodeId, + NodeAddressing::SubstrateNodeId, true, true, true, false, )); assert!(active_copy_node_target_is_complete( - CopyNodeEndpointRequirement::SubstrateNodeId, + NodeAddressing::SubstrateNodeId, true, true, true, @@ -667,7 +667,7 @@ fn kubernetes_copy_target_requires_a_substrate_node_id() { #[test] fn direct_copy_target_accepts_a_direct_endpoint() { assert!(active_copy_node_target_is_complete( - CopyNodeEndpointRequirement::DirectEndpoint, + NodeAddressing::DirectEndpoint, true, true, true, diff --git a/crates/libs/api/src/routes/admin/node_maintenance.rs b/crates/libs/api/src/routes/admin/node_maintenance.rs index d2fd531849..a713aed433 100644 --- a/crates/libs/api/src/routes/admin/node_maintenance.rs +++ b/crates/libs/api/src/routes/admin/node_maintenance.rs @@ -428,8 +428,14 @@ fn ensure_template_capable_node( Ok(()) } -fn require_direct_substrate(state: &AppState) -> Result<(), (StatusCode, Json)> { - if state.sandbox_controller.managed_by() != "space" { +fn require_admin_node_registry( + state: &AppState, +) -> Result<(), (StatusCode, Json)> { + if !state + .sandbox_controller + .substrate_capabilities() + .registers_nodes_via_admin_api + { return Err(api_error( StatusCode::NOT_IMPLEMENTED, "direct node registry is only available for the direct/self-managed substrate", @@ -995,7 +1001,7 @@ pub async fn create_node( if let Err(err) = require_managing_org_admin(&state.pool, &actor).await { return err.into_response(); } - if let Err(err) = require_direct_substrate(&state) { + if let Err(err) = require_admin_node_registry(&state) { return err.into_response(); } let role = req.role; @@ -1189,7 +1195,7 @@ pub async fn update_node( if let Err(err) = require_managing_org_admin(&state.pool, &actor).await { return err.into_response(); } - if let Err(err) = require_direct_substrate(&state) { + if let Err(err) = require_admin_node_registry(&state) { return err.into_response(); } if let Some(metadata) = req.metadata.as_ref() @@ -1326,7 +1332,7 @@ pub async fn delete_node( if let Err(err) = require_managing_org_admin(&state.pool, &actor).await { return err.into_response(); } - if let Err(err) = require_direct_substrate(&state) { + if let Err(err) = require_admin_node_registry(&state) { return err.into_response(); } @@ -1578,7 +1584,7 @@ pub async fn pull_node_template( if let Err(err) = require_managing_org_admin(&state.pool, &actor).await { return err.into_response(); } - if let Err(err) = require_direct_substrate(&state) { + if let Err(err) = require_admin_node_registry(&state) { return err.into_response(); } @@ -1709,7 +1715,7 @@ pub async fn delete_node_template( if let Err(err) = require_managing_org_admin(&state.pool, &actor).await { return err.into_response(); } - if let Err(err) = require_direct_substrate(&state) { + if let Err(err) = require_admin_node_registry(&state) { return err.into_response(); } diff --git a/crates/libs/api/src/routes/admin/node_maintenance/tests.rs b/crates/libs/api/src/routes/admin/node_maintenance/tests.rs index dc321dfb65..43e5b077e9 100644 --- a/crates/libs/api/src/routes/admin/node_maintenance/tests.rs +++ b/crates/libs/api/src/routes/admin/node_maintenance/tests.rs @@ -124,3 +124,71 @@ async fn hard_delete_direct_node_removes_node_and_projection_rows() -> TestResul assert_eq!(deleted_again, 0); Ok(()) } + +/// The node registry routes are open only where the running substrate keeps +/// its nodes in the admin registry. Deleting a node that does not exist +/// reaches the registry lookup on Direct and is refused before it on +/// Kubernetes. +#[tokio::test(flavor = "current_thread")] +async fn node_registry_delete_follows_the_running_substrate() -> TestResult { + // `AppState::from_pool` chooses its substrate from the process + // environment. Keep it Direct whatever the host sets, so construction + // never reaches a real cluster before the Kubernetes fixture is swapped in. + if !pplx_space_test_env::enter_test_process(&[( + "SPACE_K8S_CONTROL_PLANE", + Some(std::ffi::OsStr::new("false")), + )]) { + return Ok(()); + } + let Some(db) = try_pg_test_db().await? else { + return Ok(()); + }; + let actor = sqlx::query_as::<_, crate::models::User>(sql!( + "SELECT u.* + FROM users u + JOIN clusters c ON c.managing_org_id = u.org_id + WHERE c.singleton + AND u.role = 'organization_admin' + ORDER BY u.id + LIMIT 1" + )) + .fetch_one(&db.pool) + .await?; + let direct = AppState::from_pool(db.pool.clone()).await?; + assert_eq!(direct.sandbox_controller.managed_by(), "space"); + + #[cfg_attr(not(feature = "kubernetes"), allow(unused_mut))] + let mut cases = vec![(direct, StatusCode::NOT_FOUND, "Node not found")]; + #[cfg(feature = "kubernetes")] + { + let mut kubernetes = AppState::from_pool(db.pool.clone()).await?; + kubernetes.sandbox_controller = crate::controllers::sandbox::SandboxController::new( + crate::controllers::substrate::Substrate::Kubernetes(Box::new( + crate::controllers::substrate::kubernetes::tests::substrate_with_store( + db.pool.clone(), + ), + )), + ); + cases.push(( + kubernetes, + StatusCode::NOT_IMPLEMENTED, + "direct node registry is only available for the direct/self-managed substrate", + )); + } + + for (state, expected_status, expected_error) in cases { + let substrate = state.sandbox_controller.managed_by(); + let response = delete_node( + State(state.clone()), + Extension(AuthUser(actor.clone())), + Path(NodeId::new()), + ) + .await + .into_response(); + assert_eq!(response.status(), expected_status, "{substrate} substrate"); + let body: serde_json::Value = + serde_json::from_slice(&axum::body::to_bytes(response.into_body(), 100_000).await?)?; + assert_eq!(body["error"], expected_error, "{substrate} substrate"); + } + Ok(()) +} diff --git a/crates/libs/api/src/routes/admin/sandboxes.rs b/crates/libs/api/src/routes/admin/sandboxes.rs index 352999bb15..ad00bd9bc0 100644 --- a/crates/libs/api/src/routes/admin/sandboxes.rs +++ b/crates/libs/api/src/routes/admin/sandboxes.rs @@ -242,7 +242,10 @@ pub async fn create_sandbox( &state.pool, request, &idempotency_key, - state.sandbox_controller.supports_node_pool_placement(), + state + .sandbox_controller + .substrate_capabilities() + .supports_node_pool_placement, ) .await { diff --git a/crates/libs/api/src/routes/admin/sandboxes/tests.rs b/crates/libs/api/src/routes/admin/sandboxes/tests.rs index a7ce651f21..01a1546110 100644 --- a/crates/libs/api/src/routes/admin/sandboxes/tests.rs +++ b/crates/libs/api/src/routes/admin/sandboxes/tests.rs @@ -532,3 +532,108 @@ fn admin_sandbox_reads_keep_the_controller_contention_disposition() { "the non-contention path must still run the call site's log; it is the only record" ); } + +/// With placement enabled, an admin create pinned to a node pool is refused +/// on a substrate without node-pool placement (Direct) and accepted, with the +/// pin persisted in the queued create, on one that has it (Kubernetes). +#[tokio::test] +async fn admin_pinned_create_follows_the_running_substrate() -> TestResult { + // Turn the placement gate on, and keep the substrate `AppState::from_pool` + // chooses from the process environment Direct whatever the host sets. + if !pplx_space_test_env::enter_test_process(&[ + ( + pplx_space_common_core::env::SANDBOX_ALLOW_PLACEMENT_ENV, + Some(std::ffi::OsStr::new("true")), + ), + ( + "SPACE_K8S_CONTROL_PLANE", + Some(std::ffi::OsStr::new("false")), + ), + ]) { + return Ok(()); + } + let Some(db) = try_pg_test_db().await? else { + return Ok(()); + }; + let actor = sqlx::query_as::<_, crate::models::User>(sql!( + "SELECT u.* + FROM users u + JOIN clusters c ON c.managing_org_id = u.org_id + WHERE c.singleton + AND u.role = 'organization_admin' + ORDER BY u.id + LIMIT 1" + )) + .fetch_one(&db.pool) + .await?; + let direct = AppState::from_pool(db.pool.clone()).await?; + assert_eq!(direct.sandbox_controller.managed_by(), "space"); + + #[cfg_attr(not(feature = "kubernetes"), allow(unused_mut))] + let mut cases = vec![(direct, StatusCode::BAD_REQUEST)]; + #[cfg(feature = "kubernetes")] + { + let mut kubernetes = AppState::from_pool(db.pool.clone()).await?; + kubernetes.sandbox_controller = crate::controllers::sandbox::SandboxController::new( + crate::controllers::substrate::Substrate::Kubernetes(Box::new( + crate::controllers::substrate::kubernetes::tests::substrate_with_store( + db.pool.clone(), + ), + )), + ); + cases.push((kubernetes, StatusCode::ACCEPTED)); + } + + for (state, expected_status) in cases { + let substrate = state.sandbox_controller.managed_by(); + let org_id = seed_org(&db.pool).await?; + seed_user(&db.pool, OrgId::from_uuid(org_id)).await?; + let request = serde_json::from_value(serde_json::json!({ + "org_id": org_id, + "sandbox": { + "name": null, + "template": "seed", + "runtime": "fc", + "env": {}, + "metadata": {}, + "resources": null + }, + "node_pool": "isolated-test-pool" + }))?; + let mut headers = HeaderMap::new(); + headers.insert("idempotency-key", Uuid::now_v7().to_string().parse()?); + + let response = create_sandbox( + State(state.clone()), + Extension(AuthUser(actor.clone())), + headers, + AdminCreateJson(request), + ) + .await; + + assert_eq!(response.status(), expected_status, "{substrate} substrate"); + let persisted_pins: Vec> = sqlx::query_scalar(sql!( + "SELECT payload_json->>'node_pool' FROM public_operations WHERE org_id = $1" + )) + .bind(org_id) + .fetch_all(&db.pool) + .await?; + if expected_status == StatusCode::ACCEPTED { + assert_eq!( + persisted_pins, + [Some("isolated-test-pool".to_owned())], + "{substrate} substrate" + ); + } else { + let body: serde_json::Value = serde_json::from_slice( + &axum::body::to_bytes(response.into_body(), 100_000).await?, + )?; + assert_eq!( + body["error"], "sandbox node-pool placement requires the Kubernetes substrate", + "{substrate} substrate" + ); + assert!(persisted_pins.is_empty(), "{substrate} substrate"); + } + } + Ok(()) +}