zombienet_orchestrator/
network.rs

Help
1pub mod chain_upgrade;
2pub mod node;
3pub mod parachain;
4pub mod relaychain;
5
6use std::{collections::HashMap, path::PathBuf};
7
8use configuration::{
9    para_states::{Initial, Running},
10    shared::node::EnvVar,
11    types::{Arg, Command, Image, Port},
12    ParachainConfig, ParachainConfigBuilder, RegistrationStrategy,
13};
14use provider::{types::TransferedFile, DynNamespace, ProviderError};
15use serde::Serialize;
16use support::fs::FileSystem;
17
18use self::{node::NetworkNode, parachain::Parachain, relaychain::Relaychain};
19use crate::{
20    generators::chain_spec::ChainSpec,
21    network_spec::{self, NetworkSpec},
22    shared::{
23        macros,
24        types::{ChainDefaultContext, RegisterParachainOptions},
25    },
26    spawner::{self, SpawnNodeCtx},
27    ScopedFilesystem, ZombieRole,
28};
29
30#[derive(Serialize)]
31pub struct Network<T: FileSystem> {
32    #[serde(skip)]
33    ns: DynNamespace,
34    #[serde(skip)]
35    filesystem: T,
36    relay: Relaychain,
37    initial_spec: NetworkSpec,
38    parachains: HashMap<u32, Parachain>,
39    #[serde(skip)]
40    nodes_by_name: HashMap<String, NetworkNode>,
41}
42
43impl<T: FileSystem> std::fmt::Debug for Network<T> {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("Network")
46            .field("ns", &"ns_skipped")
47            .field("relay", &self.relay)
48            .field("initial_spec", &self.initial_spec)
49            .field("parachains", &self.parachains)
50            .field("nodes_by_name", &self.nodes_by_name)
51            .finish()
52    }
53}
54
55macros::create_add_options!(AddNodeOptions {
56    chain_spec: Option<PathBuf>
57});
58
59macros::create_add_options!(AddCollatorOptions {
60    chain_spec: Option<PathBuf>,
61    chain_spec_relay: Option<PathBuf>
62});
63
64impl<T: FileSystem> Network<T> {
65    pub(crate) fn new_with_relay(
66        relay: Relaychain,
67        ns: DynNamespace,
68        fs: T,
69        initial_spec: NetworkSpec,
70    ) -> Self {
71        Self {
72            ns,
73            filesystem: fs,
74            relay,
75            initial_spec,
76            parachains: Default::default(),
77            nodes_by_name: Default::default(),
78        }
79    }
80
81    // Pubic API
82    pub fn ns_name(&self) -> String {
83        self.ns.name().to_string()
84    }
85
86    pub fn base_dir(&self) -> Option<&str> {
87        self.ns.base_dir().to_str()
88    }
89
90    pub fn relaychain(&self) -> &Relaychain {
91        &self.relay
92    }
93
94    // Teardown the network
95    pub async fn destroy(self) -> Result<(), ProviderError> {
96        self.ns.destroy().await
97    }
98
99    /// Add a node to the relaychain
100    ///
101    /// NOTE: name must be unique in the whole network. The new node is added to the
102    /// running network instance.
103    ///
104    /// # Example:
105    /// ```rust
106    /// # use provider::NativeProvider;
107    /// # use support::{fs::local::LocalFileSystem};
108    /// # use zombienet_orchestrator::{errors, AddNodeOptions, Orchestrator};
109    /// # use configuration::NetworkConfig;
110    /// # async fn example() -> Result<(), errors::OrchestratorError> {
111    /// #   let provider = NativeProvider::new(LocalFileSystem {});
112    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
113    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
114    /// let mut network = orchestrator.spawn(config).await?;
115    ///
116    /// // Create the options to add the new node
117    /// let opts = AddNodeOptions {
118    ///     rpc_port: Some(9444),
119    ///     is_validator: true,
120    ///     ..Default::default()
121    /// };
122    ///
123    /// network.add_node("new-node", opts).await?;
124    /// #   Ok(())
125    /// # }
126    /// ```
127    pub async fn add_node(
128        &mut self,
129        name: impl Into<String>,
130        options: AddNodeOptions,
131    ) -> Result<(), anyhow::Error> {
132        let name = name.into();
133        let relaychain = self.relaychain();
134
135        if self.nodes_iter().any(|n| n.name == name) {
136            return Err(anyhow::anyhow!("Name: {} is already used.", name));
137        }
138
139        let chain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec {
140            chain_spec_custom_path.clone()
141        } else {
142            PathBuf::from(format!(
143                "{}/{}.json",
144                self.ns.base_dir().to_string_lossy(),
145                relaychain.chain
146            ))
147        };
148
149        let chain_context = ChainDefaultContext {
150            default_command: self.initial_spec.relaychain.default_command.as_ref(),
151            default_image: self.initial_spec.relaychain.default_image.as_ref(),
152            default_resources: self.initial_spec.relaychain.default_resources.as_ref(),
153            default_db_snapshot: self.initial_spec.relaychain.default_db_snapshot.as_ref(),
154            default_args: self.initial_spec.relaychain.default_args.iter().collect(),
155        };
156
157        let mut node_spec =
158            network_spec::node::NodeSpec::from_ad_hoc(&name, options.into(), &chain_context)?;
159
160        node_spec.available_args_output = Some(
161            self.initial_spec
162                .node_available_args_output(&node_spec, self.ns.clone())
163                .await?,
164        );
165
166        let base_dir = self.ns.base_dir().to_string_lossy();
167        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
168
169        let ctx = SpawnNodeCtx {
170            chain_id: &relaychain.chain_id,
171            parachain_id: None,
172            chain: &relaychain.chain,
173            role: ZombieRole::Node,
174            ns: &self.ns,
175            scoped_fs: &scoped_fs,
176            parachain: None,
177            bootnodes_addr: &vec![],
178            wait_ready: true,
179        };
180
181        let global_files_to_inject = vec![TransferedFile::new(
182            chain_spec_path,
183            PathBuf::from(format!("/cfg/{}.json", relaychain.chain)),
184        )];
185
186        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
187
188        // TODO: register the new node as validator in the relaychain
189        // STEPS:
190        //  - check balance of `stash` derivation for validator account
191        //  - call rotate_keys on the new validator
192        //  - call setKeys on the new validator
193        // if node_spec.is_validator {
194        //     let running_node = self.relay.nodes.first().unwrap();
195        //     // tx_helper::validator_actions::register(vec![&node], &running_node.ws_uri, None).await?;
196        // }
197
198        // Add node to relaychain data
199        self.add_running_node(node.clone(), None);
200
201        Ok(())
202    }
203
204    /// Add a new collator to a parachain
205    ///
206    /// NOTE: name must be unique in the whole network.
207    ///
208    /// # Example:
209    /// ```rust
210    /// # use provider::NativeProvider;
211    /// # use support::{fs::local::LocalFileSystem};
212    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
213    /// # use configuration::NetworkConfig;
214    /// # async fn example() -> Result<(), anyhow::Error> {
215    /// #   let provider = NativeProvider::new(LocalFileSystem {});
216    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
217    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
218    /// let mut network = orchestrator.spawn(config).await?;
219    ///
220    /// let col_opts = AddCollatorOptions {
221    ///     command: Some("polkadot-parachain".try_into()?),
222    ///     ..Default::default()
223    /// };
224    ///
225    /// network.add_collator("new-col-1", col_opts, 100).await?;
226    /// #   Ok(())
227    /// # }
228    /// ```
229    pub async fn add_collator(
230        &mut self,
231        name: impl Into<String>,
232        options: AddCollatorOptions,
233        para_id: u32,
234    ) -> Result<(), anyhow::Error> {
235        let spec = self
236            .initial_spec
237            .parachains
238            .iter()
239            .find(|para| para.id == para_id)
240            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
241        let role = if spec.is_cumulus_based {
242            ZombieRole::CumulusCollator
243        } else {
244            ZombieRole::Collator
245        };
246        let chain_context = ChainDefaultContext {
247            default_command: spec.default_command.as_ref(),
248            default_image: spec.default_image.as_ref(),
249            default_resources: spec.default_resources.as_ref(),
250            default_db_snapshot: spec.default_db_snapshot.as_ref(),
251            default_args: spec.default_args.iter().collect(),
252        };
253        let parachain = self
254            .parachains
255            .get(&para_id)
256            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
257
258        let base_dir = self.ns.base_dir().to_string_lossy();
259        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
260
261        // TODO: we want to still supporting spawn a dedicated bootnode??
262        let ctx = SpawnNodeCtx {
263            chain_id: &self.relay.chain_id,
264            parachain_id: parachain.chain_id.as_deref(),
265            chain: &self.relay.chain,
266            role,
267            ns: &self.ns,
268            scoped_fs: &scoped_fs,
269            parachain: Some(spec),
270            bootnodes_addr: &vec![],
271            wait_ready: true,
272        };
273
274        let relaychain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec_relay {
275            chain_spec_custom_path.clone()
276        } else {
277            PathBuf::from(format!(
278                "{}/{}.json",
279                self.ns.base_dir().to_string_lossy(),
280                self.relay.chain
281            ))
282        };
283
284        let mut global_files_to_inject = vec![TransferedFile::new(
285            relaychain_spec_path,
286            PathBuf::from(format!("/cfg/{}.json", self.relay.chain)),
287        )];
288
289        let para_chain_spec_local_path = if let Some(para_chain_spec_custom) = &options.chain_spec {
290            Some(para_chain_spec_custom.clone())
291        } else if let Some(para_spec_path) = &parachain.chain_spec_path {
292            Some(PathBuf::from(format!(
293                "{}/{}",
294                self.ns.base_dir().to_string_lossy(),
295                para_spec_path.to_string_lossy()
296            )))
297        } else {
298            None
299        };
300
301        if let Some(para_spec_path) = para_chain_spec_local_path {
302            global_files_to_inject.push(TransferedFile::new(
303                para_spec_path,
304                PathBuf::from(format!("/cfg/{}.json", para_id)),
305            ));
306        }
307
308        let mut node_spec =
309            network_spec::node::NodeSpec::from_ad_hoc(name.into(), options.into(), &chain_context)?;
310
311        node_spec.available_args_output = Some(
312            self.initial_spec
313                .node_available_args_output(&node_spec, self.ns.clone())
314                .await?,
315        );
316
317        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
318        let para = self.parachains.get_mut(&para_id).unwrap();
319        para.collators.push(node.clone());
320        self.add_running_node(node, None);
321
322        Ok(())
323    }
324
325    /// Get a parachain config builder from a running network
326    ///
327    /// This allow you to build a new parachain config to be deployed into
328    /// the running network.
329    pub fn para_config_builder(&self) -> ParachainConfigBuilder<Initial, Running> {
330        // TODO: build the validation context from the running network
331        ParachainConfigBuilder::new_with_running(Default::default())
332    }
333
334    /// Add a new parachain to the running network
335    ///
336    /// # Arguments
337    /// * `para_config` - Parachain configuration to deploy
338    /// * `custom_relaychain_spec` - Optional path to a custom relaychain spec to use
339    /// * `custom_parchain_fs_prefix` - Optional prefix to use when artifacts are created
340    ///
341    ///
342    /// # Example:
343    /// ```rust
344    /// # use anyhow::anyhow;
345    /// # use provider::NativeProvider;
346    /// # use support::{fs::local::LocalFileSystem};
347    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
348    /// # use configuration::NetworkConfig;
349    /// # async fn example() -> Result<(), anyhow::Error> {
350    /// #   let provider = NativeProvider::new(LocalFileSystem {});
351    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
352    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
353    /// let mut network = orchestrator.spawn(config).await?;
354    /// let para_config = network
355    ///     .para_config_builder()
356    ///     .with_id(100)
357    ///     .with_default_command("polkadot-parachain")
358    ///     .with_collator(|c| c.with_name("col-100-1"))
359    ///     .build()
360    ///     .map_err(|_e| anyhow!("Building config"))?;
361    ///
362    /// network.add_parachain(&para_config, None, None).await?;
363    ///
364    /// #   Ok(())
365    /// # }
366    /// ```
367    pub async fn add_parachain(
368        &mut self,
369        para_config: &ParachainConfig,
370        custom_relaychain_spec: Option<PathBuf>,
371        custom_parchain_fs_prefix: Option<String>,
372    ) -> Result<(), anyhow::Error> {
373        // build
374        let mut para_spec = network_spec::parachain::ParachainSpec::from_config(para_config)?;
375        let base_dir = self.ns.base_dir().to_string_lossy().to_string();
376        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
377
378        let mut global_files_to_inject = vec![];
379
380        // get relaychain id
381        let relay_chain_id = if let Some(custom_path) = custom_relaychain_spec {
382            // use this file as relaychain spec
383            global_files_to_inject.push(TransferedFile::new(
384                custom_path.clone(),
385                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
386            ));
387            let content = std::fs::read_to_string(custom_path)?;
388            ChainSpec::chain_id_from_spec(&content)?
389        } else {
390            global_files_to_inject.push(TransferedFile::new(
391                PathBuf::from(format!(
392                    "{}/{}",
393                    scoped_fs.base_dir,
394                    self.relaychain().chain_spec_path.to_string_lossy()
395                )),
396                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
397            ));
398            self.relay.chain_id.clone()
399        };
400
401        let chain_spec_raw_path = para_spec
402            .build_chain_spec(&relay_chain_id, &self.ns, &scoped_fs)
403            .await?;
404
405        // Para artifacts
406        let para_path_prefix = if let Some(custom_prefix) = custom_parchain_fs_prefix {
407            custom_prefix
408        } else {
409            para_spec.id.to_string()
410        };
411
412        scoped_fs.create_dir(&para_path_prefix).await?;
413        // create wasm/state
414        para_spec
415            .genesis_state
416            .build(
417                chain_spec_raw_path.as_ref(),
418                format!("{}/genesis-state", &para_path_prefix),
419                &self.ns,
420                &scoped_fs,
421            )
422            .await?;
423        para_spec
424            .genesis_wasm
425            .build(
426                chain_spec_raw_path.as_ref(),
427                format!("{}/para_spec-wasm", &para_path_prefix),
428                &self.ns,
429                &scoped_fs,
430            )
431            .await?;
432
433        let parachain =
434            Parachain::from_spec(&para_spec, &global_files_to_inject, &scoped_fs).await?;
435        let parachain_id = parachain.chain_id.clone();
436
437        // Create `ctx` for spawn the nodes
438        let ctx_para = SpawnNodeCtx {
439            parachain: Some(&para_spec),
440            parachain_id: parachain_id.as_deref(),
441            role: if para_spec.is_cumulus_based {
442                ZombieRole::CumulusCollator
443            } else {
444                ZombieRole::Collator
445            },
446            bootnodes_addr: &vec![],
447            chain_id: &self.relaychain().chain_id,
448            chain: &self.relaychain().chain,
449            ns: &self.ns,
450            scoped_fs: &scoped_fs,
451            wait_ready: false,
452        };
453
454        // Register the parachain to the running network
455        let first_node_url = self
456            .relaychain()
457            .nodes
458            .first()
459            .ok_or(anyhow::anyhow!(
460                "At least one node of the relaychain should be running"
461            ))?
462            .ws_uri();
463
464        if para_config.registration_strategy() == Some(&RegistrationStrategy::UsingExtrinsic) {
465            let register_para_options = RegisterParachainOptions {
466                id: parachain.para_id,
467                // This needs to resolve correctly
468                wasm_path: para_spec
469                    .genesis_wasm
470                    .artifact_path()
471                    .ok_or(anyhow::anyhow!(
472                        "artifact path for wasm must be set at this point",
473                    ))?
474                    .to_path_buf(),
475                state_path: para_spec
476                    .genesis_state
477                    .artifact_path()
478                    .ok_or(anyhow::anyhow!(
479                        "artifact path for state must be set at this point",
480                    ))?
481                    .to_path_buf(),
482                node_ws_url: first_node_url.to_string(),
483                onboard_as_para: para_spec.onboard_as_parachain,
484                seed: None, // TODO: Seed is passed by?
485                finalization: false,
486            };
487
488            Parachain::register(register_para_options, &scoped_fs).await?;
489        }
490
491        // Spawn the nodes
492        let spawning_tasks = para_spec
493            .collators
494            .iter()
495            .map(|node| spawner::spawn_node(node, parachain.files_to_inject.clone(), &ctx_para));
496
497        let running_nodes = futures::future::try_join_all(spawning_tasks).await?;
498        let running_para_id = parachain.para_id;
499        self.add_para(parachain);
500        for node in running_nodes {
501            self.add_running_node(node, Some(running_para_id));
502        }
503
504        Ok(())
505    }
506
507    // deregister and stop the collator?
508    // remove_parachain()
509
510    pub fn get_node(&self, name: impl Into<String>) -> Result<&NetworkNode, anyhow::Error> {
511        let name = name.into();
512        if let Some(node) = self.nodes_iter().find(|&n| n.name == name) {
513            return Ok(node);
514        }
515
516        let list = self
517            .nodes_iter()
518            .map(|n| &n.name)
519            .cloned()
520            .collect::<Vec<_>>()
521            .join(", ");
522
523        Err(anyhow::anyhow!(
524            "can't find node with name: {name:?}, should be one of {list}"
525        ))
526    }
527
528    pub fn get_node_mut(
529        &mut self,
530        name: impl Into<String>,
531    ) -> Result<&mut NetworkNode, anyhow::Error> {
532        let name = name.into();
533        self.nodes_iter_mut()
534            .find(|n| n.name == name)
535            .ok_or(anyhow::anyhow!("can't find node with name: {name:?}"))
536    }
537
538    pub fn nodes(&self) -> Vec<&NetworkNode> {
539        self.nodes_by_name.values().collect::<Vec<&NetworkNode>>()
540    }
541
542    pub async fn detach(&self) {
543        self.ns.detach().await
544    }
545
546    // Internal API
547    pub(crate) fn add_running_node(&mut self, node: NetworkNode, para_id: Option<u32>) {
548        if let Some(para_id) = para_id {
549            if let Some(para) = self.parachains.get_mut(&para_id) {
550                para.collators.push(node.clone());
551            } else {
552                // is the first node of the para, let create the entry
553                unreachable!()
554            }
555        } else {
556            self.relay.nodes.push(node.clone());
557        }
558        // TODO: we should hold a ref to the node in the vec in the future.
559        let node_name = node.name.clone();
560        self.nodes_by_name.insert(node_name, node);
561    }
562
563    pub(crate) fn add_para(&mut self, para: Parachain) {
564        self.parachains.insert(para.para_id, para);
565    }
566
567    pub fn name(&self) -> &str {
568        self.ns.name()
569    }
570
571    pub fn parachain(&self, para_id: u32) -> Option<&Parachain> {
572        self.parachains.get(&para_id)
573    }
574
575    pub fn parachains(&self) -> Vec<&Parachain> {
576        self.parachains.values().collect()
577    }
578
579    pub(crate) fn nodes_iter(&self) -> impl Iterator<Item = &NetworkNode> {
580        self.relay
581            .nodes
582            .iter()
583            .chain(self.parachains.values().flat_map(|p| &p.collators))
584    }
585
586    pub(crate) fn nodes_iter_mut(&mut self) -> impl Iterator<Item = &mut NetworkNode> {
587        self.relay.nodes.iter_mut().chain(
588            self.parachains
589                .iter_mut()
590                .flat_map(|(_, p)| &mut p.collators),
591        )
592    }
593}
zombienet_orchestrator/network.rs

zombienet_orchestrator/
network.rs
