This commit is contained in:
Zachery Aaron Shores-Chmielewski 2026-06-28 11:57:56 +04:00
parent e39dc83a30
commit d0b6555a7c
6 changed files with 2656 additions and 0 deletions

View file

@ -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<String>) -> 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<String, String>,
pub env: BTreeMap<String, String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct DockerRunResult {
pub container_id: String,
}
pub trait DockerCli {
fn run_container(
&mut self,
request: DockerRunRequest,
) -> Result<DockerRunResult, DockerCliError>;
fn inspect_ssh_endpoint(
&mut self,
container_id: &str,
) -> Result<Option<SshEndpoint>, DockerCliError>;
fn remove_force(&mut self, container_id: &str) -> Result<(), DockerCliError>;
}
#[derive(Clone, Debug)]
pub struct DockerProvider<C> {
cli: C,
}
impl<C> DockerProvider<C> {
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<C: DockerCli> DockerProvider<C> {
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<SshEndpoint>,
) -> 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<C: DockerCli> ProviderPlugin for DockerProvider<C> {
fn create_lease(
&mut self,
request: CreateLeaseRequest,
) -> Result<CreateLeaseResult, ProviderError> {
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<Option<SshEndpoint>, 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<String>) -> 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<Vec<(BootstrapLogStream, String)>, 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<C> {
spec: BootstrapSessionSpec,
client: C,
stage: BootstrapStage,
next_seq: u64,
last_stdout_seq: Option<u64>,
last_stderr_seq: Option<u64>,
closed: bool,
}
impl<C> SshBootstrapSession<C> {
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<C: BootstrapSshClient> SshBootstrapSession<C> {
pub fn start(&mut self, sink: &mut dyn BootstrapDatastreamSink) -> Vec<BootstrapSessionEvent> {
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<BootstrapSessionEvent> {
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<BootstrapSessionEvent> {
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<LogicalNodeId>,
}
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<Vec<NodeManagerCommand>, 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<C> {
pub manager: NodeManager,
pub bootstrap: Option<SshBootstrapSession<C>>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum DockerClusterError {
Provider(String),
Node(String),
UnknownNode(LogicalNodeId),
EndpointUnavailable(ProviderLeaseId),
MissingBootstrap(LogicalNodeId),
}
pub struct DockerClusterHarness<D, F, S>
where
D: DockerCli,
F: SshBootstrapClientFactory,
S: BootstrapDatastreamSink,
{
provider: DockerProvider<D>,
client_factory: F,
datastream: S,
router: SwactorJoinRouter,
nodes: BTreeMap<LogicalNodeId, ManagedDockerNode<F::Client>>,
teardown_complete: bool,
}
impl<D, F, S> DockerClusterHarness<D, F, S>
where
D: DockerCli,
F: SshBootstrapClientFactory,
S: BootstrapDatastreamSink,
{
pub fn new(provider: DockerProvider<D>, 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<D> {
&self.provider
}
pub fn datastream(&self) -> &S {
&self.datastream
}
pub fn nodes(&self) -> &BTreeMap<LogicalNodeId, ManagedDockerNode<F::Client>> {
&self.nodes
}
pub fn records(&self) -> Vec<NodeRecord> {
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<NodeManagerCommand>,
) -> Result<Option<SshBootstrapSession<F::Client>>, 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<BootstrapSessionEvent>,
) -> 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<D, F, S> Drop for DockerClusterHarness<D, F, S>
where
D: DockerCli,
F: SshBootstrapClientFactory,
S: BootstrapDatastreamSink,
{
fn drop(&mut self) {
if !self.teardown_complete {
let _ = self.teardown();
}
}
}

View file

@ -7,6 +7,7 @@ pub mod arena_manager;
pub mod dashboard_view; pub mod dashboard_view;
pub mod device_bridge; pub mod device_bridge;
pub mod distribution_stack; pub mod distribution_stack;
pub mod docker_cluster_provisioning;
pub mod driver_pumps; pub mod driver_pumps;
pub mod edge_establisher; pub mod edge_establisher;
pub mod engine_builder; pub mod engine_builder;
@ -16,6 +17,7 @@ pub mod gpu_worker_ingress_parser;
pub mod gpu_worker_process_adapter; pub mod gpu_worker_process_adapter;
pub mod membership_pool_readiness; pub mod membership_pool_readiness;
pub mod node_boot_lifecycle; pub mod node_boot_lifecycle;
pub mod node_provisioning;
pub mod observability_surface; pub mod observability_surface;
pub mod orchestrator_run_fsm; pub mod orchestrator_run_fsm;
pub mod orchestrator_token_endpoint; pub mod orchestrator_token_endpoint;

View file

@ -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<String>,
pub min_gpu_ram_mb: Option<u64>,
pub min_down_mbps: Option<f64>,
pub min_up_mbps: Option<f64>,
pub min_reliability: Option<f64>,
pub require_verified: bool,
pub provider_labels: BTreeMap<String, String>,
}
#[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<String>,
pub start_swactor_command: String,
pub stdout_sources: Vec<String>,
pub stderr_sources: Vec<String>,
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<LogicalNodeSpec> {
(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<String>,
pub destroy_handle: DestroyHandle,
pub provider_metadata: BTreeMap<String, String>,
}
#[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<u64>,
pub last_stderr_seq: Option<u64>,
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<SystemTime>,
}
#[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<LeaseFacts>,
pub connection: Option<SshEndpoint>,
pub bootstrap: Option<BootstrapFacts>,
pub swactor: Option<SwactorFacts>,
pub failed_reason: Option<String>,
pub destroyed_at: Option<SystemTime>,
}
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<SshEndpoint>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ProviderError {
pub reason: String,
}
impl ProviderError {
pub fn new(reason: impl Into<String>) -> Self {
Self {
reason: reason.into(),
}
}
}
pub trait ProviderPlugin {
fn create_lease(
&mut self,
request: CreateLeaseRequest,
) -> Result<CreateLeaseResult, ProviderError>;
fn lookup_endpoint(&mut self, lease: &LeaseFacts)
-> Result<Option<SshEndpoint>, 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<u64>,
pub last_stderr_seq: Option<u64>,
pub marker: Option<String>,
}
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<String>) -> Self {
Self {
reason: reason.into(),
}
}
}
#[derive(Clone, Debug)]
pub struct NodeManager {
record: Option<NodeRecord>,
active_bootstrap: Option<BootstrapSessionId>,
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<BootstrapSessionId> {
self.active_bootstrap
}
pub fn handle(
&mut self,
msg: NodeManagerMsg,
) -> Result<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, 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<Vec<NodeManagerCommand>, NodeManagerError> {
self.fail(format!("destroy: {reason}"))
}
fn fail(&mut self, reason: String) -> Result<Vec<NodeManagerCommand>, 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<BootstrapLogRecord>,
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<u64>,
last_stderr_seq: Option<u64>,
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<BootstrapSessionEvent> {
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<BootstrapSessionEvent> {
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<BootstrapSessionEvent> {
self.stage = BootstrapStage::JoinTimeout;
vec![BootstrapSessionEvent::Failed("join timeout".into())]
}
pub fn cancel(&mut self) -> Vec<BootstrapSessionEvent> {
self.stage = BootstrapStage::Cancelled;
self.closed = true;
vec![BootstrapSessionEvent::Closed]
}
}
#[derive(Default, Debug, Clone)]
pub struct MockProviderPlugin {
next_contract_id: u64,
create_results: VecDeque<Result<CreateLeaseResult, ProviderError>>,
endpoint_results: VecDeque<Result<Option<SshEndpoint>, ProviderError>>,
create_requests: Vec<CreateLeaseRequest>,
lookup_requests: Vec<ProviderLeaseId>,
destroyed_handles: Vec<DestroyHandle>,
destroy_failures: BTreeMap<ProviderLeaseId, String>,
}
impl MockProviderPlugin {
pub fn new() -> Self {
Self {
next_contract_id: 1,
..Self::default()
}
}
pub fn queue_create_result(&mut self, result: Result<CreateLeaseResult, ProviderError>) {
self.create_results.push_back(result);
}
pub fn queue_endpoint_result(&mut self, result: Result<Option<SshEndpoint>, ProviderError>) {
self.endpoint_results.push_back(result);
}
pub fn fail_destroy(&mut self, lease_id: ProviderLeaseId, reason: impl Into<String>) {
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<SshEndpoint>) -> 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<CreateLeaseResult, ProviderError> {
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<Option<SshEndpoint>, 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(())
}
}

View file

@ -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<docker::DockerRunRequest>,
inspect_requests: Vec<String>,
removed: Vec<String>,
endpoints: BTreeMap<String, Option<provision::SshEndpoint>>,
}
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<docker::DockerRunResult, docker::DockerCliError> {
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<Option<provision::SshEndpoint>, 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<provision::SshEndpoint>,
probed: bool,
stdout_sources: Vec<String>,
stderr_sources: Vec<String>,
verify_commands: Vec<String>,
started_command: Option<String>,
started_join: Option<provision::SwarmJoinSpec>,
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<Vec<(provision::BootstrapLogStream, String)>, 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<provision::LogicalNodeId>,
}
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
);
}

View file

@ -1,5 +1,6 @@
mod arena_manager_guarantees; mod arena_manager_guarantees;
mod device_bridge_guarantees; mod device_bridge_guarantees;
mod docker_cluster_provisioning_guarantees;
mod edge_establisher_guarantees; mod edge_establisher_guarantees;
mod engine_builder_guarantees; mod engine_builder_guarantees;
mod gpu_worker_ctl_guarantees; mod gpu_worker_ctl_guarantees;
@ -10,6 +11,7 @@ mod local_mock;
mod local_mock_pipeline_integration; mod local_mock_pipeline_integration;
mod membership_pool_readiness_guarantees; mod membership_pool_readiness_guarantees;
mod node_boot_lifecycle_guarantees; mod node_boot_lifecycle_guarantees;
mod node_provisioning_guarantees;
mod observability_surface_guarantees; mod observability_surface_guarantees;
mod orchestrator_run_fsm_guarantees; mod orchestrator_run_fsm_guarantees;
mod orchestrator_token_endpoint_guarantees; mod orchestrator_token_endpoint_guarantees;

View file

@ -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::NodeManagerCommand>,
) -> 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);
}