Repository: bottlerocket-os/bottlerocket
File: sources/api/shibaken/src/main.rs (Rust)

Code before:
mod partition;

/// Returns information gathered from the AWS instance metadata service (IMDS).
#[derive(FromArgs, Debug)]
struct Args {
    #[argh(option, default = "LevelFilter::Info")]
    /// filter level for log messages
    log_level: LevelFilter,

    #[argh(subcommand)]
    command: Commands,
}

#[derive(FromArgs, Debug)]
#[argh(subcommand)]
enum Commands {
    /// Fetch and populate the admin container's user-data with authorized ssh keys.
    GenerateAdminUserdata(admin_userdata::GenerateAdminUserdata),

    /// Fetch and return whether or not this host is in the given partition.
    /// Accepts multiple partitions, returning `true` if the host is in any of the given partitions.
    IsPartition(partition::IsPartition),
}

async fn run() -> Result<()> {
    let args: Args = argh::from_env();

    // TerminalMode::Stderr will send all logs to stderr, as sundog only expects the json output of
    // the setting on stdout.
    TermLogger::init(
        args.log_level,
        LogConfig::default(),
        TerminalMode::Stderr,
        ColorChoice::Auto,
    )
    .context(error::LoggerSnafu)?;

    log::info!("shibaken started");

    match args.command {
        Commands::GenerateAdminUserdata(generate_admin_userdata) => {
            generate_admin_userdata.run().await
        }
        Commands::IsPartition(is_partition) => is_partition.run().await,
    }
}

// Returning a Result from main makes it print a Debug representation of the error, but with Snafu
// we have nice Display representations of the error, so we wrap "main" (run) and print any error.
// https://github.com/shepmaster/snafu/issues/110
#[tokio::main]

Review comment (suggestion):
It might be better to say something like "Poll lifecycle state and wait until instance to be marked as `InService`".

Code after:
pub(crate) mod error;
mod partition;
mod warmpool;

/// Returns information gathered from the AWS instance metadata service (IMDS).
#[derive(FromArgs, Debug)]
struct Args {
    #[argh(option, default = "LevelFilter::Info")]
    /// filter level for log messages
    log_level: LevelFilter,

    #[argh(subcommand)]
    command: Commands,
}

#[derive(FromArgs, Debug)]
#[argh(subcommand)]
enum Commands {
    /// Fetch and populate the admin container's user-data with authorized ssh keys.
    GenerateAdminUserdata(admin_userdata::GenerateAdminUserdata),

    /// Fetch and return whether or not this host is in the given partition.
    /// Accepts multiple partitions, returning `true` if the host is in any of the given partitions.
    IsPartition(partition::IsPartition),

    /// Poll lifecycle state and wait until instance to be marked as InService
    WarmPoolWait(warmpool::autoscaling_warm_pool::WarmPoolWait),
}

async fn run() -> Result<()> {
    let args: Args = argh::from_env();

    // TerminalMode::Stderr will send all logs to stderr, as sundog only expects the json output of
    // the setting on stdout.
    TermLogger::init(
        args.log_level,
        LogConfig::default(),
        TerminalMode::Stderr,
        ColorChoice::Auto,
    )
    .context(error::LoggerSnafu)?;

    log::info!("shibaken started");

    match args.command {
        Commands::GenerateAdminUserdata(generate_admin_userdata) => {
            generate_admin_userdata.run().await
        }
        Commands::IsPartition(is_partition) => is_partition.run().await,
        Commands::WarmPoolWait(warm_pool_wait) => warm_pool_wait
            .run()