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, relaychain.chain.clone()) {
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                    None,
182                )
183                .await?;
184            debug!("parachain genesis state built!");
185            para.genesis_wasm
186                .build(
187                    chain_spec_raw_path,
188                    format!("{}/genesis-wasm", para.unique_id),
189                    &ns,
190                    scoped_fs,
191                    None,
192                )
193                .await?;
194            debug!("parachain genesis wasm built!");
195        }
196
197        Ok(())
198    }
199
200    // collect mutable references to all nodes from relaychain and parachains
201    fn collect_network_nodes(&mut self) -> Vec<&mut NodeSpec> {
202        vec![
203            self.relaychain.nodes.iter_mut().collect::<Vec<_>>(),
204            self.parachains
205                .iter_mut()
206                .flat_map(|para| para.collators.iter_mut())
207                .collect(),
208        ]
209        .into_iter()
210        .flatten()
211        .collect::<Vec<_>>()
212    }
213
214    // initialize the mapping of all possible node image/commands to corresponding nodes
215    fn create_image_command_to_nodes_mapping(
216        network_nodes: Vec<&mut NodeSpec>,
217    ) -> HashMap<(Option<String>, String), Vec<&mut NodeSpec>> {
218        network_nodes.into_iter().fold(
219            HashMap::new(),
220            |mut acc: HashMap<(Option<String>, String), Vec<&mut node::NodeSpec>>, node| {
221                // build mapping key using image and command if image is present or command only
222                let key = node
223                    .image
224                    .as_ref()
225                    .map(|image| {
226                        (
227                            Some(image.as_str().to_string()),
228                            node.command.as_str().to_string(),
229                        )
230                    })
231                    .unwrap_or_else(|| (None, node.command.as_str().to_string()));
232
233                // append the node to the vector of nodes for this image/command tuple
234                if let Entry::Vacant(entry) = acc.entry(key.clone()) {
235                    entry.insert(vec![node]);
236                } else {
237                    acc.get_mut(&key).unwrap().push(node);
238                }
239
240                acc
241            },
242        )
243    }
244
245    async fn retrieve_all_nodes_available_args_output(
246        ns: Arc<dyn ProviderNamespace + Send + Sync>,
247        image_command_to_nodes_mapping: &HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
248    ) -> Result<Vec<(Option<String>, String, String)>, OrchestratorError> {
249        try_join_all(
250            image_command_to_nodes_mapping
251                .keys()
252                .cloned()
253                .map(|(image, command)| async {
254                    // get node available args output from image/command
255                    let available_args = ns
256                        .get_node_available_args((command.clone(), image.clone()))
257                        .await?;
258                    debug!(
259                        "retrieved available args for image: {:?}, command: {}",
260                        image, command
261                    );
262
263                    // map the result to include image and command
264                    Ok::<_, OrchestratorError>((image, command, available_args))
265                })
266                .collect::<Vec<_>>(),
267        )
268        .await
269    }
270
271    fn update_nodes_available_args_output(
272        image_command_to_nodes_mapping: &mut HashMap<(Option<String>, String), Vec<&mut NodeSpec>>,
273        available_args_outputs: Vec<(Option<String>, String, String)>,
274    ) {
275        for (image, command, available_args_output) in available_args_outputs {
276            let nodes = image_command_to_nodes_mapping
277                .get_mut(&(image, command))
278                .expect(&format!(
279                    "node image/command key should exist {THIS_IS_A_BUG}"
280                ));
281
282            for node in nodes {
283                node.available_args_output = Some(available_args_output.clone());
284            }
285        }
286    }
287}
288
289#[cfg(test)]
290mod tests {
291
292    #[tokio::test]
293    async fn small_network_config_get_spec() {
294        use configuration::NetworkConfigBuilder;
295
296        use super::*;
297
298        let config = NetworkConfigBuilder::new()
299            .with_relaychain(|r| {
300                r.with_chain("rococo-local")
301                    .with_default_command("polkadot")
302                    .with_validator(|node| node.with_name("alice"))
303                    .with_fullnode(|node| node.with_name("bob").with_command("polkadot1"))
304            })
305            .with_parachain(|p| {
306                p.with_id(100)
307                    .with_default_command("adder-collator")
308                    .with_collator(|c| c.with_name("collator1"))
309            })
310            .build()
311            .unwrap();
312
313        let network_spec = NetworkSpec::from_config(&config).await.unwrap();
314        let alice = network_spec.relaychain.nodes.first().unwrap();
315        let bob = network_spec.relaychain.nodes.get(1).unwrap();
316        assert_eq!(alice.command.as_str(), "polkadot");
317        assert_eq!(bob.command.as_str(), "polkadot1");
318        assert!(alice.is_validator);
319        assert!(!bob.is_validator);
320
321        // paras
322        assert_eq!(network_spec.parachains.len(), 1);
323        let para_100 = network_spec.parachains.first().unwrap();
324        assert_eq!(para_100.id, 100);
325    }
326}
zombienet_orchestrator/network_spec.rs

zombienet_orchestrator/
network_spec.rs
