zombienet_orchestrator/network_spec/
relaychain.rs

1use std::collections::{HashMap, HashSet};
2
3use configuration::{
4    shared::{
5        helpers::generate_unique_node_name_from_names,
6        resources::Resources,
7        types::{Arg, AssetLocation, Chain, Command, Image},
8    },
9    types::JsonOverrides,
10    NodeConfig, RelaychainConfig,
11};
12use serde::{Deserialize, Serialize};
13use support::replacer::apply_replacements;
14
15use super::node::NodeSpec;
16use crate::{
17    errors::OrchestratorError,
18    generators::chain_spec::{ChainSpec, Context},
19    shared::{constants::DEFAULT_CHAIN_SPEC_TPL_COMMAND, types::ChainDefaultContext},
20};
21
22/// A relaychain configuration spec
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RelaychainSpec {
25    /// Chain to use (e.g. rococo-local).
26    pub(crate) chain: Chain,
27
28    /// Default command to run the node. Can be overridden on each node.
29    pub(crate) default_command: Option<Command>,
30
31    /// Default image to use (only podman/k8s). Can be overridden on each node.
32    pub(crate) default_image: Option<Image>,
33
34    /// Default resources. Can be overridden on each node.
35    pub(crate) default_resources: Option<Resources>,
36
37    /// Default database snapshot. Can be overridden on each node.
38    pub(crate) default_db_snapshot: Option<AssetLocation>,
39
40    /// Default arguments to use in nodes. Can be overridden on each node.
41    pub(crate) default_args: Vec<Arg>,
42
43    // chain_spec_path: Option<AssetLocation>,
44    pub(crate) chain_spec: ChainSpec,
45
46    /// Set the count of nominators to generator (used with PoS networks).
47    pub(crate) random_nominators_count: u32,
48
49    /// Set the max nominators value (used with PoS networks).
50    pub(crate) max_nominations: u8,
51
52    /// Genesis overrides as JSON value.
53    pub(crate) runtime_genesis_patch: Option<serde_json::Value>,
54
55    /// Wasm override path/url to use.
56    pub(crate) wasm_override: Option<AssetLocation>,
57
58    /// Nodes to run.
59    pub(crate) nodes: Vec<NodeSpec>,
60
61    /// Raw chain-spec override path, url or inline json to use.
62    pub(crate) raw_spec_override: Option<JsonOverrides>,
63}
64
65impl RelaychainSpec {
66    pub fn from_config(config: &RelaychainConfig) -> Result<RelaychainSpec, OrchestratorError> {
67        // Relaychain main command to use, in order:
68        // set as `default_command` or
69        // use the command of the first node.
70        // If non of those is set, return an error.
71        let main_cmd = config
72            .default_command()
73            .or(config.nodes().first().and_then(|node| node.command()))
74            .ok_or(OrchestratorError::InvalidConfig(
75                "Relaychain, either default_command or first node with a command needs to be set."
76                    .to_string(),
77            ))?;
78
79        // TODO: internally we use image as String
80        let main_image = config
81            .default_image()
82            .or(config.nodes().first().and_then(|node| node.image()))
83            .map(|image| image.as_str().to_string());
84
85        let replacements = HashMap::from([
86            ("disableBootnodes", "--disable-default-bootnode"),
87            ("mainCommand", main_cmd.as_str()),
88        ]);
89        let tmpl = if let Some(tmpl) = config.chain_spec_command() {
90            apply_replacements(tmpl, &replacements)
91        } else {
92            apply_replacements(DEFAULT_CHAIN_SPEC_TPL_COMMAND, &replacements)
93        };
94
95        let chain_spec = ChainSpec::new(config.chain().as_str(), Context::Relay)
96            .set_chain_name(config.chain().as_str())
97            .command(
98                tmpl.as_str(),
99                config.chain_spec_command_is_local(),
100                config.chain_spec_command_output_path(),
101            )
102            .image(main_image.clone());
103
104        // Add asset location if present
105        let chain_spec = if let Some(chain_spec_path) = config.chain_spec_path() {
106            chain_spec.asset_location(chain_spec_path.clone())
107        } else {
108            chain_spec
109        };
110
111        // add chain-spec runtime  if present
112        let chain_spec = if let Some(chain_spec_runtime) = config.chain_spec_runtime() {
113            chain_spec.runtime(chain_spec_runtime.clone())
114        } else {
115            chain_spec
116        };
117
118        // build the `node_specs`
119        let chain_context = ChainDefaultContext {
120            default_command: config.default_command(),
121            default_image: config.default_image(),
122            default_resources: config.default_resources(),
123            default_db_snapshot: config.default_db_snapshot(),
124            default_args: config.default_args(),
125        };
126
127        let mut nodes: Vec<NodeConfig> = config.nodes().into_iter().cloned().collect();
128        nodes.extend(
129            config
130                .group_node_configs()
131                .into_iter()
132                .flat_map(|node_group| node_group.expand_group_configs()),
133        );
134
135        let mut names = HashSet::new();
136        let (nodes, mut errs) = nodes
137            .iter()
138            .map(|node_config| NodeSpec::from_config(node_config, &chain_context, false, false))
139            .fold((vec![], vec![]), |(mut nodes, mut errs), result| {
140                match result {
141                    Ok(mut node) => {
142                        let unique_name =
143                            generate_unique_node_name_from_names(node.name, &mut names);
144                        node.name = unique_name;
145                        nodes.push(node);
146                    },
147                    Err(err) => errs.push(err),
148                }
149                (nodes, errs)
150            });
151
152        if !errs.is_empty() {
153            // TODO: merge errs, maybe return something like Result<Sometype, Vec<OrchestratorError>>
154            return Err(errs.swap_remove(0));
155        }
156
157        Ok(RelaychainSpec {
158            chain: config.chain().clone(),
159            default_command: config.default_command().cloned(),
160            default_image: config.default_image().cloned(),
161            default_resources: config.default_resources().cloned(),
162            default_db_snapshot: config.default_db_snapshot().cloned(),
163            wasm_override: config.wasm_override().cloned(),
164            default_args: config.default_args().into_iter().cloned().collect(),
165            chain_spec,
166            random_nominators_count: config.random_nominators_count().unwrap_or(0),
167            max_nominations: config.max_nominations().unwrap_or(24),
168            runtime_genesis_patch: config.runtime_genesis_patch().cloned(),
169            nodes,
170            raw_spec_override: config.raw_spec_override().cloned(),
171        })
172    }
173
174    pub fn chain_spec(&self) -> &ChainSpec {
175        &self.chain_spec
176    }
177
178    pub fn chain_spec_mut(&mut self) -> &mut ChainSpec {
179        &mut self.chain_spec
180    }
181}