Skip to main content

zombienet_orchestrator/
network.rs

1pub mod chain_upgrade;
2pub mod jamchain;
3pub mod node;
4pub mod parachain;
5pub mod relaychain;
6
7use std::{
8    cell::RefCell,
9    collections::HashMap,
10    path::PathBuf,
11    rc::Rc,
12    sync::Arc,
13    time::{Duration, SystemTime, UNIX_EPOCH},
14};
15
16use configuration::{
17    para_states::{Initial, Running},
18    shared::{
19        helpers::{generate_unique_node_name_from_names, sanitize_node_name},
20        node::EnvVar,
21    },
22    types::{Arg, Command, Image, ParaId, Port, ValidationContext},
23    ParachainConfig, ParachainConfigBuilder, RegistrationStrategy,
24};
25use provider::{types::TransferedFile, DynNamespace, ProviderError};
26use serde::{Deserialize, Serialize};
27use support::{
28    constants::{RELAY_NOT_NONE, THIS_IS_A_BUG},
29    fs::FileSystem,
30};
31use tokio::sync::RwLock;
32use tracing::{error, warn};
33
34use self::{
35    jamchain::Jamchain,
36    node::{JamNetworkNode, NetworkNode, NodeKind, SpawnedNode},
37    parachain::Parachain,
38    relaychain::Relaychain,
39};
40use crate::{
41    generators::{self, chain_spec::ChainSpec},
42    network_spec::{self, NetworkSpec},
43    observability::{self, ObservabilityInfo, ObservabilityState},
44    shared::{
45        constants::{NODE_MONITORING_FAILURE_THRESHOLD_SECONDS, NODE_MONITORING_INTERVAL_SECONDS},
46        macros,
47        types::{ChainDefaultContext, RegisterParachainOptions},
48    },
49    spawner::{self, SpawnNodeCtx},
50    utils::write_zombie_json,
51    ScopedFilesystem, ZombieRole,
52};
53
54/// Context where the node is running
55/// RC or Para
56#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
57pub enum NodeContext {
58    Rc,
59    Para {
60        para_id: ParaId,
61        is_cumulus_based: bool,
62    },
63    Jam,
64}
65
66#[derive(Serialize)]
67pub struct Network<T: FileSystem> {
68    #[serde(skip)]
69    ns: DynNamespace,
70    #[serde(skip)]
71    filesystem: T,
72    #[serde(skip_serializing_if = "Option::is_none")]
73    relay: Option<Relaychain>,
74    initial_spec: NetworkSpec,
75    parachains: HashMap<u32, Vec<Parachain>>,
76    #[serde(skip_serializing_if = "Option::is_none")]
77    jamchain: Option<Jamchain>,
78    /// The custom processes running next to the nodes, in the order each one
79    /// met (or failed) its condition; `zombie.json` keeps that order.
80    #[serde(skip_serializing_if = "Vec::is_empty")]
81    custom_processes: Vec<Arc<node::CustomProcessNode>>,
82    /// Every node spawned in this network, by name, regardless of its kind.
83    ///
84    /// Holds the same `Arc`s as the typed collections above; downcast with
85    /// [`Network::get_node`] / [`Network::get_jam_node`] to get the concrete
86    /// type back.
87    #[serde(skip)]
88    nodes_by_name: HashMap<String, Arc<dyn SpawnedNode>>,
89    #[serde(skip)]
90    nodes_to_watch: Arc<RwLock<Vec<Arc<dyn SpawnedNode>>>>,
91    #[serde(skip_serializing_if = "Option::is_none")]
92    start_time_ts: Option<String>,
93    #[serde(skip)]
94    observability: ObservabilityState,
95}
96
97impl<T: FileSystem> std::fmt::Debug for Network<T> {
98    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
99        f.debug_struct("Network")
100            .field("ns", &"ns_skipped")
101            .field("relay", &self.relay)
102            .field("initial_spec", &self.initial_spec)
103            .field("parachains", &self.parachains)
104            .field("jamchain", &self.jamchain)
105            .field("custom_processes", &self.custom_processes)
106            .field("nodes_by_name", &self.nodes_by_name.keys())
107            .field("observability", &self.observability)
108            .finish()
109    }
110}
111
112macros::create_add_options!(AddNodeOptions {
113    chain_spec: Option<PathBuf>,
114    override_eth_key: Option<String>
115});
116
117macros::create_add_options!(AddCollatorOptions {
118    chain_spec: Option<PathBuf>,
119    chain_spec_relay: Option<PathBuf>,
120    override_eth_key: Option<String>
121});
122
123impl<T: FileSystem> Network<T> {
124    pub(crate) fn new_with_relay(
125        relay: Relaychain,
126        ns: DynNamespace,
127        fs: T,
128        initial_spec: NetworkSpec,
129    ) -> Self {
130        Self::new(Some(relay), ns, fs, initial_spec)
131    }
132
133    /// Create a network whose root chain is a JAM chain instead of a relaychain.
134    pub(crate) fn new_without_relay(ns: DynNamespace, fs: T, initial_spec: NetworkSpec) -> Self {
135        Self::new(None, ns, fs, initial_spec)
136    }
137
138    fn new(relay: Option<Relaychain>, ns: DynNamespace, fs: T, initial_spec: NetworkSpec) -> Self {
139        Self {
140            ns,
141            filesystem: fs,
142            relay,
143            initial_spec,
144            parachains: Default::default(),
145            jamchain: Default::default(),
146            custom_processes: Default::default(),
147            nodes_by_name: Default::default(),
148            nodes_to_watch: Default::default(),
149            start_time_ts: Default::default(),
150            observability: ObservabilityState::default(),
151        }
152    }
153
154    // Pubic API
155    pub fn ns_name(&self) -> String {
156        self.ns.name().to_string()
157    }
158
159    pub fn base_dir(&self) -> Option<&str> {
160        self.ns.base_dir().to_str()
161    }
162
163    pub fn relaychain(&self) -> &Relaychain {
164        self.relay
165            .as_ref()
166            .expect(&format!("{RELAY_NOT_NONE}, {THIS_IS_A_BUG}"))
167    }
168
169    /// The relaychain of the network, `None` for a JAM network.
170    pub fn try_relaychain(&self) -> Option<&Relaychain> {
171        self.relay.as_ref()
172    }
173
174    fn relaychain_mut(&mut self) -> &mut Relaychain {
175        self.relay
176            .as_mut()
177            .expect(&format!("{RELAY_NOT_NONE}, {THIS_IS_A_BUG}"))
178    }
179
180    // Teardown the network
181    pub async fn destroy(mut self) -> Result<(), ProviderError> {
182        if let Err(e) = self.stop_observability().await {
183            warn!("⚠️  Failed to cleanup observability stack: {e}");
184        }
185        self.ns.destroy().await
186    }
187
188    pub fn observability(&self) -> Option<&ObservabilityInfo> {
189        self.observability.as_runnnig()
190    }
191
192    pub fn observability_state(&self) -> &ObservabilityState {
193        &self.observability
194    }
195
196    /// Add a node to the relaychain
197    // The new node is added to the running network instance.
198    /// # Example:
199    /// ```rust
200    /// # use provider::NativeProvider;
201    /// # use support::{fs::local::LocalFileSystem};
202    /// # use zombienet_orchestrator::{errors, AddNodeOptions, Orchestrator};
203    /// # use configuration::NetworkConfig;
204    /// # async fn example() -> Result<(), errors::OrchestratorError> {
205    /// #   let provider = NativeProvider::new(LocalFileSystem {});
206    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
207    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
208    /// let mut network = orchestrator.spawn(config).await?;
209    ///
210    /// // Create the options to add the new node
211    /// let opts = AddNodeOptions {
212    ///     rpc_port: Some(9444),
213    ///     is_validator: true,
214    ///     ..Default::default()
215    /// };
216    ///
217    /// network.add_node("new-node", opts).await?;
218    /// #   Ok(())
219    /// # }
220    /// ```
221    pub async fn add_node(
222        &mut self,
223        name: impl Into<String>,
224        options: AddNodeOptions,
225    ) -> Result<(), anyhow::Error> {
226        let name = generate_unique_node_name_from_names(
227            name,
228            &mut self.nodes_by_name.keys().cloned().collect(),
229        );
230
231        let relaychain = self.relaychain();
232
233        let chain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec {
234            chain_spec_custom_path.clone()
235        } else {
236            PathBuf::from(format!(
237                "{}/{}.json",
238                self.ns.base_dir().to_string_lossy(),
239                relaychain.chain
240            ))
241        };
242
243        let chain_context = ChainDefaultContext {
244            default_command: self.initial_spec.relaychain().default_command.as_ref(),
245            default_image: self.initial_spec.relaychain().default_image.as_ref(),
246            default_resources: self.initial_spec.relaychain().default_resources.as_ref(),
247            default_db_snapshot: self.initial_spec.relaychain().default_db_snapshot.as_ref(),
248            default_args: self.initial_spec.relaychain().default_args.iter().collect(),
249        };
250
251        let mut node_spec = network_spec::node::NodeSpec::from_ad_hoc(
252            &name,
253            options.into(),
254            &chain_context,
255            false,
256            false,
257        )?;
258
259        node_spec.available_args_output = Some(
260            self.initial_spec
261                .node_available_args_output(&node_spec, self.ns.clone())
262                .await?,
263        );
264
265        let base_dir = self.ns.base_dir().to_string_lossy();
266        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
267
268        let resolved_db_snapshots = generators::resolve_db_snapshots(
269            std::iter::once(&node_spec),
270            &self.ns,
271            &self.filesystem,
272        )
273        .await?;
274
275        let ctx = SpawnNodeCtx {
276            chain_id: &relaychain.chain_id,
277            parachain_id: None,
278            chain: &relaychain.chain,
279            role: ZombieRole::Node,
280            ns: &self.ns,
281            scoped_fs: &scoped_fs,
282            parachain: None,
283            bootnodes_addr: &vec![],
284            wait_ready: true,
285            nodes_by_name: self.nodes_json()?,
286            global_settings: &self.initial_spec.global_settings,
287            resolved_db_snapshots: &resolved_db_snapshots,
288        };
289
290        let global_files_to_inject = vec![TransferedFile::new(
291            chain_spec_path,
292            PathBuf::from(format!("/cfg/{}.json", relaychain.chain)),
293        )];
294
295        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
296
297        // TODO: register the new node as validator in the relaychain
298        // STEPS:
299        //  - check balance of `stash` derivation for validator account
300        //  - call rotate_keys on the new validator
301        //  - call setKeys on the new validator
302        // if node_spec.is_validator {
303        //     let running_node = self.relay.nodes.first().unwrap();
304        //     // tx_helper::validator_actions::register(vec![&node], &running_node.ws_uri, None).await?;
305        // }
306
307        // Let's make sure node is up before adding
308        node.wait_until_is_up(self.initial_spec.global_settings.network_spawn_timeout())
309            .await?;
310
311        // Add node to relaychain data
312        self.add_running_node(node.clone(), None).await;
313
314        // Dump zombie.json
315        self.write_zombie_json().await?;
316
317        generators::cleanup_db_snapshot_cache(&resolved_db_snapshots).await;
318
319        Ok(())
320    }
321
322    /// Add a new collator to a parachain
323    ///
324    /// NOTE: if more parachains with given id available (rare corner case)
325    /// then it adds collator to the first parachain
326    ///
327    /// # Example:
328    /// ```rust
329    /// # use provider::NativeProvider;
330    /// # use support::{fs::local::LocalFileSystem};
331    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
332    /// # use configuration::NetworkConfig;
333    /// # async fn example() -> Result<(), anyhow::Error> {
334    /// #   let provider = NativeProvider::new(LocalFileSystem {});
335    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
336    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
337    /// let mut network = orchestrator.spawn(config).await?;
338    ///
339    /// let col_opts = AddCollatorOptions {
340    ///     command: Some("polkadot-parachain".try_into()?),
341    ///     ..Default::default()
342    /// };
343    ///
344    /// network.add_collator("new-col-1", col_opts, 100).await?;
345    /// #   Ok(())
346    /// # }
347    /// ```
348    pub async fn add_collator(
349        &mut self,
350        name: impl Into<String>,
351        options: AddCollatorOptions,
352        para_id: u32,
353    ) -> Result<(), anyhow::Error> {
354        let name = generate_unique_node_name_from_names(
355            name,
356            &mut self.nodes_by_name.keys().cloned().collect(),
357        );
358        let spec = self
359            .initial_spec
360            .parachains
361            .iter()
362            .find(|para| para.id == para_id)
363            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
364        let role = if spec.is_cumulus_based {
365            ZombieRole::CumulusCollator
366        } else {
367            ZombieRole::Collator
368        };
369        let chain_context = ChainDefaultContext {
370            default_command: spec.default_command.as_ref(),
371            default_image: spec.default_image.as_ref(),
372            default_resources: spec.default_resources.as_ref(),
373            default_db_snapshot: spec.default_db_snapshot.as_ref(),
374            default_args: spec.default_args.iter().collect(),
375        };
376
377        let relay_chain = self.relaychain().chain.clone();
378        let relay_chain_id = self.relaychain().chain_id.clone();
379
380        let parachain = self
381            .parachains
382            .get_mut(&para_id)
383            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?
384            .get_mut(0)
385            .ok_or(anyhow::anyhow!(format!("parachain: {para_id} not found!")))?;
386
387        let base_dir = self.ns.base_dir().to_string_lossy();
388        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
389
390        let relaychain_spec_path = if let Some(chain_spec_custom_path) = &options.chain_spec_relay {
391            chain_spec_custom_path.clone()
392        } else {
393            PathBuf::from(format!(
394                "{}/{}.json",
395                self.ns.base_dir().to_string_lossy(),
396                relay_chain
397            ))
398        };
399
400        let mut global_files_to_inject = vec![TransferedFile::new(
401            relaychain_spec_path,
402            PathBuf::from(format!("/cfg/{}.json", relay_chain)),
403        )];
404
405        let para_chain_spec_local_path = if let Some(para_chain_spec_custom) = &options.chain_spec {
406            Some(para_chain_spec_custom.clone())
407        } else if let Some(para_spec_path) = &parachain.chain_spec_path {
408            Some(PathBuf::from(format!(
409                "{}/{}",
410                self.ns.base_dir().to_string_lossy(),
411                para_spec_path.to_string_lossy()
412            )))
413        } else {
414            None
415        };
416
417        if let Some(para_spec_path) = para_chain_spec_local_path {
418            global_files_to_inject.push(TransferedFile::new(
419                para_spec_path,
420                PathBuf::from(format!("/cfg/{para_id}.json")),
421            ));
422        }
423
424        let mut node_spec = network_spec::node::NodeSpec::from_ad_hoc(
425            name,
426            options.into(),
427            &chain_context,
428            true,
429            spec.is_evm_based,
430        )?;
431
432        node_spec.available_args_output = Some(
433            self.initial_spec
434                .node_available_args_output(&node_spec, self.ns.clone())
435                .await?,
436        );
437
438        let resolved_db_snapshots = generators::resolve_db_snapshots(
439            std::iter::once(&node_spec),
440            &self.ns,
441            &self.filesystem,
442        )
443        .await?;
444
445        // TODO: we want to still supporting spawn a dedicated bootnode??
446        let ctx = SpawnNodeCtx {
447            chain_id: &relay_chain_id,
448            parachain_id: parachain.chain_id.as_deref(),
449            chain: &relay_chain,
450            role,
451            ns: &self.ns,
452            scoped_fs: &scoped_fs,
453            parachain: Some(spec),
454            bootnodes_addr: &vec![],
455            wait_ready: true,
456            nodes_by_name: serde_json::to_value(&self.nodes_by_name)?,
457            global_settings: &self.initial_spec.global_settings,
458            resolved_db_snapshots: &resolved_db_snapshots,
459        };
460
461        let node = spawner::spawn_node(&node_spec, global_files_to_inject, &ctx).await?;
462
463        // Let's make sure node is up before adding
464        node.wait_until_is_up(self.initial_spec.global_settings.network_spawn_timeout())
465            .await?;
466
467        self.add_running_node(node, Some(para_id)).await;
468
469        // Dump zombie.json
470        self.write_zombie_json().await?;
471
472        generators::cleanup_db_snapshot_cache(&resolved_db_snapshots).await;
473
474        Ok(())
475    }
476
477    /// Get a parachain config builder from a running network
478    ///
479    /// This allow you to build a new parachain config to be deployed into
480    /// the running network.
481    pub fn para_config_builder(&self) -> ParachainConfigBuilder<Initial, Running> {
482        let used_ports = self
483            .nodes_iter()
484            .map(|node| node.spec())
485            .flat_map(|spec| {
486                [
487                    spec.ws_port.0,
488                    spec.rpc_port.0,
489                    spec.prometheus_port.0,
490                    spec.p2p_port.0,
491                ]
492            })
493            .collect();
494
495        let used_nodes_names = self.nodes_by_name.keys().cloned().collect();
496
497        // need to inverse logic of generate_unique_para_id
498        let used_para_ids = self
499            .parachains
500            .iter()
501            .map(|(id, paras)| (*id, paras.len().saturating_sub(1) as u8))
502            .collect();
503
504        let context = ValidationContext {
505            used_ports,
506            used_nodes_names,
507            used_para_ids,
508        };
509        let context = Rc::new(RefCell::new(context));
510
511        ParachainConfigBuilder::new_with_running(context)
512    }
513
514    /// Add a new parachain to the running network
515    ///
516    /// # Arguments
517    /// * `para_config` - Parachain configuration to deploy
518    /// * `custom_relaychain_spec` - Optional path to a custom relaychain spec to use
519    /// * `custom_parchain_fs_prefix` - Optional prefix to use when artifacts are created
520    ///
521    ///
522    /// # Example:
523    /// ```rust
524    /// # use anyhow::anyhow;
525    /// # use provider::NativeProvider;
526    /// # use support::{fs::local::LocalFileSystem};
527    /// # use zombienet_orchestrator::{errors, AddCollatorOptions, Orchestrator};
528    /// # use configuration::NetworkConfig;
529    /// # async fn example() -> Result<(), anyhow::Error> {
530    /// #   let provider = NativeProvider::new(LocalFileSystem {});
531    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
532    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
533    /// let mut network = orchestrator.spawn(config).await?;
534    /// let para_config = network
535    ///     .para_config_builder()
536    ///     .with_id(100)
537    ///     .with_default_command("polkadot-parachain")
538    ///     .with_collator(|c| c.with_name("col-100-1"))
539    ///     .build()
540    ///     .map_err(|_e| anyhow!("Building config"))?;
541    ///
542    /// network.add_parachain(&para_config, None, None).await?;
543    ///
544    /// #   Ok(())
545    /// # }
546    /// ```
547    pub async fn add_parachain(
548        &mut self,
549        para_config: &ParachainConfig,
550        custom_relaychain_spec: Option<PathBuf>,
551        custom_parchain_fs_prefix: Option<String>,
552    ) -> Result<(), anyhow::Error> {
553        let base_dir = self.ns.base_dir().to_string_lossy().to_string();
554        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
555
556        let mut global_files_to_inject = vec![];
557
558        // get relaychain id
559        let relay_chain_id = if let Some(custom_path) = custom_relaychain_spec {
560            // use this file as relaychain spec
561            global_files_to_inject.push(TransferedFile::new(
562                custom_path.clone(),
563                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
564            ));
565            let content = std::fs::read_to_string(custom_path)?;
566            ChainSpec::chain_id_from_spec(&content)?
567        } else {
568            global_files_to_inject.push(TransferedFile::new(
569                PathBuf::from(format!(
570                    "{}/{}",
571                    scoped_fs.base_dir,
572                    self.relaychain().chain_spec_path.to_string_lossy()
573                )),
574                PathBuf::from(format!("/cfg/{}.json", self.relaychain().chain)),
575            ));
576            self.relaychain().chain_id.clone()
577        };
578
579        let mut para_spec = network_spec::parachain::ParachainSpec::from_config(
580            para_config,
581            relay_chain_id.as_str().try_into()?,
582        )?;
583
584        let chain_spec_raw_path = para_spec
585            .build_chain_spec(
586                &relay_chain_id,
587                &self.ns,
588                &scoped_fs,
589                para_spec.post_process_script.clone().as_deref(),
590            )
591            .await?;
592
593        // Para artifacts
594        let para_path_prefix = if let Some(custom_prefix) = custom_parchain_fs_prefix {
595            custom_prefix
596        } else {
597            para_spec.id.to_string()
598        };
599
600        scoped_fs.create_dir(&para_path_prefix).await?;
601        // create wasm/state
602        para_spec
603            .genesis_state
604            .build(
605                chain_spec_raw_path.as_ref(),
606                format!("{}/genesis-state", para_path_prefix),
607                &self.ns,
608                &scoped_fs,
609                None,
610            )
611            .await?;
612        para_spec
613            .genesis_wasm
614            .build(
615                chain_spec_raw_path.as_ref(),
616                format!("{}/para_spec-wasm", para_path_prefix),
617                &self.ns,
618                &scoped_fs,
619                None,
620            )
621            .await?;
622
623        let parachain =
624            Parachain::from_spec(&para_spec, &global_files_to_inject, &scoped_fs).await?;
625        let parachain_id = parachain.chain_id.clone();
626
627        let resolved_db_snapshots = generators::resolve_db_snapshots(
628            para_spec.collators.iter(),
629            &self.ns,
630            &self.filesystem,
631        )
632        .await?;
633
634        // Create `ctx` for spawn the nodes
635        let ctx_para = SpawnNodeCtx {
636            parachain: Some(&para_spec),
637            parachain_id: parachain_id.as_deref(),
638            role: if para_spec.is_cumulus_based {
639                ZombieRole::CumulusCollator
640            } else {
641                ZombieRole::Collator
642            },
643            bootnodes_addr: &para_config
644                .bootnodes_addresses()
645                .iter()
646                .map(|&a| a.to_string())
647                .collect(),
648            chain_id: &self.relaychain().chain_id,
649            chain: &self.relaychain().chain,
650            ns: &self.ns,
651            scoped_fs: &scoped_fs,
652            wait_ready: false,
653            nodes_by_name: self.nodes_json()?,
654            global_settings: &self.initial_spec.global_settings,
655            resolved_db_snapshots: &resolved_db_snapshots,
656        };
657
658        // Register the parachain to the running network
659        let first_node_url = self
660            .relaychain()
661            .nodes
662            .first()
663            .ok_or(anyhow::anyhow!(
664                "At least one node of the relaychain should be running"
665            ))?
666            .ws_uri();
667
668        if para_config.registration_strategy() == Some(&RegistrationStrategy::UsingExtrinsic) {
669            let register_para_options = RegisterParachainOptions {
670                id: parachain.para_id,
671                // This needs to resolve correctly
672                wasm_path: para_spec
673                    .genesis_wasm
674                    .artifact_path()
675                    .ok_or(anyhow::anyhow!(
676                        "artifact path for wasm must be set at this point",
677                    ))?
678                    .to_path_buf(),
679                state_path: para_spec
680                    .genesis_state
681                    .artifact_path()
682                    .ok_or(anyhow::anyhow!(
683                        "artifact path for state must be set at this point",
684                    ))?
685                    .to_path_buf(),
686                node_ws_url: first_node_url.to_string(),
687                onboard_as_para: para_spec.onboard_as_parachain,
688                seed: None, // TODO: Seed is passed by?
689                finalization: false,
690            };
691
692            Parachain::register(register_para_options, &scoped_fs).await?;
693        }
694
695        // Spawn the nodes
696        let spawning_tasks = para_spec
697            .collators
698            .iter()
699            .map(|node| spawner::spawn_node(node, parachain.files_to_inject.clone(), &ctx_para));
700
701        let running_nodes = futures::future::try_join_all(spawning_tasks).await?;
702
703        // Let's make sure nodes are up before adding them
704        let waiting_tasks = running_nodes.iter().map(|node| {
705            node.wait_until_is_up(self.initial_spec.global_settings.network_spawn_timeout())
706        });
707
708        let _ = futures::future::try_join_all(waiting_tasks).await?;
709
710        let running_para_id = parachain.para_id;
711        self.add_para(parachain);
712        for node in running_nodes {
713            self.add_running_node(node, Some(running_para_id)).await;
714        }
715
716        // Dump zombie.json
717        self.write_zombie_json().await?;
718
719        generators::cleanup_db_snapshot_cache(&resolved_db_snapshots).await;
720
721        Ok(())
722    }
723
724    /// Register a parachain, which has already been added to the network (with manual registration
725    /// strategy)
726    ///
727    /// # Arguments
728    /// * `para_id` - Parachain Id
729    ///
730    ///
731    /// # Example:
732    /// ```rust
733    /// # use anyhow::anyhow;
734    /// # use provider::NativeProvider;
735    /// # use support::{fs::local::LocalFileSystem};
736    /// # use zombienet_orchestrator::Orchestrator;
737    /// # use configuration::{NetworkConfig, NetworkConfigBuilder, RegistrationStrategy};
738    /// # async fn example() -> Result<(), anyhow::Error> {
739    /// #   let provider = NativeProvider::new(LocalFileSystem {});
740    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
741    /// #   let config = NetworkConfigBuilder::new()
742    /// #     .with_relaychain(|r| {
743    /// #       r.with_chain("rococo-local")
744    /// #         .with_default_command("polkadot")
745    /// #         .with_node(|node| node.with_name("alice"))
746    /// #     })
747    /// #     .with_parachain(|p| {
748    /// #       p.with_id(100)
749    /// #         .with_registration_strategy(RegistrationStrategy::Manual)
750    /// #         .with_default_command("test-parachain")
751    /// #         .with_collator(|n| n.with_name("dave").validator(false))
752    /// #     })
753    /// #     .build()
754    /// #     .map_err(|_e| anyhow!("Building config"))?;
755    /// let mut network = orchestrator.spawn(config).await?;
756    ///
757    /// network.register_parachain(100).await?;
758    ///
759    /// #   Ok(())
760    /// # }
761    /// ```
762    pub async fn register_parachain(&mut self, para_id: u32) -> Result<(), anyhow::Error> {
763        let para = self
764            .initial_spec
765            .parachains
766            .iter()
767            .find(|p| p.id == para_id)
768            .ok_or(anyhow::anyhow!(
769                "no parachain with id = {para_id} available",
770            ))?;
771        let para_genesis_config = para.get_genesis_config()?;
772        let first_node_url = self
773            .relaychain()
774            .nodes
775            .first()
776            .ok_or(anyhow::anyhow!(
777                "At least one node of the relaychain should be running"
778            ))?
779            .ws_uri();
780        let register_para_options: RegisterParachainOptions = RegisterParachainOptions {
781            id: para_id,
782            // This needs to resolve correctly
783            wasm_path: para_genesis_config.wasm_path.clone(),
784            state_path: para_genesis_config.state_path.clone(),
785            node_ws_url: first_node_url.to_string(),
786            onboard_as_para: para_genesis_config.as_parachain,
787            seed: None, // TODO: Seed is passed by?
788            finalization: false,
789        };
790        let base_dir = self.ns.base_dir().to_string_lossy().to_string();
791        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
792        Parachain::register(register_para_options, &scoped_fs).await?;
793
794        Ok(())
795    }
796
797    // deregister and stop the collator?
798    // remove_parachain()
799
800    /// Get a node of a specific kind by name, downcasting it from the registry.
801    fn get_node_as<'a, N: SpawnedNode>(&'a self, name: &str) -> Result<&'a N, anyhow::Error> {
802        // Node names are sanitized when configured (`Collator-1000` runs as
803        // `collator-1000`), so a lookup by the name as it was written still
804        // finds the node.
805        let node = self
806            .nodes_by_name
807            .get(name)
808            .or_else(|| self.nodes_by_name.get(&sanitize_node_name(name)));
809        let node = node.ok_or_else(|| {
810            anyhow::anyhow!(
811                "can't find node with name: {name:?}, should be one of {}",
812                self.node_names().join(", ")
813            )
814        })?;
815
816        // `dyn SpawnedNode` upcasts to `dyn Any` (Rust >= 1.86), no `as_any` needed.
817        let node_any: &dyn std::any::Any = node.as_ref();
818        node_any.downcast_ref::<N>().ok_or_else(|| {
819            anyhow::anyhow!(
820                "node {name:?} is a '{}' node, it can't be used as a {}",
821                node.kind(),
822                std::any::type_name::<N>()
823            )
824        })
825    }
826
827    /// Get a substrate node (relaychain node or collator) by name.
828    pub fn get_node(&self, name: impl Into<String>) -> Result<&NetworkNode, anyhow::Error> {
829        self.get_node_as::<NetworkNode>(&name.into())
830    }
831
832    /// Get a JAM node by name.
833    pub fn get_jam_node(&self, name: impl Into<String>) -> Result<&JamNetworkNode, anyhow::Error> {
834        self.get_node_as::<JamNetworkNode>(&name.into())
835    }
836
837    /// Get a custom process by name.
838    pub fn get_custom_process(
839        &self,
840        name: impl Into<String>,
841    ) -> Result<&node::CustomProcessNode, anyhow::Error> {
842        self.get_node_as::<node::CustomProcessNode>(&name.into())
843    }
844
845    /// Whether a node or custom process of that name is on the network.
846    pub(crate) fn has_member(&self, name: &str) -> bool {
847        self.nodes_by_name.contains_key(name)
848    }
849
850    /// The running nodes as the spawner's placeholder context: name to the
851    /// node's serialized form (`ws_uri`, `multiaddr`, ...).
852    pub(crate) fn nodes_json(&self) -> Result<serde_json::Value, serde_json::Error> {
853        serde_json::to_value(&self.nodes_by_name)
854    }
855
856    /// The custom processes running next to the nodes, in the order each one
857    /// met its condition. Each carries its
858    /// [`CustomProcessState`](node::CustomProcessState).
859    pub fn custom_processes(&self) -> Vec<&node::CustomProcessNode> {
860        self.custom_processes.iter().map(|p| p.as_ref()).collect()
861    }
862
863    /// Get any node by name, without caring about its kind.
864    ///
865    /// Only the behaviour shared by every node ([`SpawnedNode`]) is available
866    /// on the returned reference.
867    pub fn get_any_node(&self, name: impl Into<String>) -> Result<&dyn SpawnedNode, anyhow::Error> {
868        let name = name.into();
869        self.nodes_by_name
870            .get(&name)
871            .map(|node| node.as_ref())
872            .ok_or_else(|| {
873                anyhow::anyhow!(
874                    "can't find node with name: {name:?}, should be one of {}",
875                    self.node_names().join(", ")
876                )
877            })
878    }
879
880    pub fn node_names(&self) -> Vec<String> {
881        self.nodes_by_name.keys().cloned().collect::<Vec<_>>()
882    }
883
884    /// All the substrate nodes (relaychain nodes and collators) of the network.
885    pub fn nodes(&self) -> Vec<&NetworkNode> {
886        self.nodes_iter().collect()
887    }
888
889    /// All the nodes of the network, of every kind.
890    pub fn all_nodes(&self) -> Vec<&dyn SpawnedNode> {
891        self.nodes_by_name.values().map(|n| n.as_ref()).collect()
892    }
893
894    /// All the nodes of the given kind.
895    pub fn nodes_of_kind(&self, kind: NodeKind) -> Vec<&dyn SpawnedNode> {
896        self.nodes_by_name
897            .values()
898            .filter(|n| n.kind() == kind)
899            .map(|n| n.as_ref())
900            .collect()
901    }
902
903    pub async fn detach(&self) {
904        self.ns.detach().await
905    }
906
907    // Internal API
908    pub(crate) async fn add_running_node(&mut self, node: NetworkNode, para_id: Option<u32>) {
909        let node = Arc::new(node);
910        if let Some(para_id) = para_id {
911            if let Some(para) = self.parachains.get_mut(&para_id).and_then(|p| p.get_mut(0)) {
912                para.collators.push(node.clone());
913            } else {
914                // is the first node of the para, let create the entry
915                unreachable!()
916            }
917        } else {
918            self.relaychain_mut().nodes.push(node.clone());
919        }
920
921        self.register_running_node(node).await;
922    }
923
924    /// Add an already spawned JAM node to the jamchain and to the registry.
925    pub(crate) async fn add_running_jam_node(&mut self, node: JamNetworkNode) {
926        let node = Arc::new(node);
927        if let Some(jamchain) = self.jamchain.as_mut() {
928            jamchain.nodes.push(node.clone());
929        } else {
930            // the jamchain should be set before adding any of its nodes
931            unreachable!()
932        }
933
934        self.register_running_node(node).await;
935    }
936
937    /// Add an already spawned custom process to the network and to the registry.
938    ///
939    /// Registered for lookup but not for monitoring: a custom process has no
940    /// Prometheus endpoint, and a declared port that is not meant to answer a
941    /// plain TCP probe must not read as a crash that tears the network down.
942    pub(crate) fn add_running_custom_process(&mut self, process: node::CustomProcessNode) {
943        let process = Arc::new(process);
944        self.custom_processes.push(process.clone());
945        self.register(process);
946    }
947
948    /// Mark a freshly spawned node as running and register it (as the very
949    /// same `Arc` the typed collection holds) for lookup and monitoring.
950    async fn register_running_node(&mut self, node: Arc<impl SpawnedNode>) {
951        let node = self.register(node);
952        self.nodes_to_watch.write().await.push(node);
953    }
954
955    /// Mark `node` as running now and register it for lookup by name.
956    fn register(&mut self, node: Arc<impl SpawnedNode>) -> Arc<dyn SpawnedNode> {
957        node.core().set_is_running(true);
958        node.core().set_last_start_ts(
959            SystemTime::now()
960                .duration_since(UNIX_EPOCH)
961                .expect("Timestamp should be valid")
962                .as_secs(),
963        );
964
965        let node: Arc<dyn SpawnedNode> = node;
966        self.nodes_by_name
967            .insert(node.name().to_string(), node.clone());
968        node
969    }
970
971    pub(crate) fn add_para(&mut self, para: Parachain) {
972        self.parachains.entry(para.para_id).or_default().push(para);
973    }
974
975    pub(crate) async fn write_zombie_json(&self) -> Result<(), anyhow::Error> {
976        let base_dir = self.ns.base_dir().to_string_lossy();
977        let scoped_fs = ScopedFilesystem::new(&self.filesystem, &base_dir);
978        let ns_name = self.ns.name();
979
980        write_zombie_json(serde_json::to_value(self)?, scoped_fs, ns_name).await?;
981        Ok(())
982    }
983
984    pub fn name(&self) -> &str {
985        self.ns.name()
986    }
987
988    /// Get a first parachain from the list of the parachains with specified id.
989    /// NOTE!
990    /// Usually the list will contain only one parachain.
991    /// Multiple parachains with the same id is a corner case.
992    /// If this is the case then one can get such parachain with
993    /// `parachain_by_unique_id()` method
994    ///
995    /// # Arguments
996    /// * `para_id` - Parachain Id
997    pub fn parachain(&self, para_id: u32) -> Option<&Parachain> {
998        self.parachains.get(&para_id)?.first()
999    }
1000
1001    /// Get a parachain by its unique id.
1002    ///
1003    /// This is particularly useful if there are multiple parachains
1004    /// with the same id (this is a rare corner case).
1005    ///
1006    /// # Arguments
1007    /// * `unique_id` - unique id of the parachain
1008    pub fn parachain_by_unique_id(&self, unique_id: impl AsRef<str>) -> Option<&Parachain> {
1009        self.parachains
1010            .values()
1011            .flat_map(|p| p.iter())
1012            .find(|p| p.unique_id == unique_id.as_ref())
1013    }
1014
1015    pub fn parachains(&self) -> Vec<&Parachain> {
1016        self.parachains.values().flatten().collect()
1017    }
1018
1019    pub(crate) fn nodes_iter(&self) -> impl Iterator<Item = &NetworkNode> {
1020        self.relay
1021            .iter()
1022            .flat_map(|relay| relay.nodes.iter())
1023            .chain(
1024                self.parachains
1025                    .values()
1026                    .flat_map(|p| p.iter())
1027                    .flat_map(|p| &p.collators),
1028            )
1029            .map(|n| n.as_ref())
1030    }
1031
1032    /// Waits given number of seconds until all nodes in the network report that they are
1033    /// up and running.
1034    ///
1035    /// # Arguments
1036    /// * `timeout_secs` - The number of seconds to wait.
1037    ///
1038    /// # Returns
1039    /// * `Ok()` if the node is up before timeout occured.
1040    /// * `Err(e)` if timeout or other error occurred while waiting.
1041    pub async fn wait_until_is_up(&self, timeout_secs: u64) -> Result<(), anyhow::Error> {
1042        let handles = self
1043            .nodes_by_name
1044            .values()
1045            .map(|node| node.wait_until_is_up(timeout_secs));
1046
1047        futures::future::try_join_all(handles).await?;
1048
1049        Ok(())
1050    }
1051
1052    /// Pause every node in the network (SIGSTOP). Issued in parallel.
1053    ///
1054    /// Note: with the native provider on an attached network, the live
1055    /// network has to be running with global setting `teardown_on_failure`
1056    /// disabled.
1057    pub async fn pause(&self) -> Result<(), anyhow::Error> {
1058        futures::future::try_join_all(self.nodes_iter().map(|n| n.pause())).await?;
1059        Ok(())
1060    }
1061
1062    /// Resume every node in the network (SIGCONT). Issued in parallel.
1063    ///
1064    /// Note: with the native provider on an attached network, the live
1065    /// network has to be running with global setting `teardown_on_failure`
1066    /// disabled.
1067    pub async fn resume(&self) -> Result<(), anyhow::Error> {
1068        futures::future::try_join_all(self.nodes_iter().map(|n| n.resume())).await?;
1069        Ok(())
1070    }
1071
1072    /// Start the observability stack (Prometheus + Grafana) as an add-on
1073    ///
1074    /// This can be called on any running network — whether freshly spawned or
1075    /// re-attached via [`Orchestrator::attach_to_live`]. If observability is
1076    /// already running, it will be stopped first
1077    ///
1078    /// # Example:
1079    /// ```rust
1080    /// # use provider::NativeProvider;
1081    /// # use support::{fs::local::LocalFileSystem};
1082    /// # use zombienet_orchestrator::{errors, Orchestrator};
1083    /// # use configuration::{NetworkConfig, ObservabilityConfigBuilder};
1084    /// # async fn example() -> Result<(), errors::OrchestratorError> {
1085    /// #   let provider = NativeProvider::new(LocalFileSystem {});
1086    /// #   let orchestrator = Orchestrator::new(LocalFileSystem {}, provider);
1087    /// #   let config = NetworkConfig::load_from_toml("config.toml")?;
1088    /// let mut network = orchestrator.spawn(config).await?;
1089    ///
1090    /// let obs_config = ObservabilityConfigBuilder::new()
1091    ///     .with_enabled(true)
1092    ///     .with_grafana_port(3000)
1093    ///     .build();
1094    ///
1095    /// let info = network.start_observability(&obs_config).await?;
1096    /// println!("Grafana: {}", info.grafana_url);
1097    /// #   Ok(())
1098    /// # }
1099    /// ```
1100    pub async fn start_observability(
1101        &mut self,
1102        config: &configuration::ObservabilityConfig,
1103    ) -> Result<&ObservabilityInfo, anyhow::Error> {
1104        if self.observability().is_some() {
1105            self.stop_observability().await?;
1106        }
1107
1108        let nodes = self.nodes();
1109        let info = observability::spawn_observability_stack(
1110            config,
1111            &nodes,
1112            self.ns.name(),
1113            self.ns.base_dir(),
1114            &self.filesystem,
1115        )
1116        .await?;
1117
1118        self.observability = ObservabilityState::Running(info);
1119        self.observability()
1120            .ok_or_else(|| anyhow::anyhow!("observability state was just set but is not running"))
1121    }
1122
1123    /// Stop the observability stack if running
1124    ///
1125    /// Removes the Prometheus and Grafana containers. This is safe to call
1126    /// even if no observability stack is running (it will be a no-op)
1127    pub async fn stop_observability(&mut self) -> Result<(), anyhow::Error> {
1128        if let ObservabilityState::Running(info) =
1129            std::mem::replace(&mut self.observability, ObservabilityState::Stopped)
1130        {
1131            observability::cleanup_observability_stack(&info).await?;
1132        }
1133        Ok(())
1134    }
1135
1136    pub(crate) fn spawn_watching_task(&self) {
1137        let nodes_to_watch = Arc::clone(&self.nodes_to_watch);
1138        let ns = Arc::clone(&self.ns);
1139        let node_bootstrap_timeout = self.initial_spec.global_settings.node_spawn_timeout();
1140
1141        tokio::spawn(async move {
1142            loop {
1143                tokio::time::sleep(Duration::from_secs(NODE_MONITORING_INTERVAL_SECONDS)).await;
1144
1145                let guard = nodes_to_watch.read().await;
1146                let nodes = guard.iter().filter(|n| n.is_running()).collect::<Vec<_>>();
1147
1148                let all_running = {
1149                    let all_running =
1150                        futures::future::try_join_all(nodes.iter().map(|n| {
1151                            n.wait_until_is_up(NODE_MONITORING_FAILURE_THRESHOLD_SECONDS)
1152                        }))
1153                        .await;
1154
1155                    // Re-check `is_running` to make sure we don't kill the network unnecessarily
1156                    if nodes.iter().any(|n| !n.is_running()) {
1157                        continue;
1158                    } else {
1159                        all_running
1160                    }
1161                };
1162
1163                if let Err(e) = all_running {
1164                    // check if the node was restarted and we need to give it more time
1165                    if let Some(node) = nodes.iter().find(|n| n.name() == e.to_string()) {
1166                        let now = SystemTime::now()
1167                            .duration_since(UNIX_EPOCH)
1168                            .expect("get current ts should work.")
1169                            .as_secs();
1170                        if node_bootstrap_timeout as u64 > (now - node.last_start_ts()) {
1171                            warn!("[{}] still in bootstrap window from last starting ({}), continue waiting...", node.name(), node.last_start_ts());
1172                            continue;
1173                        }
1174                    }
1175
1176                    warn!("\n\t🧟 One of the nodes crashed: {e}. tearing the network down...");
1177
1178                    if let Err(e) = ns.destroy().await {
1179                        error!("an error occurred during network teardown: {}", e);
1180                    }
1181
1182                    std::process::exit(1);
1183                }
1184            }
1185        });
1186    }
1187
1188    pub(crate) fn set_parachains(&mut self, parachains: HashMap<u32, Vec<Parachain>>) {
1189        self.parachains = parachains;
1190    }
1191
1192    pub(crate) fn insert_node(&mut self, node: Arc<dyn SpawnedNode>) {
1193        self.nodes_by_name.insert(node.name().to_string(), node);
1194    }
1195
1196    pub(crate) fn set_jamchain(&mut self, jamchain: Jamchain) {
1197        self.jamchain = Some(jamchain);
1198    }
1199
1200    /// Set the custom processes recreated from a `zombie.json`, registering
1201    /// each for lookup (but not monitoring, see `add_running_custom_process`).
1202    pub(crate) fn set_custom_processes(&mut self, processes: Vec<Arc<node::CustomProcessNode>>) {
1203        for process in &processes {
1204            self.nodes_by_name
1205                .insert(process.name().to_string(), process.clone());
1206        }
1207        self.custom_processes = processes;
1208    }
1209
1210    /// The JAM chain of this network, if any.
1211    pub fn jamchain(&self) -> Option<&Jamchain> {
1212        self.jamchain.as_ref()
1213    }
1214
1215    pub(crate) fn set_start_time_ts(&mut self, start_time: SystemTime) {
1216        if let Ok(start_time_ts) = start_time.duration_since(SystemTime::UNIX_EPOCH) {
1217            self.start_time_ts = Some(start_time_ts.as_millis().to_string());
1218        } else {
1219            // Just warn, do not propagate the err (this should not happens)
1220            warn!("⚠️ Error getting start_time timestamp");
1221        }
1222    }
1223}