MrBoec/cli/src/command.rs

use crate::cli::{Cli, Subcommand};
use frame_benchmarking_cli::{BenchmarkCmd, ExtrinsicFactory};
use futures::future::TryFutureExt;
use sc_cli::SubstrateCli;
use service::{
    self, IdentifyVariant,
    benchmarking::{
        benchmark_inherent_data, RemarkBuilder, TransferKeepAliveBuilder
    },
};
use keyring::Sr25519Keyring;

pub use crate::{error::Error, service::BlockId};
#[cfg(feature = "pyroscope")]
use pyroscope_pprofrs::{pprof_backend, PprofConfig};

type Result<T> = std::result::Result<T, Error>;

fn get_exec_name() -> Option<String> {
    std::env::current_exe()
        .ok()
        .and_then(|pb| pb.file_name().map(|s| s.to_os_string()))
        .and_then(|s| s.into_string().ok())
}

impl SubstrateCli for Cli {
    fn impl_name() -> String {
        "Ghost Node".into()
    }

    fn impl_version() -> String {
        env!("SUBSTRATE_CLI_IMPL_VERSION").into()
    }

    fn description() -> String {
        env!("CARGO_PKG_DESCRIPTION").into()
    }

    fn author() -> String {
        env!("CARGO_PKG_AUTHORS").into()
    }

    fn support_url() -> String {
        format!("{}/issues/new", env!("CARGO_PKG_REPOSITORY")).into()
    }

    fn copyright_start_year() -> i32 {
        2024
    }

    fn executable_name() -> String {
        get_exec_name().unwrap_or("ghost".into())
    }

    fn load_spec(&self, id: &str) -> std::result::Result<Box<dyn sc_service::ChainSpec>, String> {
        let id = if id == "" {
            let n = get_exec_name().unwrap_or_default();
            ["casper"]
                .iter()
                .cloned()
                .find(|&chain| n.starts_with(chain))
                .unwrap_or("casper")
        } else {
            id
        };

        Ok(match id {
            #[cfg(feature = "casper-native")]
            "casper" => Box::new(service::chain_spec::casper_config()?),
            #[cfg(feature = "casper-native")]
            "casper-dev" | "dev" | "development" => Box::new(service::chain_spec::casper_development_config()?),
            #[cfg(feature = "casper-native")]
            "casper-local" | "local" => Box::new(service::chain_spec::casper_local_testnet_config()?),
            #[cfg(feature = "casper-native")]
            "casper-staging" | "staging" => Box::new(service::chain_spec::casper_staging_testnet_config()?),
            #[cfg(not(feature = "casper-native"))]
            name if name.starts_with("casper-") && !name.ends_with(".json") =>
                Err(format!("`{}` only supported with `casper-native` feature enabled.", name))?,
            #[cfg(feature = "casper-native")]
            path => {
                let path = std::path::PathBuf::from(path);

                let chain_spec = Box::new(
                   service::GenericChainSpec::from_json_file(path.clone())?
                ) as Box<dyn service::ChainSpec>;

                if chain_spec.is_casper() {
                   Box::new(service::CasperChainSpec::from_json_file(path)?)
                } else {
                   chain_spec
                }
            },
        })
    }
}

fn set_ss58_version(spec: &Box<dyn service::ChainSpec>) {
    use sp_core::crypto::Ss58AddressFormat;

    let ss58_version = if spec.is_ghost() {
        Ss58AddressFormat::custom(1995)
    } else if spec.is_casper() {
        Ss58AddressFormat::custom(1996)
    } else {
        panic!("Invalid network")
    }
    .into();

    sp_core::crypto::set_default_ss58_version(ss58_version);
}

