zombienet_orchestrator/
network_spec.rs

1use std::{
2    collections::{hash_map::Entry, HashMap},
3    sync::Arc,
4};
5
6use configuration::{GlobalSettings, HrmpChannelConfig, NetworkConfig};
7use futures::future::try_join_all;
8use provider::{DynNamespace, ProviderError, ProviderNamespace};
9use serde::{Deserialize, Serialize};
10use support::{constants::THIS_IS_A_BUG, fs::FileSystem};
11use tracing::{debug, trace};
12
13use crate::{errors::OrchestratorError, ScopedFilesystem};
14
15pub mod node;
16pub mod parachain;
17pub mod relaychain;
18
19use self::{node::NodeSpec, parachain::ParachainSpec, relaychain::RelaychainSpec};
20
21#[derive(Debug, Clone, Serialize, Deserialize)]
22pub struct NetworkSpec {
23    /// Relaychain configuration.
24    pub(crate) relaychain: RelaychainSpec,
25
26    /// Parachains configurations.
27    pub(crate) parachains: Vec<ParachainSpec>,
28
29    /// HRMP channels configurations.
30    pub(crate) hrmp_channels: Vec<HrmpChannelConfig>,
31
32    /// Global settings
33    pub(crate) global_settings: GlobalSettings,
34}
35
36impl NetworkSpec {
37    pub async fn from_config(
38        network_config: &NetworkConfig,
39    ) -> Result<NetworkSpec, OrchestratorError> {
40        let mut errs = vec![];
41        let relaychain = RelaychainSpec::from_config(network_config.relaychain())?;
42        let mut parachains = vec![];
43
44        // TODO: move to `fold` or map+fold
45        for para_config in network_config.parachains() {
46            match ParachainSpec::from_config(para_config) {
47                Ok(para) => parachains.push(para),
48                Err(err) => errs.push(err),
49            }
50        }
51
52        if errs.is_empty() {
53            Ok(NetworkSpec {
54                relaychain,
55                parachains,
56                hrmp_channels: network_config
57                    .hrmp_channels()
58                    .into_iter()
59                    .cloned()
60                    .collect(),
61                global_settings: network_config.global_settings().clone(),
62            })
63        } else {
64            let errs_str = errs
65                .into_iter()
66                .map(|e| e.to_string())
67                .collect::<Vec<String>>()
68                .join("\n");
69            Err(OrchestratorError::InvalidConfig(errs_str))
70        }
71    }
72
73    pub async fn populate_nodes_available_args(
74        &mut self,
75        ns: Arc<dyn ProviderNamespace + Send + Sync>,
76    ) -> Result<(), OrchestratorError> {
77        let network_nodes = self.collect_network_nodes();
78
79        let mut image_command_to_nodes_mapping =
80            Self::create_image_command_to_nodes_mapping(network_nodes);
81
82        let available_args_outputs =
83            Self::retrieve_all_nodes_available_args_output(ns, &image_command_to_nodes_mapping)
84                .await?;
85
86        Self::update_nodes_available_args_output(
87            &mut image_command_to_nodes_mapping,
88            available_args_outputs,
89        );
90
91        Ok(())
92    }
93
94    //
95    pub async fn node_available_args_output(
96        &self,
97        node_spec: &NodeSpec,
98        ns: Arc<dyn ProviderNamespace + Send + Sync>,
99    ) -> Result<String, ProviderError> {
100        // try to find a node that use the same combination of image/cmd
101        let cmp_fn = |ad_hoc: &&NodeSpec| -> bool {
102            ad_hoc.image == node_spec.image && ad_hoc.command == node_spec.command
103        };
104
105        // check if we already had computed the args output for this cmd/[image]
106        let node = self.relaychain.nodes.iter().find(cmp_fn);
107        let node = if let Some(node) = node {
108            Some(node)
109        } else {
110            let node = self
111                .parachains
112                .iter()
113                .find_map(|para| para.collators.iter().find(cmp_fn));
114
115            node
116        };
117
118        let output = if let Some(node) = node {
119            node.available_args_output.clone().expect(&format!(
120                "args_output should be set for running nodes {THIS_IS_A_BUG}"
121            ))
122        } else {
123            // we need to compute the args output
124            let image = node_spec
125                .image
126                .as_ref()
127                .map(|image| image.as_str().to_string());
128            let command = node_spec.command.as_str().to_string();
129
130            ns.get_node_available_args((command, image)).await?
131        };
132
133        Ok(output)
134    }
135
136    pub fn relaychain(&self) -> &RelaychainSpec {
137        &self.relaychain
138    }
139
140    pub fn relaychain_mut(&mut self) -> &mut RelaychainSpec {
141        &mut self.relaychain
142    }
143
144    pub fn parachains_iter(&self) -> impl Iterator<Item = &ParachainSpec> {
145        self.parachains.iter()
146    }
147
148    pub fn parachains_iter_mut(&mut self) -> impl Iterator<Item = &mut ParachainSpec> {
149        self.parachains.iter_mut()
150    }
151
152    pub fn set_global_settings(&mut self, global_settings: GlobalSettings) {
153        self.global_settings = global_settings;
154    }
155
156    pub async fn build_parachain_artifacts<'a, T: FileSystem>(
157        &mut self,
158        ns: DynNamespace,
159        scoped_fs: &ScopedFilesystem<'a, T>,
160        relaychain_id: &str,
161        base_dir_exists: bool,
162    ) -> Result<(), anyhow::Error> {
163        for para in self.parachains.iter_mut() {
164            let chain_spec_raw_path = para.build_chain_spec(relaychain_id, &ns, scoped_fs).await?;
165
166            trace!("creating dirs for {}", &para.unique_id);
167            if base_dir_exists {
168                scoped_fs.create_dir_all(&para.unique_id).await?;
169            } else {
170                scoped_fs.create_dir(&para.unique_id).await?;
171            };
172            trace!("created dirs for {}", &para.unique_id);
173
174            // create wasm/state
175            para.genesis_state
176                .build(
177                    chain_spec_raw_path.clone(),
178                    format!("{}/genesis-state", para.unique_id),
179                    &ns,
180                    scoped_fs,
181                )
182                .await?;
183            debug!("parachain genesis state built!");
184            para.genesis_wasm
185                .build(
186                    chain_spec_raw_path,
187                    format!("{}/genesis-wasm", para.unique_id),
188                    &ns,
189                    scoped_fs,
190                )
191                .await?;
192            debug!("parachain genesis wasm built!");
193        }
194
195        Ok(())
196    }
197
198    // collect mutable references to all nodes from relaychain and parachains
199    fn collect_network_nodes(&mut self) -> Vec<&mut NodeSpec> {
200        vec![
201            self.relaychain.nodes.iter_mut().collect::<Vec<_>>(),
202            self.parachains
203                .iter_mut()
204                .flat_map(|para| para.collators.iter_mut())
205                .collect(),
206        ]
207        .into_iter()
208        .flatten()
209        .collect::<Vec<_>>()
210    }
211
212    // initialize the mapping of all possible node image/commands to corresponding nodes
213    fn create_image_command_to_nodes_mapping(
214        network_nodes: Vec<&mut NodeSpec>,
215    ) -> HashMap<(Option<String>, String), Vec<&mut NodeSpec>> {
216        network_nodes.into_iter().fold(
217            HashMap::new(),
218            |mut acc: HashMap<(Option<String>, String), Vec<&mut node::NodeSpec>>, node| {
219                // build mapping key using image and command if image is present or command only
220                let key = node
221                    .image
222                    .as_ref()
223                    .map(|image| {
224                        (
225                            Some(image.as_str().to_string()),
226                            node.command.as_str().to_string(),
227                        )
228                    })
229                    .unwrap_or_else(|| (None, node.command.as_str().to_string()));
230
231                // append the node to the vector of nodes for this image/command tuple
232                if let Entry::Vacant(entry) = acc.entry(key.clone()) {
233                    entry.insert(vec![node]);
234                } else {
235                    acc.get_mut(&key).unwrap().push(node);
236                }
237
238                acc
239            },
240        )
241    }
242
243    async fn retrieve_all_nodes_available_args_output(
244        ns: Arc<dyn ProviderNamespace + Send + Sync>,
245        image_command_to_nodes_mapping: &HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
246    ) -> Result<Vec<(Option<String>, String, String)>, OrchestratorError> {
247        try_join_all(
248            image_command_to_nodes_mapping
249                .keys()
250                .cloned()
251                .map(|(image, command)| async {
252                    // get node available args output from image/command
253                    let available_args = ns
254                        .get_node_available_args((command.clone(), image.clone()))
255                        .await?;
256                    debug!(
257                        "retrieved available args for image: {:?}, command: {}",
258                        image, command
259                    );
260
261                    // map the result to include image and command
262                    Ok::<_, OrchestratorError>((image, command, available_args))
263                })
264                .collect::<Vec<_>>(),
265        )
266        .await
267    }
268
269    fn update_nodes_available_args_output(
270        image_command_to_nodes_mapping: &mut HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
271        available_args_outputs: Vec<(Option<String>, String, String)>,
272    ) {
273        for (image, command, available_args_output) in available_args_outputs {
274            let nodes = image_command_to_nodes_mapping
275                .get_mut(&(image, command))
276                .expect(&format!(
277                    "node image/command key should exist {THIS_IS_A_BUG}"
278                ));
279
280            for node in nodes {
281                node.available_args_output = Some(available_args_output.clone());
282            }
283        }
284    }
285}
286
287#[cfg(test)]
288mod tests {
289
290    #[tokio::test]
291    async fn small_network_config_get_spec() {
292        use configuration::NetworkConfigBuilder;
293
294        use super::*;
295
296        let config = NetworkConfigBuilder::new()
297            .with_relaychain(|r| {
298                r.with_chain("rococo-local")
299                    .with_default_command("polkadot")
300                    .with_node(|node| node.with_name("alice"))
301                    .with_node(|node| {
302                        node.with_name("bob")
303                            .with_command("polkadot1")
304                            .validator(false)
305                    })
306            })
307            .with_parachain(|p| {
308                p.with_id(100)
309                    .with_default_command("adder-collator")
310                    .with_collator(|c| c.with_name("collator1"))
311            })
312            .build()
313            .unwrap();
314
315        let network_spec = NetworkSpec::from_config(&config).await.unwrap();
316        let alice = network_spec.relaychain.nodes.first().unwrap();
317        let bob = network_spec.relaychain.nodes.get(1).unwrap();
318        assert_eq!(alice.command.as_str(), "polkadot");
319        assert_eq!(bob.command.as_str(), "polkadot1");
320        assert!(alice.is_validator);
321        assert!(!bob.is_validator);
322
323        // paras
324        assert_eq!(network_spec.parachains.len(), 1);
325        let para_100 = network_spec.parachains.first().unwrap();
326        assert_eq!(para_100.id, 100);
327    }
328}
zombienet_orchestrator/network_spec.rs

zombienet_orchestrator/
network_spec.rs
