zombienet_orchestrator/

network_spec.rs

use std::{
    collections::{hash_map::Entry, HashMap},
    sync::Arc,
};

use configuration::{CustomProcess, GlobalSettings, HrmpChannelConfig, NetworkConfig};
use futures::future::try_join_all;
use provider::{DynNamespace, ProviderError, ProviderNamespace};
use serde::{Deserialize, Serialize};
use support::{
    constants::{RELAY_NOT_NONE, THIS_IS_A_BUG},
    fs::FileSystem,
};
use tracing::{debug, trace};

use crate::{
    errors::{merge_errs, OrchestratorError},
    network_spec::jamchain::JamchainSpec,
    ScopedFilesystem,
};

pub mod jamchain;
pub mod jamnode;
pub mod node;
pub mod parachain;
pub mod relaychain;

use self::{node::NodeSpec, parachain::ParachainSpec, relaychain::RelaychainSpec};

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NetworkSpec {
    /// Relaychain configuration, `None` for a JAM network.
    pub(crate) relaychain: Option<RelaychainSpec>,

    /// Jamchain configuration, takes the place of the relaychain
    pub(crate) jamchain: Option<JamchainSpec>,

    /// Parachains configurations.
    pub(crate) parachains: Vec<ParachainSpec>,

    /// HRMP channels configurations.
    pub(crate) hrmp_channels: Vec<HrmpChannelConfig>,

    /// Global settings
    pub(crate) global_settings: GlobalSettings,

    /// Custom processes
    #[serde(default)]
    pub(crate) custom_processes: Vec<CustomProcess>,
}

impl NetworkSpec {
    pub async fn from_config(
        network_config: &NetworkConfig,
    ) -> Result<NetworkSpec, OrchestratorError> {
        let mut errs = vec![];
        let relaychain = network_config
            .try_relaychain()
            .map(RelaychainSpec::from_config)
            .transpose()?;
        let jamchain = if let Some(jamchain_config) = network_config.jamchain() {
            Some(JamchainSpec::from_config(jamchain_config)?)
        } else {
            None
        };

        // The chain the parachains are anchored to, either the relaychain or the JAM chain.
        let root_chain = relaychain
            .as_ref()
            .map(|relaychain| relaychain.chain.clone())
            .or_else(|| jamchain.as_ref().map(|jamchain| jamchain.id.clone()));

        let mut parachains = vec![];

        if let Some(root_chain) = root_chain {
            // TODO: move to `fold` or map+fold
            for para_config in network_config.parachains() {
                match ParachainSpec::from_config(para_config, root_chain.clone()) {
                    Ok(para) => parachains.push(para),
                    Err(err) => errs.push(err),
                }
            }
        }

        if errs.is_empty() {
            Ok(NetworkSpec {
                relaychain,
                jamchain,
                parachains,
                hrmp_channels: network_config
                    .hrmp_channels()
                    .into_iter()
                    .cloned()
                    .collect(),
                global_settings: network_config.global_settings().clone(),
                custom_processes: network_config
                    .custom_processes()
                    .into_iter()
                    .cloned()
                    .collect(),
            })
        } else {
            Err(merge_errs(&errs))
        }
    }

    pub async fn populate_nodes_available_args(
        &mut self,
        ns: Arc<dyn ProviderNamespace + Send + Sync>,
    ) -> Result<(), OrchestratorError> {
        let network_nodes = self.collect_network_nodes();

        let mut image_command_to_nodes_mapping =
            Self::create_image_command_to_nodes_mapping(network_nodes);

        let available_args_outputs =
            Self::retrieve_all_nodes_available_args_output(ns, &image_command_to_nodes_mapping)
                .await?;

        Self::update_nodes_available_args_output(
            &mut image_command_to_nodes_mapping,
            available_args_outputs,
        );

        Ok(())
    }

    //
    pub async fn node_available_args_output(
        &self,
        node_spec: &NodeSpec,
        ns: Arc<dyn ProviderNamespace + Send + Sync>,
    ) -> Result<String, ProviderError> {
        // try to find a node that use the same combination of image/cmd
        let cmp_fn = |ad_hoc: &&NodeSpec| -> bool {
            ad_hoc.image == node_spec.image && ad_hoc.command == node_spec.command
        };

        // check if we already had computed the args output for this cmd/[image]
        let node = self
            .relaychain
            .iter()
            .flat_map(|relaychain| relaychain.nodes.iter())
            .find(cmp_fn);
        let node = if let Some(node) = node {
            Some(node)
        } else {
            let node = self
                .parachains
                .iter()
                .find_map(|para| para.collators.iter().find(cmp_fn));

            node
        };

        let output = if let Some(node) = node {
            node.available_args_output.clone().expect(&format!(
                "args_output should be set for running nodes {THIS_IS_A_BUG}"
            ))
        } else {
            // we need to compute the args output
            let image = node_spec
                .image
                .as_ref()
                .map(|image| image.as_str().to_string());
            let command = node_spec.command.as_str().to_string();

            ns.get_node_available_args((command, image)).await?
        };

        Ok(output)
    }

