Line data    Source code

       1              : use std::fs::File;
       2              : 
       3              : use clap::Parser;
       4              : use tracing::info;
       5              : use utils::logging;
       6              : 
       7              : /// Re-usable pieces of code that aren't CLI-specific.
       8              : mod util {
       9              :     pub(crate) mod request_stats;
      10              :     #[macro_use]
      11              :     pub(crate) mod tokio_thread_local_stats;
      12              :     /// Re-usable pieces of CLI-specific code.
      13              :     pub(crate) mod cli {
      14              :         pub(crate) mod targets;
      15              :     }
      16              : }
      17              : 
      18              : /// The pagebench CLI sub-commands, dispatched in [`main`] below.
      19              : mod cmd {
      20              :     pub(super) mod aux_files;
      21              :     pub(super) mod basebackup;
      22              :     pub(super) mod getpage_latest_lsn;
      23              :     pub(super) mod idle_streams;
      24              :     pub(super) mod ondemand_download_churn;
      25              :     pub(super) mod trigger_initial_size_calculation;
      26              : }
      27              : 
      28              : /// Component-level performance test for pageserver.
      29              : #[derive(clap::Parser)]
      30              : struct Args {
      31              :     /// Takes a client CPU profile into profile.svg. The benchmark must exit cleanly before it's
      32              :     /// written, e.g. via --runtime.
      33              :     #[arg(long)]
      34              :     profile: bool,
      35              : 
      36              :     #[command(subcommand)]
      37              :     subcommand: Subcommand,
      38              : }
      39              : 
      40              : #[derive(clap::Subcommand)]
      41              : enum Subcommand {
      42              :     Basebackup(cmd::basebackup::Args),
      43              :     GetPageLatestLsn(cmd::getpage_latest_lsn::Args),
      44              :     TriggerInitialSizeCalculation(cmd::trigger_initial_size_calculation::Args),
      45              :     OndemandDownloadChurn(cmd::ondemand_download_churn::Args),
      46              :     AuxFiles(cmd::aux_files::Args),
      47              :     IdleStreams(cmd::idle_streams::Args),
      48              : }
      49              : 
      50            0 : fn main() -> anyhow::Result<()> {
      51            0 :     logging::init(
      52            0 :         logging::LogFormat::Plain,
      53            0 :         logging::TracingErrorLayerEnablement::Disabled,
      54            0 :         logging::Output::Stderr,
      55            0 :     )?;
      56            0 :     logging::replace_panic_hook_with_tracing_panic_hook().forget();
      57              : 
      58            0 :     let args = Args::parse();
      59              : 
      60              :     // Start a CPU profile if requested.
      61            0 :     let mut profiler = None;
      62            0 :     if args.profile {
      63            0 :         profiler = Some(
      64            0 :             pprof::ProfilerGuardBuilder::default()
      65            0 :                 .frequency(1000)
      66            0 :                 .blocklist(&["libc", "libgcc", "pthread", "vdso"])
      67            0 :                 .build()?,
      68              :         );
      69            0 :     }
      70              : 
      71            0 :     match args.subcommand {
      72            0 :         Subcommand::Basebackup(args) => cmd::basebackup::main(args),
      73            0 :         Subcommand::GetPageLatestLsn(args) => cmd::getpage_latest_lsn::main(args),
      74            0 :         Subcommand::TriggerInitialSizeCalculation(args) => {
      75            0 :             cmd::trigger_initial_size_calculation::main(args)
      76              :         }
      77            0 :         Subcommand::OndemandDownloadChurn(args) => cmd::ondemand_download_churn::main(args),
      78            0 :         Subcommand::AuxFiles(args) => cmd::aux_files::main(args),
      79            0 :         Subcommand::IdleStreams(args) => cmd::idle_streams::main(args),
      80            0 :     }?;
      81              : 
      82              :     // Generate a CPU flamegraph if requested.
      83            0 :     if let Some(profiler) = profiler {
      84            0 :         let report = profiler.report().build()?;
      85            0 :         drop(profiler); // stop profiling
      86            0 :         let file = File::create("profile.svg")?;
      87            0 :         report.flamegraph(file)?;
      88            0 :         info!("wrote CPU profile flamegraph to profile.svg")
      89            0 :     }
      90              : 
      91            0 :     Ok(())
      92            0 : }