add runtime_upgrade for releaychain (#259)

Added:
  - Helper to perform a runtime upgrade on the relaychain
 
```rust
    let wasm = "file_path__or__url_to_*.compact.compressed.wasm";
    network
        .relaychain()
        .runtime_upgrade(RuntimeUpgradeOptions::new(wasm.into()))
        .await?;
```

crates/configuration/Cargo.toml

@@ -20,6 +20,8 @@ anyhow = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 toml = { workspace = true }
 serde_json = { workspace = true }
+reqwest = { workspace = true }
+tokio = { workspace = true, features = ["fs"] }
 
 # zombienet deps
 support = { workspace = true }

crates/configuration/src/shared/types.rs

@@ -5,6 +5,7 @@ use std::{
     str::FromStr,
 };
 
+use anyhow::anyhow;
 use lazy_static::lazy_static;
 use regex::Regex;
 use serde::{de, Deserialize, Deserializer, Serialize};
@@ -267,6 +268,35 @@ impl Display for AssetLocation {
     }
 }
 
+impl AssetLocation {
+    pub async fn get_asset(&self) -> Result<Vec<u8>, anyhow::Error> {
+        let contents = match self {
+            AssetLocation::Url(location) => {
+                let res = reqwest::get(location.as_ref()).await.map_err(|err| {
+                    anyhow!(
+                        "Error dowinloding asset from url {} - {}",
+                        location,
+                        err.to_string()
+                    )
+                })?;
+
+                res.bytes().await.unwrap().into()
+            },
+            AssetLocation::FilePath(filepath) => {
+                tokio::fs::read(filepath).await.map_err(|err| {
+                    anyhow!(
+                        "Error reading asset from path {} - {}",
+                        filepath.to_string_lossy(),
+                        err.to_string()
+                    )
+                })?
+            },
+        };
+
+        Ok(contents)
+    }
+}
+
 impl Serialize for AssetLocation {
     fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
     where

crates/orchestrator/src/lib.rs

@@ -5,6 +5,8 @@ pub mod errors;
 pub mod generators;
 pub mod network;
 pub mod network_helper;
+pub mod tx_helper;
+
 mod network_spec;
 #[cfg(feature = "pjs")]
 pub mod pjs_helper;

crates/orchestrator/src/network/node.rs

@@ -36,6 +36,18 @@ pub struct NetworkNode {
     metrics_cache: Arc<RwLock<MetricMap>>,
 }
 
+// #[derive(Clone, Debug)]
+// pub struct QueryMetricOptions {
+//     use_cache: bool,
+//     treat_not_found_as_zero: bool,
+// }
+
+// impl Default for QueryMetricOptions {
+//     fn default() -> Self {
+//         Self { use_cache: false, treat_not_found_as_zero: true }
+//     }
+// }
+
 impl NetworkNode {
     /// Create a new NetworkNode
     pub(crate) fn new<T: Into<String>>(
@@ -148,7 +160,8 @@ impl NetworkNode {
         let metric_name = metric_name.into();
         // force cache reload
         self.fetch_metrics().await?;
-        self.metric(&metric_name).await
+        // by default we treat not found as 0 (same in v1)
+        self.metric(&metric_name, true).await
     }
 
     /// Assert on a metric value 'by name' from Prometheus (exposed by the node)
@@ -176,13 +189,13 @@ impl NetworkNode {
         predicate: impl Fn(f64) -> bool,
     ) -> Result<bool, anyhow::Error> {
         let metric_name = metric_name.into();
-        let val = self.metric(&metric_name).await?;
+        let val = self.metric(&metric_name, true).await?;
         if predicate(val) {
             Ok(true)
         } else {
             // reload metrics
             self.fetch_metrics().await?;
-            let val = self.metric(&metric_name).await?;
+            let val = self.metric(&metric_name, true).await?;
             trace!("🔎 Current value passed to the predicated: {val}");
             Ok(predicate(val))
         }
@@ -370,12 +383,12 @@ impl NetworkNode {
         Ok(())
     }
 
+    /// Query individual metric by name
     async fn metric(
         &self,
-        metric_name: &str, // treat_not_found_as_zero: bool
+        metric_name: &str,
+        treat_not_found_as_zero: bool,
     ) -> Result<f64, anyhow::Error> {
-        // TODO: allow to pass as arg
-        let treat_not_found_as_zero = true;
         let mut metrics_map = self.metrics_cache.read().await;
         if metrics_map.is_empty() {
             // reload metrics

crates/orchestrator/src/network/relaychain.rs

@@ -1,8 +1,11 @@
-use std::path::PathBuf;
+use std::{path::PathBuf, str::FromStr};
 
+use anyhow::anyhow;
 use serde::Serialize;
+use subxt_signer::{sr25519::Keypair, SecretUri};
 
 use super::node::NetworkNode;
+use crate::{shared::types::RuntimeUpgradeOptions, tx_helper};
 
 #[derive(Debug, Serialize)]
 pub struct Relaychain {
@@ -30,4 +33,47 @@ impl Relaychain {
     pub fn nodes(&self) -> Vec<&NetworkNode> {
         self.nodes.iter().collect()
     }
+
+    /// Perform a runtime upgrade (with sudo)
+    ///
+    /// This call 'System.set_code_without_checks' wrapped in
+    /// 'Sudo.sudo_unchecked_weight'
+    pub async fn runtime_upgrade(
+        &self,
+        options: RuntimeUpgradeOptions,
+    ) -> Result<(), anyhow::Error> {
+        // check if the node is valid first
+        let ws_url = if let Some(node_name) = options.node_name {
+            if let Some(node) = self
+                .nodes()
+                .into_iter()
+                .find(|node| node.name() == node_name)
+            {
+                node.ws_uri()
+            } else {
+                return Err(anyhow!("Node: {} is not part of the relaychain", node_name));
+            }
+        } else {
+            // take the first node
+            if let Some(node) = self.nodes().first() {
+                node.ws_uri()
+            } else {
+                return Err(anyhow!("Relaychain doesn't have any node!"));
+            }
+        };
+
+        let sudo = if let Some(possible_seed) = options.seed {
+            Keypair::from_secret_key(possible_seed)
+                .map_err(|_| anyhow!("seed should return a Keypair"))?
+        } else {
+            let uri = SecretUri::from_str("//Alice")?;
+            Keypair::from_uri(&uri).map_err(|_| anyhow!("'//Alice' should return a Keypair"))?
+        };
+
+        let wasm_data = options.wasm.get_asset().await?;
+
+        tx_helper::runtime_upgrade::upgrade(ws_url, &wasm_data, &sudo).await?;
+
+        Ok(())
+    }
 }

crates/orchestrator/src/shared/types.rs

@@ -73,6 +73,24 @@ pub struct RegisterParachainOptions {
     pub finalization: bool,
 }
 
+pub struct RuntimeUpgradeOptions {
+    /// Location of the wasm file (could be either a local file or an url)
+    pub wasm: AssetLocation,
+    /// Name of the node to use as rpc endpoint
+    pub node_name: Option<String>,
+    /// Seed to use to sign and submit (default to //Alice)
+    pub seed: Option<[u8; 32]>,
+}
+
+impl RuntimeUpgradeOptions {
+    pub fn new(wasm: AssetLocation) -> Self {
+        Self {
+            wasm,
+            node_name: None,
+            seed: None,
+        }
+    }
+}
 #[derive(Debug, Clone)]
 pub struct ParachainGenesisArgs {
     pub genesis_head: String,

crates/orchestrator/src/tx_helper.rs

@@ -1,2 +1,3 @@
-pub mod register_para;
-pub mod validator_actions;
+// pub mod register_para;
+// pub mod validator_actions;
+pub mod runtime_upgrade;

crates/orchestrator/src/tx_helper/runtime_upgrade.rs

@@ -0,0 +1,66 @@
+use subxt::{dynamic::Value, tx::TxStatus, OnlineClient, SubstrateConfig};
+use subxt_signer::sr25519::Keypair;
+use tracing::{debug, info};
+
+pub async fn upgrade(
+    // options: RuntimeUpgradeOptions,
+    ws_url: &str,
+    wasm_data: &[u8],
+    sudo: &Keypair,
+    // scoped_fs: &ScopedFilesystem<'_, impl FileSystem>,
+) -> Result<(), anyhow::Error> {
+    debug!("Upgrading runtime, using {} as endpoint ", ws_url);
+    let api = OnlineClient::<SubstrateConfig>::from_url(ws_url).await?;
+
+    let upgrade = subxt::dynamic::tx(
+        "System",
+        "set_code_without_checks",
+        vec![Value::from_bytes(wasm_data)],
+    );
+
+    let sudo_call = subxt::dynamic::tx(
+        "Sudo",
+        "sudo_unchecked_weight",
+        vec![
+            upgrade.into_value(),
+            Value::named_composite([
+                ("ref_time", Value::primitive(1.into())),
+                ("proof_size", Value::primitive(1.into())),
+            ]),
+        ],
+    );
+
+    let mut tx = api
+        .tx()
+        .sign_and_submit_then_watch_default(&sudo_call, sudo)
+        .await?;
+
+    // Below we use the low level API to replicate the `wait_for_in_block` behaviour
+    // which was removed in subxt 0.33.0. See https://github.com/paritytech/subxt/pull/1237.
+    while let Some(status) = tx.next().await {
+        let status = status?;
+        match &status {
+            TxStatus::InBestBlock(tx_in_block) | TxStatus::InFinalizedBlock(tx_in_block) => {
+                let _result = tx_in_block.wait_for_success().await?;
+                let block_status = if status.as_finalized().is_some() {
+                    "Finalized"
+                } else {
+                    "Best"
+                };
+                info!(
+                    "[{}] In block: {:#?}",
+                    block_status,
+                    tx_in_block.block_hash()
+                );
+            },
+            TxStatus::Error { message }
+            | TxStatus::Invalid { message }
+            | TxStatus::Dropped { message } => {
+                return Err(anyhow::format_err!("Error submitting tx: {message}"));
+            },
+            _ => continue,
+        }
+    }
+
+    Ok(())
+}

crates/sdk/src/lib.rs

@@ -7,6 +7,12 @@ pub use orchestrator::{
     network::{node::NetworkNode, Network},
     AddCollatorOptions, AddNodeOptions, Orchestrator,
 };
+
+// Helpers used for interact with the network
+pub mod tx_helper {
+    pub use orchestrator::shared::types::RuntimeUpgradeOptions;
+}
+
 use provider::{DockerProvider, KubernetesProvider, NativeProvider};
 pub use support::fs::local::LocalFileSystem;