use std::{ sync::{ Arc, atomic::{AtomicUsize, Ordering}, }, time::Duration, }; use testing_framework_core::scenario::{ ClusterControlProfile, ClusterControlRequest, ClusterProvisioner, ClusterRequest, ClusterSource, ClusterStartMode, ClusterUnit, DeploymentPolicy, DynError, HttpReadinessRequirement, NodeClients, ProvisionedCluster, RetryPolicy, }; use tokio_retry::{ RetryIf, strategy::{ExponentialBackoff, jitter}, }; use tracing::{debug, info, warn}; use crate::{ LocalCluster, LocalDeployerEnv, NodeManager, env::{Node, wait_local_readiness}, external::build_external_client, keep_tempdir_from_env, }; const READINESS_ATTEMPTS: usize = 3; const READINESS_BACKOFF_BASE_MS: u64 = 250; const READINESS_BACKOFF_MAX_SECS: u64 = 2; #[derive(Clone, Copy)] struct RetryExecutionConfig { max_attempts: usize, keep_tempdir: bool, readiness_enabled: bool, readiness_requirement: HttpReadinessRequirement, } #[derive(Debug, thiserror::Error)] enum RetryAttemptError { #[error("failed to spawn local topology: {source}")] Spawn { #[source] source: DynError, }, #[error("readiness probe failed: {source}")] Readiness { #[source] source: DynError, }, } #[derive(Debug, thiserror::Error)] pub enum LocalClusterProvisionerError { #[error("failed to spawn local topology: {source}")] Spawn { #[source] source: DynError, }, #[error("readiness probe failed: {source}")] Readiness { #[source] source: DynError, }, #[error("source orchestration failed: {source}")] Source { #[source] source: DynError, }, #[error("local cluster provisioner does not support attached clusters")] AttachedUnsupported, } impl From for LocalClusterProvisionerError { fn from(value: RetryAttemptError) -> Self { match value { RetryAttemptError::Spawn { source } => Self::Spawn { source }, RetryAttemptError::Readiness { source } => Self::Readiness { source }, } } } /// Local implementation of the shared cluster provisioning operation. #[derive(Clone, Copy, Debug, Default)] pub struct LocalClusterProvisioner; /// Concrete local handle plus the backend-independent runtime surfaces. pub type ProvisionedLocalCluster = ProvisionedCluster>; impl LocalClusterProvisioner { #[must_use] pub fn manual_cluster( &self, deployment: E::Deployment, ) -> LocalCluster { LocalCluster::empty(deployment, keep_tempdir_from_env()) } pub async fn provision( &self, request: ClusterRequest, membership_check: bool, ) -> Result, LocalClusterProvisionerError> { match request.source().clone() { ClusterSource::Managed { deployment, external, } => { let keep_tempdir = request.policy().cleanup_policy.preserve_artifacts || keep_tempdir_from_env(); let cluster = LocalCluster::empty(deployment.clone(), keep_tempdir); if request.start_mode() == ClusterStartMode::Eager { let nodes = spawn_with_readiness_retry::( &deployment, membership_check, request.policy(), ) .await?; cluster.initialize_with_nodes(nodes); } cluster .add_external_sources(external) .map_err(|source| LocalClusterProvisionerError::Source { source })?; let profile = if request.start_mode() == ClusterStartMode::OnDemand { ClusterControlProfile::ManualControlled } else { ClusterControlProfile::FrameworkManaged }; let mut unit = ClusterUnit::new(Some(deployment), cluster.node_clients(), profile) .with_cluster_wait(Arc::new(cluster.clone())) .with_cleanup(cluster.cleanup_guard()); if request.control() == ClusterControlRequest::Full { unit = unit.with_node_control(Arc::new(cluster.clone())); } Ok(ProvisionedCluster::new(Some(cluster), unit)) } ClusterSource::External { nodes } => { let clients = NodeClients::default(); for source in nodes { let client = E::external_node_client(&source) .or_else(|_| build_external_client::(&source)) .map_err(|source| LocalClusterProvisionerError::Source { source })?; clients.add_node(client); } Ok(ProvisionedCluster::new( None, ClusterUnit::new(None, clients, ClusterControlProfile::ExternalUncontrolled), )) } ClusterSource::Attached { .. } => { Err(LocalClusterProvisionerError::AttachedUnsupported) } } } } #[async_trait::async_trait] impl ClusterProvisioner for LocalClusterProvisioner { async fn provision_cluster( &self, request: ClusterRequest, ) -> Result, DynError> { let (_, unit) = self .provision(request, true) .await .map_err(DynError::from)? .into_parts(); Ok(unit) } } async fn spawn_with_readiness_retry( deployment: &E::Deployment, membership_check: bool, policy: DeploymentPolicy, ) -> Result>, LocalClusterProvisionerError> { let (retry_policy, execution) = build_retry_execution_config(policy, membership_check); let attempts = Arc::new(AtomicUsize::new(0)); let strategy = ExponentialBackoff::from_millis(retry_policy.base_delay.as_millis() as u64) .max_delay(retry_policy.max_delay) .map(jitter) .take(execution.max_attempts.saturating_sub(1)); let operation = { let attempts = Arc::clone(&attempts); move || { let attempts = Arc::clone(&attempts); async move { run_retry_attempt::(deployment, execution, attempts).await } } }; let should_retry = { let attempts = Arc::clone(&attempts); move |error: &RetryAttemptError| { let attempt = attempts.load(Ordering::Relaxed); if attempt < execution.max_attempts { warn!( attempt, max_attempts = execution.max_attempts, error = %error, "local spawn/readiness failed; retrying with backoff" ); true } else { false } } }; RetryIf::spawn(strategy, operation, should_retry) .await .map_err(Into::into) } fn build_retry_execution_config( policy: DeploymentPolicy, membership_check: bool, ) -> (RetryPolicy, RetryExecutionConfig) { let retry_policy = policy.retry_policy.unwrap_or_else(|| { RetryPolicy::new( READINESS_ATTEMPTS, Duration::from_millis(READINESS_BACKOFF_BASE_MS), Duration::from_secs(READINESS_BACKOFF_MAX_SECS), ) }); let execution = RetryExecutionConfig { max_attempts: retry_policy.max_attempts.max(1), keep_tempdir: policy.cleanup_policy.preserve_artifacts || keep_tempdir_from_env(), readiness_enabled: policy.readiness_enabled && membership_check, readiness_requirement: policy.readiness_requirement, }; (retry_policy, execution) } async fn run_retry_attempt( deployment: &E::Deployment, execution: RetryExecutionConfig, attempts: Arc, ) -> Result>, RetryAttemptError> { let attempt = attempts.fetch_add(1, Ordering::Relaxed) + 1; let nodes = NodeManager::::spawn_initial_nodes(deployment, execution.keep_tempdir) .await .map_err(|source| RetryAttemptError::Spawn { source: source.into(), })?; if !execution.readiness_enabled { info!("skipping local readiness checks"); return Ok(nodes); } match wait_local_readiness::(&nodes, execution.readiness_requirement).await { Ok(()) => { info!(attempt, "local nodes are ready"); Ok(nodes) } Err(source) => { let source: DynError = source.into(); debug!(attempt, error = ?source, "local readiness failed"); drop(nodes); Err(RetryAttemptError::Readiness { source }) } } } #[cfg(test)] mod tests { use std::{collections::HashMap, path::Path}; use testing_framework_core::{ scenario::{ Application, ClusterControlProfile, ClusterControlRequest, ClusterRequest, ClusterStartMode, DynError, ExistingCluster, ExternalNodeSource, StartNodeOptions, }, topology::DeploymentDescriptor, }; use super::LocalClusterProvisioner; use crate::{ BuiltNodeConfig, LaunchSpec, LocalDeployerEnv, NodeConfigEntry, NodeEndpoints, ProcessSpawnError, }; #[derive(Clone)] struct EmptyDeployment; impl DeploymentDescriptor for EmptyDeployment { fn node_count(&self) -> usize { 0 } } #[derive(Clone)] struct EmptyConfig; struct TestEnv; #[async_trait::async_trait] impl Application for TestEnv { type Deployment = EmptyDeployment; type NodeClient = String; type NodeConfig = EmptyConfig; fn external_node_client(source: &ExternalNodeSource) -> Result { Ok(source.endpoint().to_owned()) } } #[async_trait::async_trait] impl LocalDeployerEnv for TestEnv { fn build_node_config( _topology: &Self::Deployment, _index: usize, _peer_ports_by_name: &HashMap, _options: &StartNodeOptions, _peer_ports: &[u16], ) -> Result, DynError> { unreachable!("empty deployment never builds node configs") } fn build_initial_node_configs( _topology: &Self::Deployment, ) -> Result>, ProcessSpawnError> { Ok(Vec::new()) } fn build_launch_spec( _config: &EmptyConfig, _dir: &Path, _label: &str, ) -> Result { unreachable!("empty deployment never launches nodes") } fn node_endpoints(_config: &EmptyConfig) -> Result { unreachable!("empty deployment has no endpoints") } fn node_client(_endpoints: &NodeEndpoints) -> Result { unreachable!("empty deployment has no clients") } } #[tokio::test] async fn on_demand_managed_unit_exposes_shared_control_and_cleanup() { let request = ClusterRequest::managed(EmptyDeployment) .with_start_mode(ClusterStartMode::OnDemand) .with_control(ClusterControlRequest::Full); let provisioned = LocalClusterProvisioner .provision::(request, true) .await .expect("manual unit should provision"); let (cluster, mut unit) = provisioned.into_parts(); let cluster = cluster.expect("managed unit should expose its concrete local handle"); assert_eq!( unit.control_profile(), ClusterControlProfile::ManualControlled ); assert!(unit.node_control().is_some()); assert!(unit.cluster_wait().is_some()); unit.take_cleanup() .expect("managed unit should own cleanup") .cleanup(); assert!(cluster.stop_all().is_err(), "cleanup must lock out clones"); } #[tokio::test] async fn external_unit_is_borrowed_and_uncontrolled() { let source = ExternalNodeSource::new("external-0".into(), "http://node-0".into()); let provisioned = LocalClusterProvisioner .provision::(ClusterRequest::external(vec![source]), true) .await .expect("external unit should resolve"); let (cluster, mut unit) = provisioned.into_parts(); assert!(cluster.is_none()); assert_eq!(unit.node_clients().snapshot(), vec!["http://node-0"]); assert_eq!( unit.control_profile(), ClusterControlProfile::ExternalUncontrolled ); assert!(unit.node_control().is_none()); assert!(unit.cluster_wait().is_none()); assert!(unit.take_cleanup().is_none()); } #[tokio::test] async fn attached_unit_reports_backend_support_gap() { let request = ClusterRequest::::attached(ExistingCluster::for_compose_project( "existing".into(), )); let error = LocalClusterProvisioner .provision(request, true) .await .err() .expect("local attached provisioning should be rejected"); assert_eq!( error.to_string(), "local cluster provisioner does not support attached clusters" ); } }