Add migration idempotency checks (#48)

Closes #42 Checks that `on_runtime_upgrade` succeeds when called twice, and if it does succeed after being called twice, the storage state root did not change. --------- Co-authored-by: Oliver Tale-Yazdi <[email protected]>
paritytech · Nov 10, 2023 · f621966 · f621966
1 parent 7f66e57
commit f621966
Show file tree

Hide file tree

Showing 16 changed files with 275 additions and 38 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -7,7 +7,7 @@ members = [
 ]
 
 [workspace.package]
-version = "0.3.5"
+version = "0.4.0"
 authors = ["Parity Technologies <[email protected]>"]
 description = "Substrate's programmatic testing framework."
 edition = "2021"

diff --git a/core/src/commands/execute_block.rs b/core/src/commands/execute_block.rs
@@ -141,6 +141,7 @@ where
 
     let _ = state_machine_call_with_proof::<Block, HostFns>(
         &ext,
+        &mut Default::default(),
         &executor,
         "TryRuntime_execute_block",
         &payload,

diff --git a/core/src/commands/fast_forward.rs b/core/src/commands/fast_forward.rs
@@ -247,6 +247,7 @@ where
         log::info!("Running migrations...");
         state_machine_call_with_proof::<Block, HostFns>(
             &ext,
+            &mut Default::default(),
             &executor,
             "TryRuntime_on_runtime_upgrade",
             command.try_state.encode().as_ref(),

diff --git a/core/src/commands/follow_chain.rs b/core/src/commands/follow_chain.rs
@@ -158,8 +158,10 @@ where
             .as_mut()
             .expect("state_ext either existed or was just created");
 
+        let mut overlayed_changes = Default::default();
         let result = state_machine_call_with_proof::<Block, HostFns>(
             state_ext,
+            &mut overlayed_changes,
             &executor,
             "TryRuntime_execute_block",
             (
@@ -187,9 +189,7 @@ where
             continue;
         }
 
-        let (mut changes, _, _) = result.expect("checked to be Ok; qed");
-
-        let storage_changes = changes
+        let storage_changes = overlayed_changes
             .drain_storage_changes(
                 &state_ext.backend,
                 // Note that in case a block contains a runtime upgrade, state version could

diff --git a/core/src/commands/on_runtime_upgrade.rs b/core/src/commands/on_runtime_upgrade.rs
@@ -57,11 +57,13 @@ pub struct Command {
     )]
     pub checks: UpgradeCheckSelect,
 
-    /// Whether to assume that the runtime is a relay chain runtime.
-    ///
-    /// This is used to adjust the behavior of weight measurement warnings.
+    /// Whether to disable weight warnings, useful if the runtime is for a relay chain.
     #[clap(long, default_value = "false", default_missing_value = "true")]
     pub no_weight_warnings: bool,
+
+    /// Whether to disable migration idempotency checks
+    #[clap(long, default_value = "false", default_missing_value = "true")]
+    pub no_idempotency_checks: bool,
 }
 
 // Runs the `on-runtime-upgrade` command.
@@ -99,15 +101,18 @@ where
         "🔬 Running TryRuntime_on_runtime_upgrade with checks: {:?}",
         command.checks
     );
-    let (_, proof, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
+    // Save the overlayed changes from the first run, so we can use them later for idempotency
+    // checks.
+    let mut overlayed_changes = Default::default();
+    let (proof, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
         &ext,
+        &mut overlayed_changes,
         &executor,
         "TryRuntime_on_runtime_upgrade",
         command.checks.encode().as_ref(),
         Default::default(), // we don't really need any extensions here.
         shared.export_proof.clone(),
     )?;
-
     let ref_time_results = encoded_result.try_into()?;
 
     // If the above call ran with checks then we need to run the call again without checks to
@@ -120,43 +125,93 @@ where
             log::info!(
                 "🔬 TryRuntime_on_runtime_upgrade succeeded! Running it again without checks for weight measurements."
             );
-            let (_, proof, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
+            let (proof, encoded_result) = state_machine_call_with_proof::<Block, HostFns>(
                 &ext,
+                &mut Default::default(),
                 &executor,
                 "TryRuntime_on_runtime_upgrade",
                 UpgradeCheckSelect::None.encode().as_ref(),
                 Default::default(), // we don't really need any extensions here.
-                shared.export_proof,
+                shared.export_proof.clone(),
             )?;
             let ref_time_results = encoded_result.try_into()?;
             (proof, ref_time_results)
         }
     };
 
+    // Check idempotency
+    let idempotency_ok = match command.no_idempotency_checks {
+        true => {
+            log::info!("ℹ Skipping idempotency check");
+            true
+        }
+        false => {
+            log::info!(
+                "🔬 Running TryRuntime_on_runtime_upgrade again to check idempotency: {:?}",
+                command.checks
+            );
+            let (oc_pre_root, _) = overlayed_changes.storage_root(&ext.backend, ext.state_version);
+            match state_machine_call_with_proof::<Block, HostFns>(
+                &ext,
+                &mut overlayed_changes,
+                &executor,
+                "TryRuntime_on_runtime_upgrade",
+                command.checks.encode().as_ref(),
+                Default::default(), // we don't really need any extensions here.
+                shared.export_proof.clone(),
+            ) {
+                Ok(_) => {
+                    // Execution was OK, check if the storage root changed.
+                    let (oc_post_root, _) =
+                        overlayed_changes.storage_root(&ext.backend, ext.state_version);
+                    if oc_pre_root != oc_post_root {
+                        log::error!("❌ Migrations are not idempotent. Selectively remove migrations from Executive until you find the culprit.");
+                        false
+                    } else {
+                        // Exeuction was OK and state root didn't change
+                        true
+                    }
+                }
+                Err(e) => {
+                    log::error!(
+                        "❌ Migrations are not idempotent, they failed during the second execution.",
+                    );
+                    log::debug!("{:?}", e);
+                    false
+                }
+            }
+        }
+    };
+
+    // Check weight safety
     let pre_root = ext.backend.root();
     let pov_safety = analyse_pov::<HashingFor<Block>>(proof, *pre_root, command.no_weight_warnings);
     let ref_time_safety = analyse_ref_time(ref_time_results, command.no_weight_warnings);
-
-    match (pov_safety, ref_time_safety, command.no_weight_warnings) {
+    let weight_ok = match (pov_safety, ref_time_safety, command.no_weight_warnings) {
         (_, _, true) => {
-            log::info!("✅ TryRuntime_on_runtime_upgrade executed without errors")
+            log::info!("ℹ Skipped checking weight safety");
+            true
         }
         (WeightSafety::ProbablySafe, WeightSafety::ProbablySafe, _) => {
             log::info!(
                 target: LOG_TARGET,
-                "✅ TryRuntime_on_runtime_upgrade executed without errors or weight safety \
-                warnings. Please note this does not guarantee a successful runtime upgrade. \
+                "✅ No weight safety issues detected. \
+                Please note this does not guarantee a successful runtime upgrade. \
                 Always test your runtime upgrade with recent state, and ensure that the weight usage \
                 of your migrations will not drastically differ between testing and actual on-chain \
                 execution."
             );
+            true
         }
         _ => {
-            log::warn!(target: LOG_TARGET, "⚠️  TryRuntime_on_runtime_upgrade executed \
-            successfully but with weight safety warnings.");
-            // Exit with a non-zero exit code to indicate that the runtime upgrade may not be safe.
-            std::process::exit(1);
+            log::error!(target: LOG_TARGET, "❌ Weight safety issues detected.");
+            false
         }
+    };
+
+    if !weight_ok || !idempotency_ok {
+        log::error!("❌ Issues detected, exiting non-zero. See logs.");
+        std::process::exit(1);
     }
 
     Ok(())

diff --git a/core/src/lib.rs b/core/src/lib.rs
@@ -152,23 +152,22 @@ impl TryFrom<Vec<u8>> for RefTimeInfo {
 /// Make sure [`LOG_TARGET`] is enabled in logging.
 pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunctions>(
     ext: &TestExternalities<HashingFor<Block>>,
+    storage_overlay: &mut OverlayedChanges<HashingFor<Block>>,
     executor: &WasmExecutor<HostFns>,
     method: &'static str,
     data: &[u8],
     mut extensions: Extensions,
     maybe_export_proof: Option<PathBuf>,
-) -> sc_cli::Result<(OverlayedChanges<HashingFor<Block>>, StorageProof, Vec<u8>)> {
-    let mut changes = Default::default();
-    let backend = ext.backend.clone();
-    let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&backend);
-    let proving_backend = TrieBackendBuilder::wrap(&backend)
+) -> sc_cli::Result<(StorageProof, Vec<u8>)> {
+    let runtime_code_backend = sp_state_machine::backend::BackendRuntimeCode::new(&ext.backend);
+    let proving_backend = TrieBackendBuilder::wrap(&ext.backend)
         .with_recorder(Default::default())
         .build();
     let runtime_code = runtime_code_backend.runtime_code()?;
 
     let encoded_result = StateMachine::new(
         &proving_backend,
-        &mut changes,
+        storage_overlay,
         executor,
         method,
         data,
@@ -210,7 +209,7 @@ pub(crate) fn state_machine_call_with_proof<Block: BlockT, HostFns: HostFunction
             })?;
     }
 
-    Ok((changes, proof, encoded_result))
+    Ok((proof, encoded_result))
 }
 
 /// Converts a [`sp_state_machine::StorageProof`] into a JSON string.

diff --git a/core/tests/create_snapshot.rs b/core/tests/create_snapshot.rs
@@ -25,7 +25,6 @@ use std::{
 use assert_cmd::cargo::cargo_bin;
 use frame_remote_externalities::{Builder, Mode, OfflineConfig, SnapshotConfig};
 use node_primitives::{Block, Hash};
-use regex::Regex;
 use substrate_cli_test_utils as common;
 use tokio::process::Command;
 
@@ -74,15 +73,10 @@ async fn create_snapshot_works() {
         let block_hash = common::block_hash(block_number, &ws_url).await.unwrap();
 
         // Try to create a snapshot.
-        let mut snapshot_creation = create_snapshot(&ws_url, &snap_file_path, block_hash);
+        let child = create_snapshot(&ws_url, &snap_file_path, block_hash);
+        let out = child.wait_with_output().await.unwrap();
 
-        let re = Regex::new(r".*writing snapshot of (\d+) bytes to .*").unwrap();
-        let matched =
-            common::wait_for_stream_pattern_match(snapshot_creation.stderr.take().unwrap(), re)
-                .await;
-
-        // Assert that the snapshot creation succeded.
-        assert!(matched.is_ok(), "Failed to create snapshot");
+        assert!(out.status.success());
 
         let snapshot_is_on_disk = Path::new(&snap_file_path).exists();
         assert!(snapshot_is_on_disk, "Snapshot was not written to disk");

diff --git a/core/tests/execute_block.rs b/core/tests/execute_block.rs
@@ -77,6 +77,9 @@ async fn execute_block_works() {
 
         // Assert that the block-execution process has executed a block.
         assert!(matched.is_ok());
+
+        let out = block_execution.wait_with_output().await.unwrap();
+        assert!(out.status.success());
     })
     .await
 }