feat(testing): make sure scenario is saved during crash and ctrl-c

Author: binier

node/testing/Cargo.toml

@@ -27,24 +27,25 @@ axum = "0.6"
 tower-http = { version = "0.4.4", features = ["cors"] }
 strum = "0.24.1"
 strum_macros = "0.24.3"
+tracing-log = "0.2.0"
 documented = { version = "0.1", optional = true }
 redux = { git = "https://github.com/openmina/redux-rs.git", branch="feat/global-time", features = ["serde"] }
 ledger = { workspace = true }
 mina-p2p-messages = { workspace = true }
 libp2p = { workspace = true, features = ["macros", "serde", "tcp", "dns", "tokio", "yamux", "pnet", "noise", "gossipsub"] }
 vrf = { workspace = true }
 
+nix = { version = "0.27.1", features = ["process", "signal"] }
+ctrlc = "3.4.2"
+temp-dir = "0.1.11"
 console = "0.15.5"
 clap = { version = "4.3", features = [ "derive", "env" ] }
+reqwest = { version = "0.11.22", features = ["blocking", "json"] }
 
 openmina-core = { path = "../../core" }
 node = { path = "../../node" }
 openmina-node-invariants = { path = "../../node/invariants" }
 openmina-node-native = { path = "../../node/native" }
-reqwest = { version = "0.11.22", features = ["blocking", "json"] }
-temp-dir = "0.1.11"
-nix = { version = "0.27.1", features = ["process", "signal"] }
-tracing-log = "0.2.0"
 
 [features]
 default = ["scenario-generators"]

node/testing/src/cluster/mod.rs

@@ -11,8 +11,8 @@ use std::collections::BTreeMap;
 use std::time::Duration;
 use std::{collections::VecDeque, sync::Arc};
 
-use libp2p::futures::{stream::FuturesUnordered, StreamExt};
 use ledger::proofs::{VerifierIndex, VerifierSRS};
+use libp2p::futures::{stream::FuturesUnordered, StreamExt};
 use node::core::channels::mpsc;
 use node::core::requests::RpcId;
 use node::{
@@ -270,6 +270,9 @@ impl Cluster {
 
         let state = node::State::new(config);
         fn effects(store: &mut node::Store<NodeTestingService>, action: node::ActionWithMeta) {
+            // if action.action().kind().to_string().starts_with("BlockProducer") {
+            //     dbg!(action.action());
+            // }
             store.service.dyn_effects(store.state.get(), &action);
             let peer_id = store.state().p2p.my_id();
             openmina_core::log::trace!(action.time(); "{peer_id}: {:?}", action.action().kind());

node/testing/src/main.rs

@@ -39,6 +39,19 @@ impl Command {
         let rt = setup();
         let _rt_guard = rt.enter();
 
+        let (shutdown_tx, shutdown_rx) = openmina_core::channels::oneshot::channel();
+        let mut shutdown_tx = Some(shutdown_tx);
+
+        ctrlc::set_handler(move || match shutdown_tx.take() {
+            Some(tx) => {
+                let _ = tx.send(());
+            }
+            None => {
+                std::process::exit(1);
+            }
+        })
+        .expect("Error setting Ctrl-C handler");
+
         match self {
             Self::Server(args) => {
                 server(args.port);
@@ -56,22 +69,32 @@ impl Command {
                         config
                     };
 
-                    if let Some(name) = cmd.name {
-                        if let Some(scenario) = Scenarios::iter()
-                            .into_iter()
-                            .find(|s| <&'static str>::from(s) == name)
-                        {
-                            rt.block_on(scenario.run_and_save_from_scratch(config));
+                    let fut = async move {
+                        if let Some(name) = cmd.name {
+                            if let Some(scenario) = Scenarios::iter()
+                                .into_iter()
+                                .find(|s| <&'static str>::from(s) == name)
+                            {
+                                scenario.run_and_save_from_scratch(config).await;
+                            } else {
+                                anyhow::bail!("no such scenario: \"{name}\"");
+                            }
                         } else {
-                            anyhow::bail!("no such scenario: \"{name}\"");
+                            for scenario in Scenarios::iter() {
+                                scenario.run_and_save_from_scratch(config.clone()).await;
+                            }
                         }
-                    } else {
-                        for scenario in Scenarios::iter() {
-                            rt.block_on(scenario.run_and_save_from_scratch(config.clone()));
-                        }
-                    }
+                        Ok(())
+                    };
 
-                    Ok(())
+                    rt.block_on(async {
+                        tokio::select! {
+                            res = fut => res,
+                            _ = shutdown_rx => {
+                                anyhow::bail!("Received ctrl-c signal! shutting down...");
+                            }
+                        }
+                    })
                 }
                 #[cfg(not(feature = "scenario-generators"))]
                 Err("binary not compiled with `scenario-generators` feature"

node/testing/src/scenario/mod.rs

@@ -101,4 +101,12 @@ impl Scenario {
         tokio::fs::write(&tmp_file, encoded).await?;
         Ok(tokio::fs::rename(tmp_file, self.file_path()).await?)
     }
+
+    pub fn save_sync(&self) -> Result<(), anyhow::Error> {
+        let tmp_file = self.tmp_file_path();
+        let encoded = serde_json::to_vec_pretty(self)?;
+        std::fs::create_dir_all(Self::PATH)?;
+        std::fs::write(&tmp_file, encoded)?;
+        Ok(std::fs::rename(tmp_file, self.file_path())?)
+    }
 }

node/testing/src/scenarios/mod.rs

@@ -140,14 +140,30 @@ impl Scenarios {
     }
 
     pub async fn run_and_save(self, cluster: &mut Cluster) {
+        struct ScenarioSaveOnExit(Scenario);
+
+        impl Drop for ScenarioSaveOnExit {
+            fn drop(&mut self) {
+                let info = self.0.info.clone();
+                let steps = std::mem::take(&mut self.0.steps);
+                let scenario = Scenario { info, steps };
+
+                eprintln!("saving scenario({}) before exit...", scenario.info.id);
+                if let Err(err) = scenario.save_sync() {
+                    eprintln!(
+                        "failed to save scenario({})! error: {}",
+                        scenario.info.id, err
+                    );
+                }
+            }
+        }
+
         eprintln!("run_and_save: {}", self.to_str());
-        let mut scenario = self.blank_scenario();
-        self.run(cluster, |step| scenario.add_step(step.clone()).unwrap())
+        let mut scenario = ScenarioSaveOnExit(self.blank_scenario());
+        self.run(cluster, |step| scenario.0.add_step(step.clone()).unwrap())
             .await;
-        scenario
-            .save()
-            .await
-            .expect("failed to save scenario after run");
+        // drop to save it.
+        let _ = scenario;
     }
 
     pub async fn run_only(self, cluster: &mut Cluster) {