From 789bf1fe191f89b3ec65481acb5bdf8829d259aa Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Sun, 28 Jun 2026 11:57:56 +0400 Subject: [PATCH] feat(mvp-system): implement node and docker cluster provisioning node_provisioning: NodeManager record FSM and transient BootstrapSession SSH core per the provisioning spec. docker_cluster_provisioning provider; add guarantee tests for both. Signed-off-by: Zachery Aaron Shores-Chmielewski --- .../src/docker_cluster_provisioning.rs | 689 +++++++++++++ crates/mvp-system/src/lib.rs | 2 + crates/mvp-system/src/node_provisioning.rs | 964 ++++++++++++++++++ .../docker_cluster_provisioning_guarantees.rs | 466 +++++++++ crates/mvp-system/src/tests/mod.rs | 2 + .../src/tests/node_provisioning_guarantees.rs | 533 ++++++++++ 6 files changed, 2656 insertions(+) create mode 100644 crates/mvp-system/src/docker_cluster_provisioning.rs create mode 100644 crates/mvp-system/src/node_provisioning.rs create mode 100644 crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs create mode 100644 crates/mvp-system/src/tests/node_provisioning_guarantees.rs diff --git a/crates/mvp-system/src/docker_cluster_provisioning.rs b/crates/mvp-system/src/docker_cluster_provisioning.rs new file mode 100644 index 0000000..bac478e --- /dev/null +++ b/crates/mvp-system/src/docker_cluster_provisioning.rs @@ -0,0 +1,689 @@ +//! Docker-backed provider and bootstrap wiring for node provisioning tests. +//! +//! Docker is treated as a concrete provider adapter here: it creates and destroys +//! real Docker container leases through a `DockerCli` boundary. Unit tests use a +//! deterministic CLI implementation, but the provider behavior remains the same +//! provider contract as a remote adapter. + +use std::collections::{BTreeMap, BTreeSet}; +use std::time::SystemTime; + +use crate::node_provisioning::{ + BootstrapDatastreamSink, BootstrapLogRecord, BootstrapLogSource, BootstrapLogStream, + BootstrapObservation, BootstrapSessionEvent, BootstrapSessionSpec, BootstrapStage, + CreateLeaseRequest, CreateLeaseResult, DesiredNodeShape, DestroyHandle, LeaseFacts, + LogicalNodeId, NodeManager, NodeManagerCommand, NodeManagerMsg, NodeRecord, ProviderError, + ProviderKind, ProviderLeaseId, ProviderPlugin, RunId, RunNodeGroupSpec, SshEndpoint, SwactorId, + SwarmJoinSpec, expand_node_group, +}; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DockerCliError { + pub reason: String, +} + +impl DockerCliError { + pub fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DockerRunRequest { + pub container_name: String, + pub image: String, + pub run_id: RunId, + pub logical_node_id: LogicalNodeId, + pub ssh_user: String, + pub labels: BTreeMap, + pub env: BTreeMap, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DockerRunResult { + pub container_id: String, +} + +pub trait DockerCli { + fn run_container( + &mut self, + request: DockerRunRequest, + ) -> Result; + + fn inspect_ssh_endpoint( + &mut self, + container_id: &str, + ) -> Result, DockerCliError>; + + fn remove_force(&mut self, container_id: &str) -> Result<(), DockerCliError>; +} + +#[derive(Clone, Debug)] +pub struct DockerProvider { + cli: C, +} + +impl DockerProvider { + pub fn new(cli: C) -> Self { + Self { cli } + } + + pub fn cli(&self) -> &C { + &self.cli + } + + pub fn cli_mut(&mut self) -> &mut C { + &mut self.cli + } + + pub fn into_cli(self) -> C { + self.cli + } +} + +impl DockerProvider { + fn build_run_request(request: &CreateLeaseRequest) -> DockerRunRequest { + let spec = &request.spec; + let mut labels = spec.shape.provider_labels.clone(); + labels.insert("mvp.provider".into(), "docker".into()); + labels.insert("mvp.run_id".into(), spec.run_id.0.to_string()); + labels.insert("mvp.logical_node_id".into(), spec.logical_node_id.0.clone()); + + let mut env = BTreeMap::new(); + env.insert("MVP_RUN_ID".into(), spec.run_id.0.to_string()); + env.insert("MVP_LOGICAL_NODE_ID".into(), spec.logical_node_id.0.clone()); + env.insert( + "MVP_ORCH_SWACTOR_ADDR".into(), + spec.swarm_join.orch_swactor_addr.clone(), + ); + env.insert( + "MVP_JOIN_TOKEN_REF".into(), + spec.swarm_join.join_token_ref.clone(), + ); + + DockerRunRequest { + container_name: format!("mvp-{}-{}", spec.run_id.0, spec.logical_node_id.0), + image: spec.shape.image.clone(), + run_id: spec.run_id.clone(), + logical_node_id: spec.logical_node_id.clone(), + ssh_user: spec.boot.ssh_user.clone(), + labels, + env, + } + } + + fn lease_from_container( + shape: &DesiredNodeShape, + logical_node_id: &LogicalNodeId, + container_id: String, + endpoint: &Option, + ) -> LeaseFacts { + let lease_id = ProviderLeaseId(format!("docker:{container_id}")); + let mut provider_metadata = BTreeMap::new(); + provider_metadata.insert("container_id".into(), container_id.clone()); + provider_metadata.insert("image".into(), shape.image.clone()); + provider_metadata.insert("logical_node_id".into(), logical_node_id.0.clone()); + if let Some(endpoint) = endpoint { + provider_metadata.insert("ssh_host".into(), endpoint.host.clone()); + provider_metadata.insert("ssh_port".into(), endpoint.port.to_string()); + } + + LeaseFacts { + provider: ProviderKind::Docker, + lease_id: lease_id.clone(), + provider_contract_id: container_id.clone(), + offer_id: None, + destroy_handle: DestroyHandle { + provider: ProviderKind::Docker, + lease_id, + provider_contract_id: container_id, + }, + provider_metadata, + } + } +} + +impl ProviderPlugin for DockerProvider { + fn create_lease( + &mut self, + request: CreateLeaseRequest, + ) -> Result { + if request.spec.provider != ProviderKind::Docker { + return Err(ProviderError::new( + "docker provider received non-docker node spec", + )); + } + let run_request = Self::build_run_request(&request); + let run_result = self + .cli + .run_container(run_request) + .map_err(|error| ProviderError::new(error.reason))?; + let endpoint = self + .cli + .inspect_ssh_endpoint(&run_result.container_id) + .map_err(|error| ProviderError::new(error.reason))?; + let lease = Self::lease_from_container( + &request.spec.shape, + &request.spec.logical_node_id, + run_result.container_id, + &endpoint, + ); + Ok(CreateLeaseResult { lease, endpoint }) + } + + fn lookup_endpoint( + &mut self, + lease: &LeaseFacts, + ) -> Result, ProviderError> { + self.cli + .inspect_ssh_endpoint(&lease.provider_contract_id) + .map_err(|error| ProviderError::new(error.reason)) + } + + fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError> { + if handle.provider != ProviderKind::Docker { + return Err(ProviderError::new( + "docker provider received non-docker destroy handle", + )); + } + self.cli + .remove_force(&handle.provider_contract_id) + .map_err(|error| ProviderError::new(error.reason)) + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct BootstrapSshError { + pub reason: String, +} + +impl BootstrapSshError { + pub fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +pub trait BootstrapSshClient { + fn connect(&mut self, endpoint: &SshEndpoint) -> Result<(), BootstrapSshError>; + fn probe_stdout(&mut self) -> Result<(), BootstrapSshError>; + fn read_bootstrap_logs( + &mut self, + stdout_sources: &[String], + stderr_sources: &[String], + ) -> Result, BootstrapSshError>; + fn run_verify_commands(&mut self, commands: &[String]) -> Result<(), BootstrapSshError>; + fn start_swactor( + &mut self, + command: &str, + join: &SwarmJoinSpec, + ) -> Result<(), BootstrapSshError>; + fn close(&mut self); +} + +#[derive(Clone, Debug)] +pub struct SshBootstrapSession { + spec: BootstrapSessionSpec, + client: C, + stage: BootstrapStage, + next_seq: u64, + last_stdout_seq: Option, + last_stderr_seq: Option, + closed: bool, +} + +impl SshBootstrapSession { + pub fn new(spec: BootstrapSessionSpec, client: C) -> Self { + Self { + spec, + client, + stage: BootstrapStage::Created, + next_seq: 1, + last_stdout_seq: None, + last_stderr_seq: None, + closed: false, + } + } + + pub fn client(&self) -> &C { + &self.client + } + + pub fn client_mut(&mut self) -> &mut C { + &mut self.client + } + + pub fn stage(&self) -> BootstrapStage { + self.stage + } + + pub fn is_closed(&self) -> bool { + self.closed + } +} + +impl SshBootstrapSession { + pub fn start(&mut self, sink: &mut dyn BootstrapDatastreamSink) -> Vec { + let mut events = Vec::new(); + self.stage = BootstrapStage::SshConnecting; + if let Err(error) = self.client.connect(&self.spec.ssh) { + self.stage = BootstrapStage::SshTimeout; + return vec![BootstrapSessionEvent::Failed(format!( + "ssh connect: {}", + error.reason + ))]; + } + if let Err(error) = self.client.probe_stdout() { + self.stage = BootstrapStage::SshTimeout; + return vec![BootstrapSessionEvent::Failed(format!( + "ssh probe: {}", + error.reason + ))]; + } + self.stage = BootstrapStage::SshReady; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::SshReady), + )); + + self.stage = BootstrapStage::StdoutStreaming; + let records = match self.client.read_bootstrap_logs( + &self.spec.boot.stdout_sources, + &self.spec.boot.stderr_sources, + ) { + Ok(records) => records, + Err(error) => { + self.stage = BootstrapStage::StreamError; + events.push(BootstrapSessionEvent::Failed(format!( + "bootstrap log stream: {}", + error.reason + ))); + return events; + } + }; + for (stream, line) in records { + let seq = self.next_seq; + self.next_seq += 1; + match stream { + BootstrapLogStream::Stdout => self.last_stdout_seq = Some(seq), + BootstrapLogStream::Stderr => self.last_stderr_seq = Some(seq), + } + sink.record(BootstrapLogRecord { + run_id: self.spec.run_id.clone(), + logical_node_id: self.spec.logical_node_id.clone(), + lease_id: self.spec.lease_id.clone(), + source: BootstrapLogSource::SshBootstrap, + stream, + seq, + timestamp: SystemTime::now(), + line, + }); + } + events.push(BootstrapSessionEvent::Observed(BootstrapObservation { + stage: BootstrapStage::StdoutStreaming, + last_stdout_seq: self.last_stdout_seq, + last_stderr_seq: self.last_stderr_seq, + marker: None, + })); + + self.stage = BootstrapStage::BootChecking; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::BootChecking), + )); + if let Err(error) = self + .client + .run_verify_commands(&self.spec.boot.verify_commands) + { + self.stage = BootstrapStage::BootCheckFailed; + events.push(BootstrapSessionEvent::Failed(format!( + "boot check failed: {}", + error.reason + ))); + return events; + } + + self.stage = BootstrapStage::SwactorStarting; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::SwactorStarting), + )); + if let Err(error) = self + .client + .start_swactor(&self.spec.boot.start_swactor_command, &self.spec.swarm_join) + { + self.stage = BootstrapStage::StartFailed; + events.push(BootstrapSessionEvent::Failed(format!( + "swactor start failed: {}", + error.reason + ))); + return events; + } + + self.stage = BootstrapStage::WaitingForSwactorJoin; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::WaitingForSwactorJoin), + )); + events + } + + pub fn convergence_observed( + &mut self, + _swactor_id: SwactorId, + sink: &mut dyn BootstrapDatastreamSink, + ) -> Vec { + self.stage = BootstrapStage::Converged; + sink.flush(); + self.client.close(); + self.closed = true; + self.stage = BootstrapStage::Closed; + vec![ + BootstrapSessionEvent::Observed(BootstrapObservation::stage(BootstrapStage::Converged)), + BootstrapSessionEvent::Closed, + ] + } + + pub fn cancel(&mut self) -> Vec { + self.stage = BootstrapStage::Cancelled; + self.client.close(); + self.closed = true; + vec![BootstrapSessionEvent::Closed] + } +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SwactorJoinEvent { + pub logical_node_id: LogicalNodeId, + pub swactor_id: SwactorId, +} + +#[derive(Clone, Debug, Default)] +pub struct SwactorJoinRouter { + expected_nodes: BTreeSet, +} + +impl SwactorJoinRouter { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, logical_node_id: LogicalNodeId) { + self.expected_nodes.insert(logical_node_id); + } + + pub fn route( + &self, + event: SwactorJoinEvent, + manager: &mut NodeManager, + ) -> Result, DockerClusterError> { + if !self.expected_nodes.contains(&event.logical_node_id) { + return Err(DockerClusterError::UnknownNode(event.logical_node_id)); + } + manager + .handle(NodeManagerMsg::SwactorJoined { + logical_node_id: event.logical_node_id, + swactor_id: event.swactor_id, + }) + .map_err(|error| DockerClusterError::Node(error.reason)) + } +} + +pub trait SshBootstrapClientFactory { + type Client: BootstrapSshClient; + + fn client_for(&mut self, spec: &BootstrapSessionSpec) -> Self::Client; +} + +#[derive(Clone, Debug)] +pub struct ManagedDockerNode { + pub manager: NodeManager, + pub bootstrap: Option>, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DockerClusterError { + Provider(String), + Node(String), + UnknownNode(LogicalNodeId), + EndpointUnavailable(ProviderLeaseId), + MissingBootstrap(LogicalNodeId), +} + +pub struct DockerClusterHarness +where + D: DockerCli, + F: SshBootstrapClientFactory, + S: BootstrapDatastreamSink, +{ + provider: DockerProvider, + client_factory: F, + datastream: S, + router: SwactorJoinRouter, + nodes: BTreeMap>, + teardown_complete: bool, +} + +impl DockerClusterHarness +where + D: DockerCli, + F: SshBootstrapClientFactory, + S: BootstrapDatastreamSink, +{ + pub fn new(provider: DockerProvider, client_factory: F, datastream: S) -> Self { + Self { + provider, + client_factory, + datastream, + router: SwactorJoinRouter::new(), + nodes: BTreeMap::new(), + teardown_complete: false, + } + } + + pub fn provider(&self) -> &DockerProvider { + &self.provider + } + + pub fn datastream(&self) -> &S { + &self.datastream + } + + pub fn nodes(&self) -> &BTreeMap> { + &self.nodes + } + + pub fn records(&self) -> Vec { + self.nodes + .values() + .filter_map(|node| node.manager.record().cloned()) + .collect() + } + + pub fn all_ready(&self) -> bool { + !self.nodes.is_empty() && self.nodes.values().all(|node| node.manager.is_ready()) + } + + pub fn start_group(&mut self, group: &RunNodeGroupSpec) -> Result<(), DockerClusterError> { + for spec in expand_node_group(group) { + let logical_node_id = spec.logical_node_id.clone(); + self.router.register(logical_node_id.clone()); + let mut manager = NodeManager::new(); + let commands = manager + .handle(NodeManagerMsg::Start(spec)) + .map_err(|error| DockerClusterError::Node(error.reason))?; + match self.process_start_commands(&mut manager, commands) { + Ok(bootstrap) => { + self.nodes + .insert(logical_node_id, ManagedDockerNode { manager, bootstrap }); + } + Err(error) => { + self.nodes.insert( + logical_node_id, + ManagedDockerNode { + manager, + bootstrap: None, + }, + ); + let _ = self.teardown(); + return Err(error); + } + } + } + self.teardown_complete = false; + Ok(()) + } + + pub fn route_join(&mut self, event: SwactorJoinEvent) -> Result<(), DockerClusterError> { + let logical_node_id = event.logical_node_id.clone(); + let node = self + .nodes + .get_mut(&logical_node_id) + .ok_or_else(|| DockerClusterError::UnknownNode(logical_node_id.clone()))?; + let commands = self.router.route(event, &mut node.manager)?; + for command in commands { + match command { + NodeManagerCommand::BootstrapConvergenceObserved { swactor_id, .. } => { + let bootstrap = node.bootstrap.as_mut().ok_or_else(|| { + DockerClusterError::MissingBootstrap(logical_node_id.clone()) + })?; + let events = bootstrap.convergence_observed(swactor_id, &mut self.datastream); + Self::feed_bootstrap_events(&mut node.manager, events)?; + } + other => { + return Err(DockerClusterError::Node(format!( + "unexpected command {other:?}" + ))); + } + } + } + Ok(()) + } + + pub fn teardown(&mut self) -> Result<(), DockerClusterError> { + for node in self.nodes.values_mut() { + let commands = node + .manager + .handle(NodeManagerMsg::Destroy) + .map_err(|error| DockerClusterError::Node(error.reason))?; + for command in commands { + match command { + NodeManagerCommand::CancelBootstrap { .. } => { + if let Some(bootstrap) = node.bootstrap.as_mut() { + let _ = bootstrap.cancel(); + } + } + NodeManagerCommand::DestroyLease(handle) => { + self.provider + .destroy_lease(&handle) + .map_err(|error| DockerClusterError::Provider(error.reason))?; + node.manager + .handle(NodeManagerMsg::LeaseDestroyed) + .map_err(|error| DockerClusterError::Node(error.reason))?; + } + other => { + return Err(DockerClusterError::Node(format!( + "unexpected command {other:?}" + ))); + } + } + } + } + self.teardown_complete = true; + Ok(()) + } + + fn process_start_commands( + &mut self, + manager: &mut NodeManager, + commands: Vec, + ) -> Result>, DockerClusterError> { + let mut pending = commands; + let mut bootstrap = None; + while let Some(command) = pending.pop() { + match command { + NodeManagerCommand::CreateLease(request) => { + let result = self + .provider + .create_lease(request) + .map_err(|error| DockerClusterError::Provider(error.reason))?; + let more = manager + .handle(NodeManagerMsg::LeaseCreated(result)) + .map_err(|error| DockerClusterError::Node(error.reason))?; + pending.extend(more); + } + NodeManagerCommand::LookupEndpoint(lease) => { + let lease_id = lease.lease_id.clone(); + let endpoint = match self + .provider + .lookup_endpoint(&lease) + .map_err(|error| DockerClusterError::Provider(error.reason))? + { + Some(endpoint) => endpoint, + None => { + let _ = manager.handle(NodeManagerMsg::EndpointFailed( + "docker ssh endpoint unavailable".into(), + )); + return Err(DockerClusterError::EndpointUnavailable(lease_id)); + } + }; + let more = manager + .handle(NodeManagerMsg::EndpointKnown(endpoint)) + .map_err(|error| DockerClusterError::Node(error.reason))?; + pending.extend(more); + } + NodeManagerCommand::StartBootstrap(spec) => { + let client = self.client_factory.client_for(&spec); + let mut session = SshBootstrapSession::new(spec, client); + let events = session.start(&mut self.datastream); + Self::feed_bootstrap_events(manager, events)?; + bootstrap = Some(session); + } + other => { + return Err(DockerClusterError::Node(format!( + "unexpected command {other:?}" + ))); + } + } + } + Ok(bootstrap) + } + + fn feed_bootstrap_events( + manager: &mut NodeManager, + events: Vec, + ) -> Result<(), DockerClusterError> { + for event in events { + match event { + BootstrapSessionEvent::Observed(observation) => { + manager + .handle(NodeManagerMsg::BootstrapObserved(observation)) + .map_err(|error| DockerClusterError::Node(error.reason))?; + } + BootstrapSessionEvent::Failed(reason) => { + manager + .handle(NodeManagerMsg::BootstrapFailed(reason)) + .map_err(|error| DockerClusterError::Node(error.reason))?; + } + BootstrapSessionEvent::Closed => { + manager + .handle(NodeManagerMsg::BootstrapClosed) + .map_err(|error| DockerClusterError::Node(error.reason))?; + } + } + } + Ok(()) + } +} + +impl Drop for DockerClusterHarness +where + D: DockerCli, + F: SshBootstrapClientFactory, + S: BootstrapDatastreamSink, +{ + fn drop(&mut self) { + if !self.teardown_complete { + let _ = self.teardown(); + } + } +} diff --git a/crates/mvp-system/src/lib.rs b/crates/mvp-system/src/lib.rs index 910ae64..4971001 100644 --- a/crates/mvp-system/src/lib.rs +++ b/crates/mvp-system/src/lib.rs @@ -7,6 +7,7 @@ pub mod arena_manager; pub mod dashboard_view; pub mod device_bridge; pub mod distribution_stack; +pub mod docker_cluster_provisioning; pub mod driver_pumps; pub mod edge_establisher; pub mod engine_builder; @@ -16,6 +17,7 @@ pub mod gpu_worker_ingress_parser; pub mod gpu_worker_process_adapter; pub mod membership_pool_readiness; pub mod node_boot_lifecycle; +pub mod node_provisioning; pub mod observability_surface; pub mod orchestrator_run_fsm; pub mod orchestrator_token_endpoint; diff --git a/crates/mvp-system/src/node_provisioning.rs b/crates/mvp-system/src/node_provisioning.rs new file mode 100644 index 0000000..cb20809 --- /dev/null +++ b/crates/mvp-system/src/node_provisioning.rs @@ -0,0 +1,964 @@ +//! Node-local provisioning and bootstrap state machines for the MVP node +//! provisioning spec. +//! +//! The module is intentionally in-process and deterministic. `NodeManager` is the +//! actor core: it owns one node record and emits commands for provider and +//! bootstrap effects. `BootstrapSession` is the transient pre-swactor SSH core. +//! Tests can drive both without Vast.ai, Docker, or real SSH. + +use std::collections::{BTreeMap, VecDeque}; +use std::time::SystemTime; + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct RunId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct LogicalNodeId(pub String); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct NodeGroupId(pub String); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct RoleId(pub String); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct ProviderLeaseId(pub String); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct SwactorId(pub String); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct BootstrapSessionId(pub u64); + +#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub struct DatastreamStreamId(pub String); + +#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +pub enum ProviderKind { + Mock, + Docker, + VastAi, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct DesiredNodeShape { + pub image: String, + pub disk_gb: u32, + pub gpu_name: Option, + pub min_gpu_ram_mb: Option, + pub min_down_mbps: Option, + pub min_up_mbps: Option, + pub min_reliability: Option, + pub require_verified: bool, + pub provider_labels: BTreeMap, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct BootstrapTimeoutPolicy { + pub ssh_connect_secs: u64, + pub boot_check_secs: u64, + pub swactor_join_secs: u64, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct BootSpec { + pub ssh_user: String, + pub verify_commands: Vec, + pub start_swactor_command: String, + pub stdout_sources: Vec, + pub stderr_sources: Vec, + pub timeout_policy: BootstrapTimeoutPolicy, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SwarmJoinTemplate { + pub orch_swactor_addr: String, + pub join_token_ref: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SwarmJoinSpec { + pub orch_swactor_addr: String, + pub join_token_ref: String, + pub expected_logical_node_id: LogicalNodeId, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct RunNodeGroupSpec { + pub run_id: RunId, + pub group_id: NodeGroupId, + pub role: RoleId, + pub count: u32, + pub provider: ProviderKind, + pub shape: DesiredNodeShape, + pub boot: BootSpec, + pub swarm_join: SwarmJoinTemplate, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct LogicalNodeSpec { + pub run_id: RunId, + pub logical_node_id: LogicalNodeId, + pub group_id: NodeGroupId, + pub role: RoleId, + pub provider: ProviderKind, + pub shape: DesiredNodeShape, + pub boot: BootSpec, + pub swarm_join: SwarmJoinSpec, +} + +pub fn expand_node_group(group: &RunNodeGroupSpec) -> Vec { + (0..group.count) + .map(|index| { + let logical_node_id = LogicalNodeId(format!("{}-{index}", group.group_id.0)); + LogicalNodeSpec { + run_id: group.run_id.clone(), + logical_node_id: logical_node_id.clone(), + group_id: group.group_id.clone(), + role: group.role.clone(), + provider: group.provider, + shape: group.shape.clone(), + boot: group.boot.clone(), + swarm_join: SwarmJoinSpec { + orch_swactor_addr: group.swarm_join.orch_swactor_addr.clone(), + join_token_ref: group.swarm_join.join_token_ref.clone(), + expected_logical_node_id: logical_node_id, + }, + } + }) + .collect() +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum NodeStage { + New, + LeaseRequested, + LeaseCreated, + EndpointKnown, + BootstrapRunning, + SwactorJoined, + HandedOff, + Dormant, + Failed, + Destroyed, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum BootstrapStage { + Created, + SshConnecting, + SshReady, + StdoutStreaming, + BootChecking, + SwactorStarting, + WaitingForSwactorJoin, + Converged, + Closed, + SshTimeout, + BootCheckFailed, + StartFailed, + JoinTimeout, + StreamError, + Cancelled, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct DestroyHandle { + pub provider: ProviderKind, + pub lease_id: ProviderLeaseId, + pub provider_contract_id: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct LeaseFacts { + pub provider: ProviderKind, + pub lease_id: ProviderLeaseId, + pub provider_contract_id: String, + pub offer_id: Option, + pub destroy_handle: DestroyHandle, + pub provider_metadata: BTreeMap, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SshEndpoint { + pub host: String, + pub port: u16, + pub user: String, + pub auth_ref: String, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct BootstrapFacts { + pub session_id: BootstrapSessionId, + pub last_stage: BootstrapStage, + pub last_stdout_seq: Option, + pub last_stderr_seq: Option, + pub last_observed_at: SystemTime, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct SwactorFacts { + pub swactor_id: SwactorId, + pub joined_at: SystemTime, + pub handed_off_at: Option, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct NodeRecord { + pub logical_node_id: LogicalNodeId, + pub run_id: RunId, + pub group_id: NodeGroupId, + pub role: RoleId, + pub desired: LogicalNodeSpec, + pub stage: NodeStage, + pub ready: bool, + pub lease: Option, + pub connection: Option, + pub bootstrap: Option, + pub swactor: Option, + pub failed_reason: Option, + pub destroyed_at: Option, +} + +impl NodeRecord { + pub fn from_spec(spec: LogicalNodeSpec) -> Self { + Self { + logical_node_id: spec.logical_node_id.clone(), + run_id: spec.run_id.clone(), + group_id: spec.group_id.clone(), + role: spec.role.clone(), + desired: spec, + stage: NodeStage::New, + ready: false, + lease: None, + connection: None, + bootstrap: None, + swactor: None, + failed_reason: None, + destroyed_at: None, + } + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CreateLeaseRequest { + pub spec: LogicalNodeSpec, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct CreateLeaseResult { + pub lease: LeaseFacts, + pub endpoint: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProviderError { + pub reason: String, +} + +impl ProviderError { + pub fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +pub trait ProviderPlugin { + fn create_lease( + &mut self, + request: CreateLeaseRequest, + ) -> Result; + + fn lookup_endpoint(&mut self, lease: &LeaseFacts) + -> Result, ProviderError>; + + fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError>; +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct BootstrapObservation { + pub stage: BootstrapStage, + pub last_stdout_seq: Option, + pub last_stderr_seq: Option, + pub marker: Option, +} + +impl BootstrapObservation { + pub fn stage(stage: BootstrapStage) -> Self { + Self { + stage, + last_stdout_seq: None, + last_stderr_seq: None, + marker: None, + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub enum NodeManagerMsg { + Start(LogicalNodeSpec), + LeaseCreated(CreateLeaseResult), + LeaseFailed(String), + EndpointKnown(SshEndpoint), + EndpointFailed(String), + BootstrapObserved(BootstrapObservation), + BootstrapFailed(String), + BootstrapClosed, + SwactorJoined { + logical_node_id: LogicalNodeId, + swactor_id: SwactorId, + }, + Destroy, + LeaseDestroyed, + DestroyFailed(String), +} + +#[derive(Clone, Debug, PartialEq)] +pub enum NodeManagerCommand { + CreateLease(CreateLeaseRequest), + LookupEndpoint(LeaseFacts), + StartBootstrap(BootstrapSessionSpec), + BootstrapConvergenceObserved { + session_id: BootstrapSessionId, + swactor_id: SwactorId, + }, + CancelBootstrap { + session_id: BootstrapSessionId, + }, + DestroyLease(DestroyHandle), +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct NodeManagerError { + pub reason: String, +} + +impl NodeManagerError { + fn new(reason: impl Into) -> Self { + Self { + reason: reason.into(), + } + } +} + +#[derive(Clone, Debug)] +pub struct NodeManager { + record: Option, + active_bootstrap: Option, + next_bootstrap_session_id: u64, + lease_destroyed: bool, +} + +impl Default for NodeManager { + fn default() -> Self { + Self { + record: None, + active_bootstrap: None, + next_bootstrap_session_id: 1, + lease_destroyed: false, + } + } +} + +impl NodeManager { + pub fn new() -> Self { + Self::default() + } + + pub fn record(&self) -> Option<&NodeRecord> { + self.record.as_ref() + } + + pub fn is_ready(&self) -> bool { + self.record.as_ref().is_some_and(|record| record.ready) + } + + pub fn active_bootstrap(&self) -> Option { + self.active_bootstrap + } + + pub fn handle( + &mut self, + msg: NodeManagerMsg, + ) -> Result, NodeManagerError> { + match msg { + NodeManagerMsg::Start(spec) => self.start(spec), + NodeManagerMsg::LeaseCreated(result) => self.lease_created(result), + NodeManagerMsg::LeaseFailed(reason) => self.fail(reason), + NodeManagerMsg::EndpointKnown(endpoint) => self.endpoint_known(endpoint), + NodeManagerMsg::EndpointFailed(reason) => self.fail(reason), + NodeManagerMsg::BootstrapObserved(observation) => self.bootstrap_observed(observation), + NodeManagerMsg::BootstrapFailed(reason) => self.fail(reason), + NodeManagerMsg::BootstrapClosed => self.bootstrap_closed(), + NodeManagerMsg::SwactorJoined { + logical_node_id, + swactor_id, + } => self.swactor_joined(logical_node_id, swactor_id), + NodeManagerMsg::Destroy => self.destroy(), + NodeManagerMsg::LeaseDestroyed => self.lease_destroyed(), + NodeManagerMsg::DestroyFailed(reason) => self.destroy_failed(reason), + } + } + + fn start( + &mut self, + spec: LogicalNodeSpec, + ) -> Result, NodeManagerError> { + if self.record.is_some() { + return Err(NodeManagerError::new("node manager already started")); + } + let mut record = NodeRecord::from_spec(spec.clone()); + record.stage = NodeStage::LeaseRequested; + self.record = Some(record); + Ok(vec![NodeManagerCommand::CreateLease(CreateLeaseRequest { + spec, + })]) + } + + fn lease_created( + &mut self, + result: CreateLeaseResult, + ) -> Result, NodeManagerError> { + let record = self.record_mut()?; + record.lease = Some(result.lease.clone()); + record.stage = NodeStage::LeaseCreated; + if let Some(endpoint) = result.endpoint { + self.begin_bootstrap(endpoint) + } else { + Ok(vec![NodeManagerCommand::LookupEndpoint(result.lease)]) + } + } + + fn endpoint_known( + &mut self, + endpoint: SshEndpoint, + ) -> Result, NodeManagerError> { + if self.require_record()?.lease.is_none() { + return Err(NodeManagerError::new("endpoint cannot arrive before lease")); + } + self.begin_bootstrap(endpoint) + } + + fn begin_bootstrap( + &mut self, + endpoint: SshEndpoint, + ) -> Result, NodeManagerError> { + if self.active_bootstrap.is_some() { + return Err(NodeManagerError::new("bootstrap already active")); + } + let session_id = BootstrapSessionId(self.next_bootstrap_session_id); + self.next_bootstrap_session_id = self.next_bootstrap_session_id.wrapping_add(1).max(1); + let (run_id, logical_node_id, lease_id, boot, swarm_join, timeout_policy) = { + let record = self.record_mut()?; + let lease_id = record + .lease + .as_ref() + .ok_or_else(|| NodeManagerError::new("bootstrap requires known lease"))? + .lease_id + .clone(); + record.connection = Some(endpoint.clone()); + record.stage = NodeStage::EndpointKnown; + record.bootstrap = Some(BootstrapFacts { + session_id, + last_stage: BootstrapStage::Created, + last_stdout_seq: None, + last_stderr_seq: None, + last_observed_at: SystemTime::now(), + }); + record.stage = NodeStage::BootstrapRunning; + ( + record.run_id.clone(), + record.logical_node_id.clone(), + lease_id, + record.desired.boot.clone(), + record.desired.swarm_join.clone(), + record.desired.boot.timeout_policy, + ) + }; + self.active_bootstrap = Some(session_id); + Ok(vec![NodeManagerCommand::StartBootstrap( + BootstrapSessionSpec { + datastream: DatastreamStreamId(format!( + "run/{}/node/{}/bootstrap", + run_id.0, logical_node_id.0 + )), + run_id, + logical_node_id, + lease_id, + ssh: endpoint, + boot, + swarm_join, + timeout_policy, + }, + )]) + } + + fn bootstrap_observed( + &mut self, + observation: BootstrapObservation, + ) -> Result, NodeManagerError> { + if self.active_bootstrap.is_none() { + return Err(NodeManagerError::new( + "bootstrap observation without active session", + )); + } + let record = self.record_mut()?; + let facts = record + .bootstrap + .as_mut() + .ok_or_else(|| NodeManagerError::new("missing bootstrap facts"))?; + facts.last_stage = observation.stage; + facts.last_observed_at = SystemTime::now(); + if observation.last_stdout_seq.is_some() { + facts.last_stdout_seq = observation.last_stdout_seq; + } + if observation.last_stderr_seq.is_some() { + facts.last_stderr_seq = observation.last_stderr_seq; + } + Ok(Vec::new()) + } + + fn swactor_joined( + &mut self, + logical_node_id: LogicalNodeId, + swactor_id: SwactorId, + ) -> Result, NodeManagerError> { + let expected = self.require_record()?.logical_node_id.clone(); + if logical_node_id != expected { + return Err(NodeManagerError::new(format!( + "swactor join for {}, expected {}", + logical_node_id.0, expected.0 + ))); + } + let session_id = self + .active_bootstrap + .ok_or_else(|| NodeManagerError::new("swactor join without active bootstrap"))?; + let record = self.record_mut()?; + record.swactor = Some(SwactorFacts { + swactor_id: swactor_id.clone(), + joined_at: SystemTime::now(), + handed_off_at: None, + }); + record.stage = NodeStage::SwactorJoined; + Ok(vec![NodeManagerCommand::BootstrapConvergenceObserved { + session_id, + swactor_id, + }]) + } + + fn bootstrap_closed(&mut self) -> Result, NodeManagerError> { + let record = self.record_mut()?; + if record.stage != NodeStage::SwactorJoined { + return Err(NodeManagerError::new( + "bootstrap closed before swactor convergence", + )); + } + let swactor = record + .swactor + .as_mut() + .ok_or_else(|| NodeManagerError::new("handoff requires swactor facts"))?; + swactor.handed_off_at = Some(SystemTime::now()); + record.stage = NodeStage::HandedOff; + record.ready = true; + record.stage = NodeStage::Dormant; + self.active_bootstrap = None; + Ok(Vec::new()) + } + + fn destroy(&mut self) -> Result, NodeManagerError> { + let mut commands = Vec::new(); + let active_bootstrap = self.active_bootstrap.take(); + let lease_already_destroyed = self.lease_destroyed; + let lease_command = { + let record = self.record_mut()?; + record.ready = false; + if let Some(session_id) = active_bootstrap { + commands.push(NodeManagerCommand::CancelBootstrap { session_id }); + } + if lease_already_destroyed { + None + } else { + record + .lease + .as_ref() + .map(|lease| NodeManagerCommand::DestroyLease(lease.destroy_handle.clone())) + } + }; + if let Some(command) = lease_command { + commands.push(command); + } else { + let record = self.record_mut()?; + record.stage = NodeStage::Destroyed; + record.destroyed_at = Some(SystemTime::now()); + } + Ok(commands) + } + + fn lease_destroyed(&mut self) -> Result, NodeManagerError> { + let record = self.record_mut()?; + record.stage = NodeStage::Destroyed; + record.ready = false; + record.destroyed_at = Some(SystemTime::now()); + self.lease_destroyed = true; + Ok(Vec::new()) + } + + fn destroy_failed( + &mut self, + reason: String, + ) -> Result, NodeManagerError> { + self.fail(format!("destroy: {reason}")) + } + + fn fail(&mut self, reason: String) -> Result, NodeManagerError> { + let record = self.record_mut()?; + record.stage = NodeStage::Failed; + record.ready = false; + record.failed_reason = Some(reason); + Ok(Vec::new()) + } + + fn require_record(&self) -> Result<&NodeRecord, NodeManagerError> { + self.record + .as_ref() + .ok_or_else(|| NodeManagerError::new("node manager not started")) + } + + fn record_mut(&mut self) -> Result<&mut NodeRecord, NodeManagerError> { + self.record + .as_mut() + .ok_or_else(|| NodeManagerError::new("node manager not started")) + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum BootstrapLogSource { + SshBootstrap, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub enum BootstrapLogStream { + Stdout, + Stderr, +} + +#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] +pub struct BootstrapLogRecord { + pub run_id: RunId, + pub logical_node_id: LogicalNodeId, + pub lease_id: ProviderLeaseId, + pub source: BootstrapLogSource, + pub stream: BootstrapLogStream, + pub seq: u64, + pub timestamp: SystemTime, + pub line: String, +} + +pub trait BootstrapDatastreamSink { + fn record(&mut self, record: BootstrapLogRecord); + fn flush(&mut self); +} + +#[derive(Default, Debug, Clone, PartialEq, Eq)] +pub struct InMemoryBootstrapDatastream { + records: Vec, + flush_count: usize, +} + +impl InMemoryBootstrapDatastream { + pub fn records(&self) -> &[BootstrapLogRecord] { + &self.records + } + + pub fn flush_count(&self) -> usize { + self.flush_count + } +} + +impl BootstrapDatastreamSink for InMemoryBootstrapDatastream { + fn record(&mut self, record: BootstrapLogRecord) { + self.records.push(record); + } + + fn flush(&mut self) { + self.flush_count += 1; + } +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct BootstrapSessionSpec { + pub run_id: RunId, + pub logical_node_id: LogicalNodeId, + pub lease_id: ProviderLeaseId, + pub ssh: SshEndpoint, + pub boot: BootSpec, + pub swarm_join: SwarmJoinSpec, + pub datastream: DatastreamStreamId, + pub timeout_policy: BootstrapTimeoutPolicy, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BootstrapSessionEvent { + Observed(BootstrapObservation), + Failed(String), + Closed, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct MockBootstrapScript { + pub ssh_ok: bool, + pub verify_ok: bool, + pub start_ok: bool, + pub records: Vec<(BootstrapLogStream, String)>, +} + +impl MockBootstrapScript { + pub fn successful(records: Vec<(BootstrapLogStream, String)>) -> Self { + Self { + ssh_ok: true, + verify_ok: true, + start_ok: true, + records, + } + } +} + +#[derive(Clone, Debug)] +pub struct BootstrapSession { + spec: BootstrapSessionSpec, + stage: BootstrapStage, + next_seq: u64, + last_stdout_seq: Option, + last_stderr_seq: Option, + closed: bool, +} + +impl BootstrapSession { + pub fn new(spec: BootstrapSessionSpec) -> Self { + Self { + spec, + stage: BootstrapStage::Created, + next_seq: 1, + last_stdout_seq: None, + last_stderr_seq: None, + closed: false, + } + } + + pub fn stage(&self) -> BootstrapStage { + self.stage + } + + pub fn is_closed(&self) -> bool { + self.closed + } + + pub fn start( + &mut self, + script: &MockBootstrapScript, + sink: &mut dyn BootstrapDatastreamSink, + ) -> Vec { + let mut events = Vec::new(); + self.stage = BootstrapStage::SshConnecting; + if !script.ssh_ok { + self.stage = BootstrapStage::SshTimeout; + return vec![BootstrapSessionEvent::Failed("ssh timeout".into())]; + } + self.stage = BootstrapStage::SshReady; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::SshReady), + )); + + self.stage = BootstrapStage::StdoutStreaming; + for (stream, line) in &script.records { + let seq = self.next_seq; + self.next_seq += 1; + match stream { + BootstrapLogStream::Stdout => self.last_stdout_seq = Some(seq), + BootstrapLogStream::Stderr => self.last_stderr_seq = Some(seq), + } + sink.record(BootstrapLogRecord { + run_id: self.spec.run_id.clone(), + logical_node_id: self.spec.logical_node_id.clone(), + lease_id: self.spec.lease_id.clone(), + source: BootstrapLogSource::SshBootstrap, + stream: *stream, + seq, + timestamp: SystemTime::now(), + line: line.clone(), + }); + } + events.push(BootstrapSessionEvent::Observed(BootstrapObservation { + stage: BootstrapStage::StdoutStreaming, + last_stdout_seq: self.last_stdout_seq, + last_stderr_seq: self.last_stderr_seq, + marker: None, + })); + + self.stage = BootstrapStage::BootChecking; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::BootChecking), + )); + if !script.verify_ok { + self.stage = BootstrapStage::BootCheckFailed; + events.push(BootstrapSessionEvent::Failed("boot check failed".into())); + return events; + } + + self.stage = BootstrapStage::SwactorStarting; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::SwactorStarting), + )); + if !script.start_ok { + self.stage = BootstrapStage::StartFailed; + events.push(BootstrapSessionEvent::Failed("swactor start failed".into())); + return events; + } + + self.stage = BootstrapStage::WaitingForSwactorJoin; + events.push(BootstrapSessionEvent::Observed( + BootstrapObservation::stage(BootstrapStage::WaitingForSwactorJoin), + )); + events + } + + pub fn convergence_observed( + &mut self, + _swactor_id: SwactorId, + sink: &mut dyn BootstrapDatastreamSink, + ) -> Vec { + self.stage = BootstrapStage::Converged; + sink.flush(); + self.closed = true; + self.stage = BootstrapStage::Closed; + vec![ + BootstrapSessionEvent::Observed(BootstrapObservation::stage(BootstrapStage::Converged)), + BootstrapSessionEvent::Closed, + ] + } + + pub fn join_timeout(&mut self) -> Vec { + self.stage = BootstrapStage::JoinTimeout; + vec![BootstrapSessionEvent::Failed("join timeout".into())] + } + + pub fn cancel(&mut self) -> Vec { + self.stage = BootstrapStage::Cancelled; + self.closed = true; + vec![BootstrapSessionEvent::Closed] + } +} + +#[derive(Default, Debug, Clone)] +pub struct MockProviderPlugin { + next_contract_id: u64, + create_results: VecDeque>, + endpoint_results: VecDeque, ProviderError>>, + create_requests: Vec, + lookup_requests: Vec, + destroyed_handles: Vec, + destroy_failures: BTreeMap, +} + +impl MockProviderPlugin { + pub fn new() -> Self { + Self { + next_contract_id: 1, + ..Self::default() + } + } + + pub fn queue_create_result(&mut self, result: Result) { + self.create_results.push_back(result); + } + + pub fn queue_endpoint_result(&mut self, result: Result, ProviderError>) { + self.endpoint_results.push_back(result); + } + + pub fn fail_destroy(&mut self, lease_id: ProviderLeaseId, reason: impl Into) { + self.destroy_failures.insert(lease_id, reason.into()); + } + + pub fn create_requests(&self) -> &[CreateLeaseRequest] { + &self.create_requests + } + + pub fn lookup_requests(&self) -> &[ProviderLeaseId] { + &self.lookup_requests + } + + pub fn destroyed_handles(&self) -> &[DestroyHandle] { + &self.destroyed_handles + } + + pub fn result_with_endpoint(id: u64, endpoint: Option) -> CreateLeaseResult { + let lease_id = ProviderLeaseId(format!("mock:{id}")); + CreateLeaseResult { + lease: LeaseFacts { + provider: ProviderKind::Mock, + lease_id: lease_id.clone(), + provider_contract_id: id.to_string(), + offer_id: Some(format!("offer-{id}")), + destroy_handle: DestroyHandle { + provider: ProviderKind::Mock, + lease_id, + provider_contract_id: id.to_string(), + }, + provider_metadata: BTreeMap::new(), + }, + endpoint, + } + } + + fn default_endpoint(id: u64) -> SshEndpoint { + SshEndpoint { + host: "127.0.0.1".into(), + port: 22000 + id as u16, + user: "root".into(), + auth_ref: format!("mock-auth-{id}"), + } + } +} + +impl ProviderPlugin for MockProviderPlugin { + fn create_lease( + &mut self, + request: CreateLeaseRequest, + ) -> Result { + self.create_requests.push(request); + if let Some(result) = self.create_results.pop_front() { + return result; + } + let id = self.next_contract_id; + self.next_contract_id = self.next_contract_id.wrapping_add(1).max(1); + Ok(Self::result_with_endpoint( + id, + Some(Self::default_endpoint(id)), + )) + } + + fn lookup_endpoint( + &mut self, + lease: &LeaseFacts, + ) -> Result, ProviderError> { + self.lookup_requests.push(lease.lease_id.clone()); + if let Some(result) = self.endpoint_results.pop_front() { + return result; + } + Ok(Some(Self::default_endpoint( + lease.provider_contract_id.parse().unwrap_or(1), + ))) + } + + fn destroy_lease(&mut self, handle: &DestroyHandle) -> Result<(), ProviderError> { + if let Some(reason) = self.destroy_failures.get(&handle.lease_id) { + return Err(ProviderError::new(reason.clone())); + } + self.destroyed_handles.push(handle.clone()); + Ok(()) + } +} diff --git a/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs b/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs new file mode 100644 index 0000000..53943ad --- /dev/null +++ b/crates/mvp-system/src/tests/docker_cluster_provisioning_guarantees.rs @@ -0,0 +1,466 @@ +use std::collections::BTreeMap; + +use mvp_system::docker_cluster_provisioning as docker; +use mvp_system::node_provisioning as provision; +use provision::ProviderPlugin; + +fn docker_group_spec(count: u32) -> provision::RunNodeGroupSpec { + provision::RunNodeGroupSpec { + run_id: provision::RunId(7), + group_id: provision::NodeGroupId("workers".into()), + role: provision::RoleId("worker".into()), + count, + provider: provision::ProviderKind::Docker, + shape: provision::DesiredNodeShape { + image: "mvp-worker-test:latest".into(), + disk_gb: 32, + gpu_name: None, + min_gpu_ram_mb: None, + min_down_mbps: None, + min_up_mbps: None, + min_reliability: None, + require_verified: false, + provider_labels: [("suite".into(), "docker-provisioning".into())] + .into_iter() + .collect(), + }, + boot: provision::BootSpec { + ssh_user: "root".into(), + verify_commands: vec!["test -x /opt/mvp/swactor".into()], + start_swactor_command: "/opt/mvp/swactor-node --join ${MVP_ORCH_SWACTOR_ADDR}".into(), + stdout_sources: vec!["/var/log/mvp/bootstrap.out".into()], + stderr_sources: vec!["/var/log/mvp/bootstrap.err".into()], + timeout_policy: provision::BootstrapTimeoutPolicy { + ssh_connect_secs: 10, + boot_check_secs: 20, + swactor_join_secs: 30, + }, + }, + swarm_join: provision::SwarmJoinTemplate { + orch_swactor_addr: "quic://orch.local:9443".into(), + join_token_ref: "secret://run-7-token".into(), + }, + } +} + +fn one_spec() -> provision::LogicalNodeSpec { + provision::expand_node_group(&docker_group_spec(1)) + .into_iter() + .next() + .expect("one logical node") +} + +fn endpoint(port: u16) -> provision::SshEndpoint { + provision::SshEndpoint { + host: "127.0.0.1".into(), + port, + user: "root".into(), + auth_ref: format!("test-key-{port}"), + } +} + +#[derive(Clone, Debug, Default)] +struct FakeDockerCli { + next_id: usize, + run_requests: Vec, + inspect_requests: Vec, + removed: Vec, + endpoints: BTreeMap>, +} + +impl FakeDockerCli { + fn new() -> Self { + Self { + next_id: 1, + ..Self::default() + } + } +} + +impl docker::DockerCli for FakeDockerCli { + fn run_container( + &mut self, + request: docker::DockerRunRequest, + ) -> Result { + self.run_requests.push(request); + let container_id = format!("container-{}", self.next_id); + let ssh_endpoint = endpoint(22000 + self.next_id as u16); + self.next_id += 1; + self.endpoints + .entry(container_id.clone()) + .or_insert_with(|| Some(ssh_endpoint)); + Ok(docker::DockerRunResult { container_id }) + } + + fn inspect_ssh_endpoint( + &mut self, + container_id: &str, + ) -> Result, docker::DockerCliError> { + self.inspect_requests.push(container_id.into()); + Ok(self.endpoints.get(container_id).cloned().flatten()) + } + + fn remove_force(&mut self, container_id: &str) -> Result<(), docker::DockerCliError> { + self.removed.push(container_id.into()); + Ok(()) + } +} + +#[derive(Clone, Debug)] +struct FakeSshClient { + logs: Vec<(provision::BootstrapLogStream, String)>, + connected_to: Option, + probed: bool, + stdout_sources: Vec, + stderr_sources: Vec, + verify_commands: Vec, + started_command: Option, + started_join: Option, + closed: bool, +} + +impl FakeSshClient { + fn new(logical_node_id: &provision::LogicalNodeId) -> Self { + Self { + logs: vec![ + ( + provision::BootstrapLogStream::Stdout, + format!("{} boot entered", logical_node_id.0), + ), + ( + provision::BootstrapLogStream::Stderr, + format!("{} stderr ready", logical_node_id.0), + ), + ], + connected_to: None, + probed: false, + stdout_sources: Vec::new(), + stderr_sources: Vec::new(), + verify_commands: Vec::new(), + started_command: None, + started_join: None, + closed: false, + } + } +} + +impl docker::BootstrapSshClient for FakeSshClient { + fn connect( + &mut self, + endpoint: &provision::SshEndpoint, + ) -> Result<(), docker::BootstrapSshError> { + self.connected_to = Some(endpoint.clone()); + Ok(()) + } + + fn probe_stdout(&mut self) -> Result<(), docker::BootstrapSshError> { + self.probed = true; + Ok(()) + } + + fn read_bootstrap_logs( + &mut self, + stdout_sources: &[String], + stderr_sources: &[String], + ) -> Result, docker::BootstrapSshError> { + self.stdout_sources = stdout_sources.to_vec(); + self.stderr_sources = stderr_sources.to_vec(); + Ok(self.logs.clone()) + } + + fn run_verify_commands( + &mut self, + commands: &[String], + ) -> Result<(), docker::BootstrapSshError> { + self.verify_commands = commands.to_vec(); + Ok(()) + } + + fn start_swactor( + &mut self, + command: &str, + join: &provision::SwarmJoinSpec, + ) -> Result<(), docker::BootstrapSshError> { + self.started_command = Some(command.into()); + self.started_join = Some(join.clone()); + Ok(()) + } + + fn close(&mut self) { + self.closed = true; + } +} + +#[derive(Clone, Debug, Default)] +struct FakeSshFactory { + created_for: Vec, +} + +impl docker::SshBootstrapClientFactory for FakeSshFactory { + type Client = FakeSshClient; + + fn client_for(&mut self, spec: &provision::BootstrapSessionSpec) -> Self::Client { + self.created_for.push(spec.logical_node_id.clone()); + FakeSshClient::new(&spec.logical_node_id) + } +} + +fn started_manager() -> (provision::NodeManager, provision::LogicalNodeId) { + let spec = one_spec(); + let logical_node_id = spec.logical_node_id.clone(); + let mut manager = provision::NodeManager::new(); + let start_commands = manager + .handle(provision::NodeManagerMsg::Start(spec.clone())) + .expect("start succeeds"); + assert!(matches!( + start_commands.as_slice(), + [provision::NodeManagerCommand::CreateLease(_)] + )); + let lease = docker::DockerProvider::new(FakeDockerCli::new()) + .create_lease(provision::CreateLeaseRequest { spec }) + .expect("docker provider creates lease"); + let bootstrap_commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + assert!(matches!( + bootstrap_commands.as_slice(), + [provision::NodeManagerCommand::StartBootstrap(_)] + )); + (manager, logical_node_id) +} + +#[test] +fn docker_provider_maps_node_spec_to_container_lease_and_destroy_handle() { + let spec = one_spec(); + let mut provider = docker::DockerProvider::new(FakeDockerCli::new()); + + let result = provider + .create_lease(provision::CreateLeaseRequest { spec: spec.clone() }) + .expect("docker provider creates lease"); + + let cli = provider.cli(); + assert_eq!(cli.run_requests.len(), 1); + let run = &cli.run_requests[0]; + assert_eq!(run.container_name, "mvp-7-workers-0"); + assert_eq!(run.image, "mvp-worker-test:latest"); + assert_eq!(run.ssh_user, "root"); + assert_eq!(run.labels["mvp.provider"], "docker"); + assert_eq!(run.labels["mvp.logical_node_id"], "workers-0"); + assert_eq!(run.env["MVP_LOGICAL_NODE_ID"], "workers-0"); + assert_eq!(run.env["MVP_ORCH_SWACTOR_ADDR"], "quic://orch.local:9443"); + assert_eq!(cli.inspect_requests, vec!["container-1"]); + + assert_eq!(result.lease.provider, provision::ProviderKind::Docker); + assert_eq!( + result.lease.lease_id, + provision::ProviderLeaseId("docker:container-1".into()) + ); + assert_eq!(result.lease.provider_contract_id, "container-1"); + assert_eq!(result.endpoint, Some(endpoint(22001))); + + provider + .destroy_lease(&result.lease.destroy_handle) + .expect("destroy succeeds"); + assert_eq!(provider.cli().removed, vec!["container-1"]); +} + +#[test] +fn ssh_bootstrap_session_runs_pre_handoff_steps_and_closes_on_convergence() { + let spec = one_spec(); + let bootstrap_spec = provision::BootstrapSessionSpec { + run_id: spec.run_id.clone(), + logical_node_id: spec.logical_node_id.clone(), + lease_id: provision::ProviderLeaseId("docker:container-1".into()), + ssh: endpoint(22001), + boot: spec.boot.clone(), + swarm_join: spec.swarm_join.clone(), + datastream: provision::DatastreamStreamId("run/7/workers-0/bootstrap".into()), + timeout_policy: spec.boot.timeout_policy, + }; + let mut session = docker::SshBootstrapSession::new( + bootstrap_spec, + FakeSshClient::new(&provision::LogicalNodeId("workers-0".into())), + ); + let mut datastream = provision::InMemoryBootstrapDatastream::default(); + + let events = session.start(&mut datastream); + + assert_eq!( + session.stage(), + provision::BootstrapStage::WaitingForSwactorJoin + ); + assert!(events.iter().any(|event| matches!( + event, + provision::BootstrapSessionEvent::Observed(obs) + if obs.stage == provision::BootstrapStage::SshReady + ))); + assert_eq!(datastream.records().len(), 2); + assert_eq!(datastream.records()[0].seq, 1); + assert_eq!( + datastream.records()[0].stream, + provision::BootstrapLogStream::Stdout + ); + assert_eq!(datastream.records()[1].seq, 2); + assert_eq!( + datastream.records()[1].stream, + provision::BootstrapLogStream::Stderr + ); + assert_eq!(session.client().connected_to, Some(endpoint(22001))); + assert!(session.client().probed); + assert_eq!( + session.client().verify_commands, + vec!["test -x /opt/mvp/swactor"] + ); + assert_eq!( + session.client().started_command.as_deref(), + Some("/opt/mvp/swactor-node --join ${MVP_ORCH_SWACTOR_ADDR}") + ); + + let events = session.convergence_observed( + provision::SwactorId("swactor-workers-0".into()), + &mut datastream, + ); + assert!(session.is_closed()); + assert!(session.client().closed); + assert_eq!(datastream.flush_count(), 1); + assert!(matches!( + events.as_slice(), + [_, provision::BootstrapSessionEvent::Closed] + )); +} + +#[test] +fn join_router_routes_known_logical_nodes_and_rejects_unknown_nodes() { + let (mut manager, logical_node_id) = started_manager(); + let mut router = docker::SwactorJoinRouter::new(); + router.register(logical_node_id.clone()); + + let unknown = router.route( + docker::SwactorJoinEvent { + logical_node_id: provision::LogicalNodeId("workers-99".into()), + swactor_id: provision::SwactorId("swactor-wrong".into()), + }, + &mut manager, + ); + assert!(matches!( + unknown, + Err(docker::DockerClusterError::UnknownNode(id)) if id == provision::LogicalNodeId("workers-99".into()) + )); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::BootstrapRunning + ); + + let commands = router + .route( + docker::SwactorJoinEvent { + logical_node_id, + swactor_id: provision::SwactorId("swactor-workers-0".into()), + }, + &mut manager, + ) + .expect("known join routes"); + assert!(matches!( + commands.as_slice(), + [provision::NodeManagerCommand::BootstrapConvergenceObserved { .. }] + )); +} + +#[test] +fn docker_cluster_harness_wires_provider_bootstrap_join_and_teardown() { + let provider = docker::DockerProvider::new(FakeDockerCli::new()); + let datastream = provision::InMemoryBootstrapDatastream::default(); + let mut harness = + docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); + + harness + .start_group(&docker_group_spec(2)) + .expect("cluster starts"); + assert!(!harness.all_ready()); + assert_eq!(harness.nodes().len(), 2); + assert_eq!(harness.provider().cli().run_requests.len(), 2); + assert_eq!(harness.datastream().records().len(), 4); + + let ids: Vec<_> = harness.nodes().keys().cloned().collect(); + for id in ids { + harness + .route_join(docker::SwactorJoinEvent { + swactor_id: provision::SwactorId(format!("swactor-{}", id.0)), + logical_node_id: id, + }) + .expect("join routes"); + } + + assert!(harness.all_ready()); + assert_eq!(harness.datastream().flush_count(), 2); + for record in harness.records() { + assert_eq!(record.stage, provision::NodeStage::Dormant); + assert!(record.ready); + assert_eq!( + record.lease.as_ref().unwrap().provider, + provision::ProviderKind::Docker + ); + } + + harness.teardown().expect("teardown succeeds"); + assert_eq!( + harness.provider().cli().removed, + vec!["container-1", "container-2"] + ); + for record in harness.records() { + assert_eq!(record.stage, provision::NodeStage::Destroyed); + assert!(!record.ready); + } +} + +#[test] +fn docker_cluster_harness_teardown_cleans_known_leases_before_join() { + let provider = docker::DockerProvider::new(FakeDockerCli::new()); + let datastream = provision::InMemoryBootstrapDatastream::default(); + let mut harness = + docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); + + harness + .start_group(&docker_group_spec(1)) + .expect("cluster starts"); + let id = provision::LogicalNodeId("workers-0".into()); + assert_eq!( + harness.nodes()[&id].manager.record().unwrap().stage, + provision::NodeStage::BootstrapRunning + ); + + harness.teardown().expect("teardown succeeds before join"); + + let node = &harness.nodes()[&id]; + assert_eq!( + node.manager.record().unwrap().stage, + provision::NodeStage::Destroyed + ); + assert!(node.bootstrap.as_ref().unwrap().is_closed()); + assert!(node.bootstrap.as_ref().unwrap().client().closed); + assert_eq!(harness.provider().cli().removed, vec!["container-1"]); +} + +#[test] +fn docker_cluster_harness_cleans_known_lease_when_endpoint_never_appears() { + let mut cli = FakeDockerCli::new(); + cli.endpoints.insert("container-1".into(), None); + let provider = docker::DockerProvider::new(cli); + let datastream = provision::InMemoryBootstrapDatastream::default(); + let mut harness = + docker::DockerClusterHarness::new(provider, FakeSshFactory::default(), datastream); + + let result = harness.start_group(&docker_group_spec(1)); + + assert!(matches!( + result, + Err(docker::DockerClusterError::EndpointUnavailable(id)) + if id == provision::ProviderLeaseId("docker:container-1".into()) + )); + assert_eq!(harness.provider().cli().removed, vec!["container-1"]); + let id = provision::LogicalNodeId("workers-0".into()); + assert_eq!( + harness.nodes()[&id].manager.record().unwrap().stage, + provision::NodeStage::Destroyed + ); +} diff --git a/crates/mvp-system/src/tests/mod.rs b/crates/mvp-system/src/tests/mod.rs index 32fec57..6481daa 100644 --- a/crates/mvp-system/src/tests/mod.rs +++ b/crates/mvp-system/src/tests/mod.rs @@ -1,5 +1,6 @@ mod arena_manager_guarantees; mod device_bridge_guarantees; +mod docker_cluster_provisioning_guarantees; mod edge_establisher_guarantees; mod engine_builder_guarantees; mod gpu_worker_ctl_guarantees; @@ -10,6 +11,7 @@ mod local_mock; mod local_mock_pipeline_integration; mod membership_pool_readiness_guarantees; mod node_boot_lifecycle_guarantees; +mod node_provisioning_guarantees; mod observability_surface_guarantees; mod orchestrator_run_fsm_guarantees; mod orchestrator_token_endpoint_guarantees; diff --git a/crates/mvp-system/src/tests/node_provisioning_guarantees.rs b/crates/mvp-system/src/tests/node_provisioning_guarantees.rs new file mode 100644 index 0000000..cc379de --- /dev/null +++ b/crates/mvp-system/src/tests/node_provisioning_guarantees.rs @@ -0,0 +1,533 @@ +//! In-process contract tests for the MVP node provisioning specification. +//! +//! These tests deliberately use a mock provider and a deterministic bootstrap +//! session. They prove the node-local FSM, bootstrap handoff, datastream log +//! routing, and known-lease teardown without Vast.ai, Docker, or real SSH. + +use mvp_system::node_provisioning as provision; +use mvp_system::node_provisioning::ProviderPlugin; + +fn group_spec(count: u32) -> provision::RunNodeGroupSpec { + provision::RunNodeGroupSpec { + run_id: provision::RunId(42), + group_id: provision::NodeGroupId("workers".into()), + role: provision::RoleId("worker".into()), + count, + provider: provision::ProviderKind::Mock, + shape: provision::DesiredNodeShape { + image: "ghcr.io/acme/mvp-worker:test".into(), + disk_gb: 80, + gpu_name: Some("RTX 4090".into()), + min_gpu_ram_mb: Some(20_000), + min_down_mbps: Some(100.0), + min_up_mbps: Some(20.0), + min_reliability: Some(0.95), + require_verified: false, + provider_labels: [("system".into(), "mvp".into())].into_iter().collect(), + }, + boot: provision::BootSpec { + ssh_user: "root".into(), + verify_commands: vec!["test -x /opt/mvp/swactor".into()], + start_swactor_command: "/opt/mvp/swactor-node --join ${ORCH_ADDR}".into(), + stdout_sources: vec!["/var/log/mvp/stdout.log".into()], + stderr_sources: vec!["/var/log/mvp/stderr.log".into()], + timeout_policy: provision::BootstrapTimeoutPolicy { + ssh_connect_secs: 30, + boot_check_secs: 60, + swactor_join_secs: 90, + }, + }, + swarm_join: provision::SwarmJoinTemplate { + orch_swactor_addr: "quic://orch.example:9443".into(), + join_token_ref: "secret://run-42-token".into(), + }, + } +} + +fn one_logical_node() -> provision::LogicalNodeSpec { + provision::expand_node_group(&group_spec(1)) + .into_iter() + .next() + .expect("fixture expands to one node") +} + +fn start_manager( + spec: provision::LogicalNodeSpec, +) -> (provision::NodeManager, provision::CreateLeaseRequest) { + let mut manager = provision::NodeManager::new(); + let commands = manager + .handle(provision::NodeManagerMsg::Start(spec)) + .expect("start succeeds"); + assert_eq!(commands.len(), 1); + let provision::NodeManagerCommand::CreateLease(request) = commands[0].clone() else { + panic!("start must emit CreateLease, got {:?}", commands[0]); + }; + (manager, request) +} + +fn start_bootstrap_command( + commands: Vec, +) -> provision::BootstrapSessionSpec { + assert_eq!(commands.len(), 1); + let provision::NodeManagerCommand::StartBootstrap(spec) = commands[0].clone() else { + panic!("expected StartBootstrap, got {:?}", commands[0]); + }; + spec +} + +#[test] +fn runplan_group_expands_to_stable_logical_node_specs() { + let nodes = provision::expand_node_group(&group_spec(3)); + + let ids: Vec<_> = nodes + .iter() + .map(|node| node.logical_node_id.0.as_str()) + .collect(); + assert_eq!(ids, vec!["workers-0", "workers-1", "workers-2"]); + + for node in nodes { + assert_eq!(node.run_id, provision::RunId(42)); + assert_eq!(node.group_id, provision::NodeGroupId("workers".into())); + assert_eq!(node.role, provision::RoleId("worker".into())); + assert_eq!( + node.swarm_join.expected_logical_node_id, node.logical_node_id, + "join spec must carry the stable logical node id" + ); + } +} + +#[test] +fn node_manager_start_records_desired_state_and_requests_lease() { + let spec = one_logical_node(); + let (manager, request) = start_manager(spec.clone()); + + let record = manager.record().expect("record exists after start"); + assert_eq!( + record.logical_node_id, + provision::LogicalNodeId("workers-0".into()) + ); + assert_eq!(record.stage, provision::NodeStage::LeaseRequested); + assert!(!record.ready); + assert_eq!(record.desired, spec); + assert_eq!( + request.spec.logical_node_id, + provision::LogicalNodeId("workers-0".into()) + ); +} + +#[test] +fn node_manager_handoff_sets_ready_only_after_bootstrap_closed() { + let spec = one_logical_node(); + let logical_node_id = spec.logical_node_id.clone(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let _bootstrap = start_bootstrap_command(commands); + + let commands = manager + .handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id, + swactor_id: provision::SwactorId("swactor-a".into()), + }) + .expect("swactor join accepted"); + assert!(matches!( + commands.as_slice(), + [provision::NodeManagerCommand::BootstrapConvergenceObserved { .. }] + )); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::SwactorJoined + ); + assert!(!manager.is_ready(), "join alone must not mark readiness"); + + manager + .handle(provision::NodeManagerMsg::BootstrapClosed) + .expect("bootstrap closes after convergence"); + let record = manager.record().expect("record exists"); + assert_eq!(record.stage, provision::NodeStage::Dormant); + assert!(record.ready); + assert!( + record + .swactor + .as_ref() + .expect("swactor facts recorded") + .handed_off_at + .is_some() + ); + assert_eq!(manager.active_bootstrap(), None); +} + +#[test] +fn node_manager_stores_compact_bootstrap_facts_not_log_bodies() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let _bootstrap = start_bootstrap_command(commands); + + manager + .handle(provision::NodeManagerMsg::BootstrapObserved( + provision::BootstrapObservation { + stage: provision::BootstrapStage::StdoutStreaming, + last_stdout_seq: Some(7), + last_stderr_seq: Some(3), + marker: Some("this full line belongs in datastream".into()), + }, + )) + .expect("observation accepted"); + + let facts = manager + .record() + .expect("record exists") + .bootstrap + .as_ref() + .expect("bootstrap facts exist"); + assert_eq!(facts.last_stage, provision::BootstrapStage::StdoutStreaming); + assert_eq!(facts.last_stdout_seq, Some(7)); + assert_eq!(facts.last_stderr_seq, Some(3)); +} + +#[test] +fn delayed_provider_endpoint_starts_bootstrap_after_endpoint_known() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let endpoint = provision::SshEndpoint { + host: "203.0.113.10".into(), + port: 22001, + user: "root".into(), + auth_ref: "mock-key".into(), + }; + provider.queue_create_result(Ok(provision::MockProviderPlugin::result_with_endpoint( + 77, None, + ))); + provider.queue_endpoint_result(Ok(Some(endpoint.clone()))); + + let lease = provider + .create_lease(request) + .expect("queued lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease.clone())) + .expect("lease accepted"); + assert_eq!( + commands, + vec![provision::NodeManagerCommand::LookupEndpoint(lease.lease)] + ); + + let endpoint_result = provider + .lookup_endpoint(manager.record().unwrap().lease.as_ref().unwrap()) + .expect("endpoint lookup succeeds") + .expect("endpoint appears"); + let commands = manager + .handle(provision::NodeManagerMsg::EndpointKnown(endpoint_result)) + .expect("endpoint accepted"); + let bootstrap = start_bootstrap_command(commands); + assert_eq!(bootstrap.ssh, endpoint); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::BootstrapRunning + ); + assert_eq!( + provider.lookup_requests(), + &[provision::ProviderLeaseId("mock:77".into())] + ); +} + +#[test] +fn pre_handoff_failures_are_terminal_and_do_not_replace() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let _ = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease.clone())) + .expect("lease accepted"); + + manager + .handle(provision::NodeManagerMsg::BootstrapFailed( + "boot check failed".into(), + )) + .expect("failure accepted"); + + let record = manager.record().expect("record exists"); + assert_eq!(record.stage, provision::NodeStage::Failed); + assert!(!record.ready); + assert_eq!(record.failed_reason.as_deref(), Some("boot check failed")); + assert_eq!(record.lease.as_ref().expect("lease retained"), &lease.lease); +} + +#[test] +fn swactor_join_for_wrong_logical_node_is_rejected() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let _ = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + + let result = manager.handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id: provision::LogicalNodeId("workers-99".into()), + swactor_id: provision::SwactorId("swactor-wrong".into()), + }); + + assert!(result.is_err()); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::BootstrapRunning + ); + assert!(!manager.is_ready()); +} + +#[test] +fn bootstrap_session_streams_logs_flushes_and_closes_on_convergence() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let bootstrap_spec = start_bootstrap_command(commands); + let mut session = provision::BootstrapSession::new(bootstrap_spec); + let mut datastream = provision::InMemoryBootstrapDatastream::default(); + let script = provision::MockBootstrapScript::successful(vec![ + (provision::BootstrapLogStream::Stdout, "boot entered".into()), + (provision::BootstrapLogStream::Stderr, "warning".into()), + ( + provision::BootstrapLogStream::Stdout, + "swactor starting".into(), + ), + ]); + + let events = session.start(&script, &mut datastream); + + assert_eq!( + session.stage(), + provision::BootstrapStage::WaitingForSwactorJoin + ); + assert_eq!(datastream.records().len(), 3); + assert_eq!(datastream.records()[0].seq, 1); + assert_eq!( + datastream.records()[0].stream, + provision::BootstrapLogStream::Stdout + ); + assert_eq!(datastream.records()[1].seq, 2); + assert_eq!( + datastream.records()[1].stream, + provision::BootstrapLogStream::Stderr + ); + assert_eq!(datastream.records()[2].seq, 3); + assert_eq!(datastream.records()[2].line, "swactor starting"); + assert!(events.iter().any(|event| matches!( + event, + provision::BootstrapSessionEvent::Observed(obs) + if obs.stage == provision::BootstrapStage::SshReady + ))); + + let events = + session.convergence_observed(provision::SwactorId("swactor-a".into()), &mut datastream); + assert!(session.is_closed()); + assert_eq!(datastream.flush_count(), 1); + assert!(matches!( + events.as_slice(), + [_, provision::BootstrapSessionEvent::Closed] + )); +} + +#[test] +fn bootstrap_session_reports_boot_check_failure_without_handoff() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let bootstrap_spec = start_bootstrap_command(commands); + let mut session = provision::BootstrapSession::new(bootstrap_spec); + let mut datastream = provision::InMemoryBootstrapDatastream::default(); + let script = provision::MockBootstrapScript { + ssh_ok: true, + verify_ok: false, + start_ok: true, + records: vec![], + }; + + let events = session.start(&script, &mut datastream); + + assert_eq!(session.stage(), provision::BootstrapStage::BootCheckFailed); + assert!(events.iter().any(|event| matches!( + event, + provision::BootstrapSessionEvent::Failed(reason) if reason == "boot check failed" + ))); + assert!(!session.is_closed()); +} + +#[test] +fn teardown_cancels_active_bootstrap_and_destroys_known_lease_only_once() { + let spec = one_logical_node(); + let (mut manager, request) = start_manager(spec); + let mut provider = provision::MockProviderPlugin::new(); + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease.clone())) + .expect("lease accepted"); + let _bootstrap = start_bootstrap_command(commands); + + let commands = manager + .handle(provision::NodeManagerMsg::Destroy) + .expect("destroy accepted"); + assert!(matches!( + commands.as_slice(), + [ + provision::NodeManagerCommand::CancelBootstrap { .. }, + provision::NodeManagerCommand::DestroyLease(_) + ] + )); + for command in commands { + if let provision::NodeManagerCommand::DestroyLease(handle) = command { + provider.destroy_lease(&handle).expect("destroy succeeds"); + manager + .handle(provision::NodeManagerMsg::LeaseDestroyed) + .expect("destroy recorded"); + } + } + + assert_eq!(provider.destroyed_handles().len(), 1); + assert_eq!(provider.destroyed_handles()[0], lease.lease.destroy_handle); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::Destroyed + ); + let commands = manager + .handle(provision::NodeManagerMsg::Destroy) + .expect("duplicate destroy accepted"); + assert!( + commands.is_empty(), + "destroy must be idempotent after lease destruction" + ); + assert_eq!(provider.destroyed_handles().len(), 1); +} + +#[test] +fn in_process_mock_provision_bootstrap_handoff_and_teardown() { + let specs = provision::expand_node_group(&group_spec(2)); + let mut provider = provision::MockProviderPlugin::new(); + let mut managers = Vec::new(); + let mut datastream = provision::InMemoryBootstrapDatastream::default(); + + for spec in specs { + let logical_node_id = spec.logical_node_id.clone(); + let swactor_id = provision::SwactorId(format!("swactor-{}", logical_node_id.0)); + let (mut manager, request) = start_manager(spec); + + let lease = provider.create_lease(request).expect("mock lease succeeds"); + let commands = manager + .handle(provision::NodeManagerMsg::LeaseCreated(lease)) + .expect("lease accepted"); + let bootstrap_spec = start_bootstrap_command(commands); + let mut session = provision::BootstrapSession::new(bootstrap_spec); + let script = provision::MockBootstrapScript::successful(vec![ + ( + provision::BootstrapLogStream::Stdout, + format!("{} boot entered", logical_node_id.0), + ), + ( + provision::BootstrapLogStream::Stdout, + format!("{} swactor starting", logical_node_id.0), + ), + ]); + + for event in session.start(&script, &mut datastream) { + match event { + provision::BootstrapSessionEvent::Observed(observation) => { + manager + .handle(provision::NodeManagerMsg::BootstrapObserved(observation)) + .expect("bootstrap observation accepted"); + } + provision::BootstrapSessionEvent::Failed(reason) => { + manager + .handle(provision::NodeManagerMsg::BootstrapFailed(reason)) + .expect("bootstrap failure recorded"); + } + provision::BootstrapSessionEvent::Closed => { + manager + .handle(provision::NodeManagerMsg::BootstrapClosed) + .expect("bootstrap closed recorded"); + } + } + } + + let commands = manager + .handle(provision::NodeManagerMsg::SwactorJoined { + logical_node_id: logical_node_id.clone(), + swactor_id: swactor_id.clone(), + }) + .expect("swactor join accepted"); + assert!(matches!( + commands.as_slice(), + [provision::NodeManagerCommand::BootstrapConvergenceObserved { + swactor_id: observed, + .. + }] if *observed == swactor_id + )); + + for event in session.convergence_observed(swactor_id, &mut datastream) { + match event { + provision::BootstrapSessionEvent::Observed(observation) => { + manager + .handle(provision::NodeManagerMsg::BootstrapObserved(observation)) + .expect("convergence observation accepted"); + } + provision::BootstrapSessionEvent::Failed(reason) => { + manager + .handle(provision::NodeManagerMsg::BootstrapFailed(reason)) + .expect("bootstrap failure recorded"); + } + provision::BootstrapSessionEvent::Closed => { + manager + .handle(provision::NodeManagerMsg::BootstrapClosed) + .expect("bootstrap closed recorded"); + } + } + } + + assert!(manager.is_ready()); + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::Dormant + ); + managers.push(manager); + } + + assert_eq!(datastream.records().len(), 4); + assert_eq!(datastream.flush_count(), 2); + assert!(managers.iter().all(provision::NodeManager::is_ready)); + + for manager in &mut managers { + let commands = manager + .handle(provision::NodeManagerMsg::Destroy) + .expect("destroy accepted"); + for command in commands { + if let provision::NodeManagerCommand::DestroyLease(handle) = command { + provider.destroy_lease(&handle).expect("destroy succeeds"); + manager + .handle(provision::NodeManagerMsg::LeaseDestroyed) + .expect("destroy recorded"); + } + } + assert_eq!( + manager.record().expect("record exists").stage, + provision::NodeStage::Destroyed + ); + } + assert_eq!(provider.destroyed_handles().len(), 2); +}