fn run_node_inner<F>(
    cli: Cli,
    logger_hook: F,
) -> Result<()> 
where
    F: FnOnce(&mut sc_cli::LoggerBuilder, &sc_service::Configuration),
{
    let runner = cli
        .create_runner_with_logger_hook::<sc_cli::RunCmd, F>(&cli.run.base, logger_hook)
        .map_err(Error::from)?;
    let chain_spec = &runner.config().chain_spec;

    set_ss58_version(chain_spec);

    runner.run_node_until_exit(move |config| async move {
        let hwbench = (!cli.run.no_hardware_benchmarks)
            .then_some(config.database.path().map(|database_path| {
                let _ = std::fs::create_dir_all(&database_path);
                sc_sysinfo::gather_hwbench(Some(database_path))
            })).flatten();

        let database_source = config.database.clone();
        let task_manager = service::build_full(
            config,
            service::NewFullParams {
                force_authoring_backoff: cli.run.force_authoring_backoff,
                telemetry_worker_handle: None,
                hwbench,
            },
        ).map(|full| full.task_manager)?;

        if let Some(path) = database_source.path() {
            sc_storage_monitor::StorageMonitorService::try_spawn(
                cli.storage_monitor,
                path.to_path_buf(),
                &task_manager.spawn_essential_handle(),
            )?;
        }

        Ok(task_manager)
    })
}

