diff --git a/crates/katana/core/src/sequencer_error.rs b/crates/katana/core/src/sequencer_error.rs index 2ceb908f45..67dae0e220 100644 --- a/crates/katana/core/src/sequencer_error.rs +++ b/crates/katana/core/src/sequencer_error.rs @@ -42,6 +42,4 @@ pub enum SequencerError { DataUnavailable, #[error("Failed to decode state")] FailedToDecodeStateDump, - #[error("Failed to set storage")] - FailedToSetStorage, } diff --git a/crates/katana/src/args.rs b/crates/katana/src/args.rs index 957c492438..ee3e325392 100644 --- a/crates/katana/src/args.rs +++ b/crates/katana/src/args.rs @@ -29,7 +29,6 @@ pub struct KatanaArgs { pub block_time: Option, #[arg(long)] - #[arg(hide = true)] #[arg(value_name = "PATH")] #[arg(help = "Dump the state of chain on exit to the given file.")] #[arg(long_help = "Dump the state of chain on exit to the given file. If the value is a \ diff --git a/crates/katana/src/main.rs b/crates/katana/src/main.rs index b8da3943cd..d14a7e2a88 100644 --- a/crates/katana/src/main.rs +++ b/crates/katana/src/main.rs @@ -1,13 +1,14 @@ -use std::io; use std::process::exit; use std::sync::Arc; +use std::{fs, io}; use clap::{CommandFactory, Parser}; use clap_complete::{generate, Shell}; use env_logger::Env; -use katana_core::sequencer::KatanaSequencer; +use katana_core::sequencer::{KatanaSequencer, Sequencer}; use katana_rpc::{spawn, KatanaApi, NodeHandle, StarknetApi}; -use log::error; +use log::{error, info}; +use tokio::signal::ctrl_c; use yansi::Paint; mod args; @@ -47,13 +48,17 @@ async fn main() { print_intro( accounts, - config.starknet.seed, + config.starknet.seed.clone(), format!("🚀 JSON-RPC server started: {}", Paint::red(format!("http://{addr}"))), ); } sequencer.start().await; - handle.stopped().await; + + // Wait until Ctrl + C is pressed, then shutdown + ctrl_c().await.unwrap(); + shutdown_handler(sequencer.clone(), config).await; + handle.stop().unwrap(); } Err(err) => { error! {"{}", err}; @@ -104,3 +109,22 @@ ACCOUNTS SEED println!("\n{address}\n\n"); } + +pub async fn shutdown_handler(sequencer: Arc, config: KatanaArgs) { + if let Some(path) = config.dump_state { + info!("Dumping state on shutdown"); + let state = (*sequencer).backend().dump_state().await; + if let Ok(state) = state { + match fs::write(path.clone(), state) { + Ok(_) => { + info!("Successfully dumped state") + } + Err(_) => { + error!("Failed to write state dump to {:?}", path) + } + }; + } else { + error!("Failed to fetch state dump.") + } + }; +}