    pub fn relaychain(&self) -> &RelaychainSpec {
        self.relaychain
            .as_ref()
            .expect(&format!("{RELAY_NOT_NONE}, {THIS_IS_A_BUG}"))
    }

    pub fn relaychain_mut(&mut self) -> &mut RelaychainSpec {
        self.relaychain
            .as_mut()
            .expect(&format!("{RELAY_NOT_NONE}, {THIS_IS_A_BUG}"))
    }

    /// The relaychain spec, `None` for a JAM network.
    pub fn try_relaychain(&self) -> Option<&RelaychainSpec> {
        self.relaychain.as_ref()
    }

    pub fn parachains_iter(&self) -> impl Iterator<Item = &ParachainSpec> {
        self.parachains.iter()
    }

    pub fn parachains_iter_mut(&mut self) -> impl Iterator<Item = &mut ParachainSpec> {
        self.parachains.iter_mut()
    }

    pub fn set_global_settings(&mut self, global_settings: GlobalSettings) {
        self.global_settings = global_settings;
    }

    pub async fn build_parachain_artifacts<'a, T: FileSystem>(
        &mut self,
        ns: DynNamespace,
        scoped_fs: &ScopedFilesystem<'a, T>,
        relaychain_id: &str,
        base_dir_exists: bool,
    ) -> Result<(), anyhow::Error> {
        for para in self.parachains.iter_mut() {
            let chain_spec_raw_path = para
                .build_chain_spec(
                    relaychain_id,
                    &ns,
                    scoped_fs,
                    para.post_process_script.clone().as_deref(),
                )
                .await?;

            trace!("creating dirs for {}", &para.unique_id);
            if base_dir_exists {
                scoped_fs.create_dir_all(&para.unique_id).await?;
            } else {
                scoped_fs.create_dir(&para.unique_id).await?;
            };
            trace!("created dirs for {}", &para.unique_id);

            // create wasm/state
            para.genesis_state
                .build(
                    chain_spec_raw_path.clone(),
                    format!("{}/genesis-state", para.unique_id),
                    &ns,
                    scoped_fs,
                    None,
                )
                .await?;
            debug!("parachain genesis state built!");
            para.genesis_wasm
                .build(
                    chain_spec_raw_path,
                    format!("{}/genesis-wasm", para.unique_id),
                    &ns,
                    scoped_fs,
                    None,
                )
                .await?;
            debug!("parachain genesis wasm built!");
        }

        Ok(())
    }

    // collect mutable references to all nodes from relaychain and parachains
    fn collect_network_nodes(&mut self) -> Vec<&mut NodeSpec> {
        vec![
            self.relaychain
                .iter_mut()
                .flat_map(|relaychain| relaychain.nodes.iter_mut())
                .collect::<Vec<_>>(),
            self.parachains
                .iter_mut()
                .flat_map(|para| para.collators.iter_mut())
                .collect(),
        ]
        .into_iter()
        .flatten()
        .collect::<Vec<_>>()
    }

    // initialize the mapping of all possible node image/commands to corresponding nodes
    fn create_image_command_to_nodes_mapping(
        network_nodes: Vec<&mut NodeSpec>,
    ) -> HashMap<(Option<String>, String), Vec<&mut NodeSpec>> {
        network_nodes.into_iter().fold(
            HashMap::new(),
            |mut acc: HashMap<(Option<String>, String), Vec<&mut node::NodeSpec>>, node| {
                // build mapping key using image and command if image is present or command only
                let key = node
                    .image
                    .as_ref()
                    .map(|image| {
                        (
                            Some(image.as_str().to_string()),
                            node.command.as_str().to_string(),
                        )
                    })
                    .unwrap_or_else(|| (None, node.command.as_str().to_string()));

                // append the node to the vector of nodes for this image/command tuple
                if let Entry::Vacant(entry) = acc.entry(key.clone()) {
                    entry.insert(vec![node]);
                } else {
                    acc.get_mut(&key).unwrap().push(node);
                }

                acc
            },
        )
    }

    async fn retrieve_all_nodes_available_args_output(
        ns: Arc<dyn ProviderNamespace + Send + Sync>,
        image_command_to_nodes_mapping: &HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
    ) -> Result<Vec<(Option<String>, String, String)>, OrchestratorError> {
        try_join_all(
            image_command_to_nodes_mapping
                .keys()
                .map(|(image, command)| async {
                    let image = image.clone();
                    let command = command.clone();
                    // get node available args output from image/command
                    let available_args = ns
                        .get_node_available_args((command.clone(), image.clone()))
                        .await?;
                    debug!(
                        "retrieved available args for image: {:?}, command: {}",
                        image, command
                    );

                    // map the result to include image and command
                    Ok::<_, OrchestratorError>((image, command, available_args))
                })
                .collect::<Vec<_>>(),
        )
        .await
    }

    fn update_nodes_available_args_output(
        image_command_to_nodes_mapping: &mut HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
        available_args_outputs: Vec<(Option<String>, String, String)>,
    ) {
        for (image, command, available_args_output) in available_args_outputs {
            let nodes = image_command_to_nodes_mapping
                .get_mut(&(image, command))
                .expect(&format!(
                    "node image/command key should exist {THIS_IS_A_BUG}"
                ));

            for node in nodes {
                node.available_args_output = Some(available_args_output.clone());
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::generators::generate_node_port;

    #[tokio::test]
    async fn small_network_config_get_spec() {
        use configuration::NetworkConfigBuilder;

        use super::*;

        let config = NetworkConfigBuilder::new()
            .with_relaychain(|r| {
                r.with_chain("rococo-local")
                    .with_default_command("polkadot")
                    .with_validator(|node| node.with_name("alice"))
                    .with_fullnode(|node| node.with_name("bob").with_command("polkadot1"))
            })
            .with_parachain(|p| {
                p.with_id(100)
                    .with_default_command("adder-collator")
                    .with_collator(|c| c.with_name("collator1"))
            })
            .with_custom_process(|c| c.with_name("eth-rpc").with_command("command"))
            .build()
            .unwrap();

        let network_spec = NetworkSpec::from_config(&config).await.unwrap();
        let alice = network_spec.relaychain().nodes.first().unwrap();
        let bob = network_spec.relaychain().nodes.get(1).unwrap();
        assert_eq!(alice.command.as_str(), "polkadot");
        assert_eq!(bob.command.as_str(), "polkadot1");
        assert!(alice.is_validator);
        assert!(!bob.is_validator);
        assert_eq!(network_spec.custom_processes.len(), 1);
        let first_custom_process = network_spec.custom_processes.first().unwrap();
        assert_eq!(first_custom_process.name(), "eth-rpc");

        // paras
        assert_eq!(network_spec.parachains.len(), 1);
        let para_100 = network_spec.parachains.first().unwrap();
        assert_eq!(para_100.id, 100);
    }

    #[tokio::test]
    async fn jam_network_with_parachain_get_spec() {
        use configuration::NetworkConfig;

        use super::*;

        let config = NetworkConfig::load_from_toml_string(
            r#"
[jamchain]
id = "dev"
default_command = "polkajam"

[[jamchain.nodes]]
name = "jam-or"
mode = "ordinary"

[[parachains]]
id = 100
default_command = "polkadot-omni-node"
default_args = ["--jam-rpc-url http://{{ZOMBIE:jam-or:rpc_uri}}"]

[[parachains.collators]]
name = "collator"
"#,
        )
        .unwrap();

        let network_spec = NetworkSpec::from_config(&config).await.unwrap();

        assert!(network_spec.try_relaychain().is_none());
        assert_eq!(network_spec.jamchain.as_ref().unwrap().nodes.len(), 1);

        // the parachain is anchored to the JAM chain
        assert_eq!(network_spec.parachains.len(), 1);
        let para = network_spec.parachains.first().unwrap();
        assert_eq!(para.id, 100);
        let collator = para.collators.first().unwrap();
        assert_eq!(collator.command.as_str(), "polkadot-omni-node");
    }

    #[tokio::test]
    async fn used_port_in_node_reports_err() {
        use configuration::NetworkConfigBuilder;

        use super::*;
        let used_port = generate_node_port(None).unwrap();

        let config = NetworkConfigBuilder::new()
            .with_relaychain(|r| {
                r.with_chain("rococo-local")
                    .with_default_command("polkadot")
                    .with_validator(|node| node.with_name("alice"))
                    .with_fullnode(|node| node.with_name("bob").with_command("polkadot1"))
            })
            .with_parachain(|p| {
                p.with_id(100)
                    .with_default_command("adder-collator")
                    .with_collator(|c| c.with_name("collator1").with_rpc_port(used_port.0))
            })
            .build()
            .unwrap();

        let network_spec_res = NetworkSpec::from_config(&config).await;
        assert!(network_spec_res.is_err());
        assert!(network_spec_res
            .unwrap_err()
            .to_string()
            .contains("err Can\'t bind in socket"));
    }

    #[tokio::test]
    async fn used_port_in_multiple_node_reports_err() {
        use configuration::NetworkConfigBuilder;

        use super::*;
        let used_port_0 = generate_node_port(None).unwrap();
        let used_port_1 = generate_node_port(None).unwrap();

        let config = NetworkConfigBuilder::new()
            .with_relaychain(|r| {
                r.with_chain("rococo-local")
                    .with_default_command("polkadot")
                    .with_validator(|node| node.with_name("alice").with_rpc_port(used_port_0.0))
                    .with_fullnode(|node| node.with_name("bob").with_rpc_port(used_port_1.0))
            })
            .with_parachain(|p| {
                p.with_id(100)
                    .with_default_command("adder-collator")
                    .with_collator(|c| c.with_name("collator1"))
            })
            .build()
            .unwrap();

        let network_spec_res = NetworkSpec::from_config(&config).await;
        assert!(network_spec_res.is_err());
        // ensure we report 2 errors
        let err_str = network_spec_res.unwrap_err().to_string();
        let socket_errs: Vec<&str> = err_str.matches("err Can\'t bind in socket").collect();
        assert_eq!(socket_errs.len(), 2);
    }
}