/// Parses ghost specific CLI arguments and run the service.
pub fn run() -> Result<()> {
    let cli: Cli = Cli::from_args();

    #[cfg(feature = "pyroscope")]
    let mut pyroscope_agent_maybe = if let Some(ref agent_addr) = cli.run.pyroscope_server {
        let address = agent_addr
            .to_socket_addrs()
            .map_err(Error::AddressResolutionFailure)?
            .next()
            .ok_or_else(|| Error::AddressResolutionMissing)?;

        let agent = pyro::Pyroscope::builder(
            "http://".to_owned() + address.to_string().as_str(),
            "ghost".to_owned(),
        ).backend(pprof_backend(PprofConfig::new().sample_rate(113))).build()?;

        Some(agent.start()?)
    } else {
        None
    };

    #[cfg(not(feature = "pyroscope"))]
    if cli.run.pyroscope_server.is_some() {
        return Err(Error::PyroscopeNotCompiledIn)
    }

    match &cli.subcommand {
        None => run_node_inner(cli, ghost_metrics::logger_hook()),
        Some(Subcommand::BuildSpec(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            Ok(runner.sync_run(|config| cmd.run(config.chain_spec, config.network))?)
        },
        Some(Subcommand::CheckBlock(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            set_ss58_version(chain_spec);

            runner.async_run(|mut config| {
                let (client, _, import_queue, task_manager) =
                    service::new_chain_ops(&mut config)?;
                Ok((cmd.run(client, import_queue).map_err(Error::SubstrateCli), task_manager))
            })
        },
        Some(Subcommand::ExportBlocks(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            set_ss58_version(chain_spec);

            Ok(runner.async_run(|mut config| {
                let ( client, _, _, task_manager ) = 
                    service::new_chain_ops(&mut config)?;
                Ok((cmd.run(client, config.database).map_err(Error::SubstrateCli), task_manager))
            })?)
        },
        Some(Subcommand::ExportState(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            set_ss58_version(chain_spec);

            Ok(runner.async_run(|mut config| {
                let ( client, _, _, task_manager ) = 
                    service::new_chain_ops(&mut config)?;
                Ok((cmd.run(client, config.chain_spec).map_err(Error::SubstrateCli), task_manager))
            })?)
        },
        Some(Subcommand::ImportBlocks(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            set_ss58_version(chain_spec);

            Ok(runner.async_run(|mut config| {
                let ( client, _, import_queue, task_manager ) = 
                    service::new_chain_ops(&mut config)?;
                Ok((cmd.run(client, import_queue).map_err(Error::SubstrateCli), task_manager))
            })?)
        },
        Some(Subcommand::PurgeChain(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            Ok(runner.sync_run(|config| cmd.run(config.database))?)
        },
        Some(Subcommand::Revert(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            set_ss58_version(chain_spec);

            Ok(runner.async_run(|mut config| {
                let (client, backend, _, task_manager) = 
                    service::new_chain_ops(&mut config)?;

                let aux_revert = Box::new(|client, backend, blocks| {
                    service::revert_backend(client, backend, blocks).map_err(|err| {
                        match err {
                            service::Error::Blockchain(err) => err.into(),
                            err => sc_cli::Error::Application(err.into()),
                        }
                    })
                });

                Ok((
                    cmd.run(client, backend, Some(aux_revert)).map_err(Error::SubstrateCli),
                    task_manager
                ))
            })?)
        },
        Some(Subcommand::Benchmark(cmd)) => {
            let runner = cli.create_runner(cmd).map_err(Error::SubstrateCli)?;
            let chain_spec = &runner.config().chain_spec;

            match cmd {
                #[cfg(not(feature = "runtime-benchmarks"))]
                BenchmarkCmd::Storage(_) => 
                    return Err(sc_cli::Error::Input(
                        "Compile with `--feature=runtime-benchmarks \
                        to enable storage benchmarks.".into()
                    ).into()),
                #[cfg(feature = "runtime-benchmarks")]
                BenchmarkCmd::Storage(cmd) => runner.sync_run(|mut config| {
                    let (client, backend, _, _,) = 
                        service::new_chain_ops(&mut config)?;
                    let db = backend.expose_db();
                    let storage = backend.expose_storage();

                    cmd.run(config, client.clone(), db, storage).map_err(Error::SubstrateCli)
                }),
                BenchmarkCmd::Block(cmd) => runner.sync_run(|mut config| {
                    let (client, _, _, _,) = 
                        service::new_chain_ops(&mut config)?;
                    cmd.run(client.clone()).map_err(Error::SubstrateCli)
                }),
                // These commands are very similar and can be handled in nearly the same way
                BenchmarkCmd::Extrinsic(_) | BenchmarkCmd::Overhead(_) =>
                    runner.sync_run(|mut config| {
                        let (client, _, _, _) = 
                            service::new_chain_ops(&mut config)?;

                        let inherent_data = benchmark_inherent_data()
                            .map_err(|e| format!("generating inherent data: {:?}", e))?;

                        let remark_builder = RemarkBuilder::new(
                            client.clone(), 
                            config.chain_spec.identify_chain(),
                        );

                        match cmd {
                            BenchmarkCmd::Extrinsic(cmd) => {
                                let tka_builder = TransferKeepAliveBuilder::new(
                                    client.clone(),
                                    Sr25519Keyring::Alice.to_account_id(),
                                    config.chain_spec.identify_chain(),
                                );

                                let ext_factory = ExtrinsicFactory(vec![
                                    Box::new(remark_builder),
                                    Box::new(tka_builder),
                                ]);

                                cmd.run(client.clone(), inherent_data, Vec::new(), &ext_factory)
                                    .map_err(Error::SubstrateCli)
                            },
                            BenchmarkCmd::Overhead(cmd) => cmd.run(
                                config,
                                client.clone(),
                                inherent_data,
                                Vec::new(),
                                &remark_builder,
                            ).map_err(Error::SubstrateCli),
                            _ => unreachable!("Ensured by the outside match; qed"),
                        }
                    }),
                BenchmarkCmd::Pallet(cmd) => {
                    set_ss58_version(chain_spec);

                    if cfg!(feature = "runtime-benchmarks") {
                        runner.sync_run(|config| {
                            cmd.run_with_spec::<sp_runtime::traits::HashingFor<service::Block>, ()>(
                                Some(config.chain_spec),
                            ).map_err(|e| Error::SubstrateCli(e))
                        })
                    } else {
                        Err(sc_cli::Error::Input(
                            "Benchmarking wasn't enabled when building the node. \
                            You can enable it with `--features=runtime-benchmarks`.".into()
                        ).into())
                    }
                },
                BenchmarkCmd::Machine(cmd) => runner.sync_run(|config| {
                    cmd.run(&config, ghost_machine_primitives::GHOST_NODE_REFERENCE_HARDWARE.clone())
                        .map_err(Error::SubstrateCli)
                }),
                // Note: this allows to implement additional new benchmark
                // commands.
                #[allow(unreachable_patterns)]
                _ => Err(Error::CommandNotImplemented),
            }
        },
        Some(Subcommand::Key(cmd)) => Ok(cmd.run(&cli)?),
        Some(Subcommand::ChainInfo(cmd)) => {
            let runner = cli.create_runner(cmd)?;
            Ok(runner.sync_run(|config| cmd.run::<service::Block>(&config))?)
        },
    }?;

    #[cfg(feature = "pyroscope")]
    if let Some(pyroscope_agent) = pyroscope_agent_maybe.take() {
        let agent = pyroscope_agent.stop()?;
        agent.shutdown();
    }
    Ok(())
}