zombienet_orchestrator/
network_spec.rs

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

zombienet_orchestrator/
network_spec.rs
