1#![recursion_limit = "512"]2#![feature(try_blocks)]34pub(crate) mod cmds;5pub(crate) mod command;6pub(crate) mod host;7pub(crate) mod keys;89pub(crate) mod extra_args;1011mod fleetdata;1213use std::{ffi::OsString, process::ExitCode};1415use anyhow::{bail, Result};16use clap::{CommandFactory, Parser};17use cmds::{18 build_systems::{BuildSystems, Deploy},19 complete::Complete,20 info::Info,21 secrets::Secret,22 tf::Tf,23};24use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};25use host::{Config, FleetOpts};26#[cfg(feature = "indicatif")]27use human_repr::HumanCount;28#[cfg(feature = "indicatif")]29use indicatif::{ProgressState, ProgressStyle};30use tracing::{error, info, info_span, Instrument};31#[cfg(feature = "indicatif")]32use tracing_indicatif::IndicatifLayer;33use tracing_subscriber::{prelude::*, EnvFilter};3435use crate::command::MyCommand;3637#[derive(Parser)]38struct Prefetch {}39impl Prefetch {40 async fn run(&self, config: &Config) -> Result<()> {41 let mut prefetch_dir = config.directory.to_path_buf();42 prefetch_dir.push("prefetch");43 if !prefetch_dir.is_dir() {44 info!("nothing to prefetch: no prefetch directory");45 return Ok(());46 }47 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();48 for entry in std::fs::read_dir(&prefetch_dir)? {49 tasks.push(Box::pin(async {50 let entry = entry?;51 if !entry.metadata()?.is_file() {52 bail!("only files should exist in prefetch directory");53 }54 let span = info_span!(55 "prefetching",56 name = entry.file_name().to_string_lossy().as_ref()57 );58 let mut path = OsString::new();59 path.push("file://");60 path.push(entry.path());6162 let mut status = config.local_host().cmd("nix").await?;63 status.args(&config.nix_args);64 status.arg("store").arg("prefetch-file").arg(path);65 status.run_nix_string().instrument(span).await?;66 Ok(())67 }));68 }69 tasks.try_collect::<Vec<()>>().await?;70 Ok(())71 }72}7374#[derive(Parser)]75enum Opts {76 77 BuildSystems(BuildSystems),7879 Deploy(Deploy),80 81 #[clap(subcommand)]82 Secret(Secret),83 84 Prefetch(Prefetch),85 86 Info(Info),87 88 #[clap(hide(true))]89 Complete(Complete),90 91 Tf(Tf),92}9394#[derive(Parser)]95#[clap(version, author)]96struct RootOpts {97 #[clap(flatten)]98 fleet_opts: FleetOpts,99 #[clap(subcommand)]100 command: Opts,101}102103async fn run_command(config: &Config, command: Opts) -> Result<()> {104 match command {105 Opts::BuildSystems(c) => c.run(config).await?,106 Opts::Deploy(d) => d.run(config).await?,107 Opts::Secret(s) => s.run(config).await?,108 Opts::Info(i) => i.run(config).await?,109 Opts::Prefetch(p) => p.run(config).await?,110 Opts::Tf(t) => t.run(config).await?,111 112 Opts::Complete(c) => {113 tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?114 }115 };116 Ok(())117}118119fn setup_logging() {120 #[cfg(feature = "indicatif")]121 let indicatif_layer = {122 use std::time::Duration;123124 IndicatifLayer::new().with_progress_style(125 ProgressStyle::with_template(126 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",127 )128 .unwrap()129 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {130 let Some(len) = state.len() else {131 return;132 };133 let pos = state.pos();134 if pos > len {135 let _ = write!(writer, "{}", pos.human_count_bare());136 } else {137 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());138 }139 })140 .with_key(141 "color_start",142 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {143 let elapsed = state.elapsed();144145 if elapsed > Duration::from_secs(60) {146 147 let _ = write!(writer, "\x1b[{}m", 1 + 30);148 } else if elapsed > Duration::from_secs(30) {149 150 let _ = write!(writer, "\x1b[{}m", 3 + 30);151 }152 },153 )154 .with_key(155 "color_end",156 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {157 if state.elapsed() > Duration::from_secs(30) {158 let _ = write!(writer, "\x1b[0m");159 }160 },161 ),162 )163 };164165 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));166167 let reg = tracing_subscriber::registry().with({168 let sub = tracing_subscriber::fmt::layer()169 .without_time()170 .with_target(false);171 #[cfg(feature = "indicatif")]172 let sub = sub.with_writer(indicatif_layer.get_stdout_writer());173 sub.with_filter(filter) 174 });175 176 #[cfg(feature = "indicatif")]177 let reg = reg.with(indicatif_layer);178 reg.init();179}180181fn main() -> ExitCode {182 let opts = RootOpts::parse();183 if let Opts::Complete(c) = &opts.command {184 c.run(RootOpts::command());185 return ExitCode::SUCCESS;186 }187188 setup_logging();189 async_main(opts)190}191192#[tokio::main]193async fn async_main(opts: RootOpts) -> ExitCode {194 if let Err(e) = main_real(opts).await {195 196 197 #[cfg(feature = "indicatif")]198 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");199 error!("{e:#}");200 return ExitCode::FAILURE;201 }202 ExitCode::SUCCESS203}204205async fn main_real(opts: RootOpts) -> Result<()> {206 nix_eval::init_tokio();207208 let nix_args = std::env::var_os("NIX_ARGS")209 .map(|a| extra_args::parse_os(&a))210 .transpose()?211 .unwrap_or_default();212 let config = opts.fleet_opts.build(nix_args).await?;213214 match run_command(&config, opts.command).await {215 Ok(()) => {216 config.save()?;217 Ok(())218 }219 Err(e) => {220 let _ = config.save();221 Err(e)222 }223 }224}225226#[cfg(test)]227mod tests {228 use super::*;229230 #[test]231 fn verify_command() {232 use clap::CommandFactory;233 RootOpts::command().debug_assert();234 }235}