logos-blockchain-testing/testing-framework/runners/compose/src/deployer/orchestrator.rs

use std::sync::Arc;

use testing_framework_core::scenario::{
    NodeControlHandle, RequiresNodeControl, RunContext, Runner, Scenario,
};
use tracing::info;

use super::{
    ComposeDeployer,
    clients::ClientBuilder,
    make_cleanup_guard,
    ports::PortManager,
    readiness::ReadinessChecker,
    setup::{DeploymentContext, DeploymentSetup},
};
use crate::{
    control::ComposeNodeControl, errors::ComposeRunnerError, ports::compose_runner_host,
    readiness::metrics_handle_from_port,
};

pub struct DeploymentOrchestrator {
    deployer: ComposeDeployer,
}

impl DeploymentOrchestrator {
    pub const fn new(deployer: ComposeDeployer) -> Self {
        Self { deployer }
    }

    pub async fn deploy<Caps>(
        &self,
        scenario: &Scenario<Caps>,
    ) -> Result<Runner, ComposeRunnerError>
    where
        Caps: RequiresNodeControl + Send + Sync,
    {
        let setup = DeploymentSetup::new(scenario.topology());
        setup.validate_environment().await?;

        let DeploymentContext {
            mut environment,
            descriptors,
        } = setup.prepare_workspace().await?;

        let host_ports = PortManager::prepare(&mut environment, &descriptors).await?;

        if self.deployer.readiness_checks {
            ReadinessChecker::wait_all(&descriptors, &host_ports, &mut environment).await?;
        } else {
            info!("readiness checks disabled; giving the stack a short grace period");
            crate::readiness::maybe_sleep_for_disabled_readiness(false).await;
        }

        let host = compose_runner_host();
        let client_builder = ClientBuilder::new();
        let node_clients = client_builder
            .build_node_clients(&descriptors, &host_ports, &host, &mut environment)
            .await?;
        let telemetry = metrics_handle_from_port(environment.prometheus_port(), &host)?;
        let node_control = self.maybe_node_control::<Caps>(&environment);

        let (block_feed, block_feed_guard) = client_builder
            .start_block_feed(&node_clients, &mut environment)
            .await?;
        let cleanup_guard = make_cleanup_guard(environment.into_cleanup(), block_feed_guard);

        let context = RunContext::new(
            descriptors,
            None,
            node_clients,
            scenario.duration(),
            telemetry,
            block_feed,
            node_control,
        );

        Ok(Runner::new(context, Some(cleanup_guard)))
    }

    fn maybe_node_control<Caps>(
        &self,
        environment: &crate::environment::StackEnvironment,
    ) -> Option<Arc<dyn NodeControlHandle>>
    where
        Caps: RequiresNodeControl + Send + Sync,
    {
        Caps::REQUIRED.then(|| {
            Arc::new(ComposeNodeControl {
                compose_file: environment.compose_path().to_path_buf(),
                project_name: environment.project_name().to_owned(),
            }) as Arc<dyn NodeControlHandle>
        })
    }
}