1#![recursion_limit = "512"]23pub(crate) mod cmds;45pub(crate) mod extra_args;67use std::{ffi::OsString, process::ExitCode};89use anyhow::{bail, Result};10use clap::{CommandFactory, Parser};11use cmds::{12 build_systems::{BuildSystems, Deploy},13 complete::Complete,14 info::Info,15 secrets::Secret,16 tf::Tf,17};18use fleet_base::{host::Config, opts::FleetOpts};19use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};2021#[cfg(feature = "indicatif")]22use human_repr::HumanCount;23#[cfg(feature = "indicatif")]24use indicatif::{ProgressState, ProgressStyle};25use tracing::{error, info, info_span, Instrument};26#[cfg(feature = "indicatif")]27use tracing_indicatif::IndicatifLayer;28use tracing_subscriber::{prelude::*, EnvFilter};2930#[derive(Parser)]31struct Prefetch {}32impl Prefetch {33 async fn run(&self, config: &Config) -> Result<()> {34 let mut prefetch_dir = config.directory.to_path_buf();35 prefetch_dir.push("prefetch");36 if !prefetch_dir.is_dir() {37 info!("nothing to prefetch: no prefetch directory");38 return Ok(());39 }40 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();41 for entry in std::fs::read_dir(&prefetch_dir)? {42 tasks.push(Box::pin(async {43 let entry = entry?;44 if !entry.metadata()?.is_file() {45 bail!("only files should exist in prefetch directory");46 }47 let span = info_span!(48 "prefetching",49 name = entry.file_name().to_string_lossy().as_ref()50 );51 let mut path = OsString::new();52 path.push("file://");53 path.push(entry.path());5455 let mut status = config.local_host().cmd("nix").await?;56 status.args(&config.nix_args);57 status.arg("store").arg("prefetch-file").arg(path);58 status.run_nix_string().instrument(span).await?;59 Ok(())60 }));61 }62 tasks.try_collect::<Vec<()>>().await?;63 Ok(())64 }65}6667#[derive(Parser)]68enum Opts {69 70 BuildSystems(BuildSystems),7172 Deploy(Deploy),73 74 #[clap(subcommand)]75 Secret(Secret),76 77 Prefetch(Prefetch),78 79 Info(Info),80 81 #[clap(hide(true))]82 Complete(Complete),83 84 Tf(Tf),85}8687#[derive(Parser)]88#[clap(version, author)]89struct RootOpts {90 #[clap(flatten)]91 fleet_opts: FleetOpts,92 #[clap(subcommand)]93 command: Opts,94}9596async fn run_command(config: &Config, opts: FleetOpts, command: Opts) -> Result<()> {97 match command {98 Opts::BuildSystems(c) => c.run(config, &opts).await?,99 Opts::Deploy(d) => d.run(config, &opts).await?,100 Opts::Secret(s) => s.run(config, &opts).await?,101 Opts::Info(i) => i.run(config).await?,102 Opts::Prefetch(p) => p.run(config).await?,103 Opts::Tf(t) => t.run(config).await?,104 105 Opts::Complete(c) => {106 tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?107 }108 };109 Ok(())110}111112fn setup_logging() {113 #[cfg(feature = "indicatif")]114 let indicatif_layer = {115 use std::time::Duration;116117 IndicatifLayer::new().with_progress_style(118 ProgressStyle::with_template(119 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",120 )121 .unwrap()122 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {123 let Some(len) = state.len() else {124 return;125 };126 let pos = state.pos();127 if pos > len {128 let _ = write!(writer, "{}", pos.human_count_bare());129 } else {130 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());131 }132 })133 .with_key(134 "color_start",135 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {136 let elapsed = state.elapsed();137138 if elapsed > Duration::from_secs(60) {139 140 let _ = write!(writer, "\x1b[{}m", 1 + 30);141 } else if elapsed > Duration::from_secs(30) {142 143 let _ = write!(writer, "\x1b[{}m", 3 + 30);144 }145 },146 )147 .with_key(148 "color_end",149 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {150 if state.elapsed() > Duration::from_secs(30) {151 let _ = write!(writer, "\x1b[0m");152 }153 },154 ),155 )156 };157158 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));159160 let reg = tracing_subscriber::registry().with({161 let sub = tracing_subscriber::fmt::layer()162 .without_time()163 .with_target(false);164 #[cfg(feature = "indicatif")]165 let sub = sub.with_writer(indicatif_layer.get_stdout_writer());166 sub.with_filter(filter) 167 });168 169 #[cfg(feature = "indicatif")]170 let reg = reg.with(indicatif_layer);171 reg.init();172}173174fn main() -> ExitCode {175 let opts = RootOpts::parse();176 if let Opts::Complete(c) = &opts.command {177 c.run(RootOpts::command());178 return ExitCode::SUCCESS;179 }180181 setup_logging();182 async_main(opts)183}184185#[tokio::main]186async fn async_main(opts: RootOpts) -> ExitCode {187 if let Err(e) = main_real(opts).await {188 189 190 #[cfg(feature = "indicatif")]191 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");192 error!("{e:#}");193 return ExitCode::FAILURE;194 }195 ExitCode::SUCCESS196}197198async fn main_real(opts: RootOpts) -> Result<()> {199 nix_eval::init_tokio();200201 let nix_args = std::env::var_os("NIX_ARGS")202 .map(|a| extra_args::parse_os(&a))203 .transpose()?204 .unwrap_or_default();205 let config = opts206 .fleet_opts207 .build(208 nix_args,209 matches!(opts.command, Opts::Deploy(_) | Opts::BuildSystems(_)),210 )211 .await?;212213 match run_command(&config, opts.fleet_opts, opts.command).await {214 Ok(()) => {215 config.save()?;216 Ok(())217 }218 Err(e) => {219 let _ = config.save();220 Err(e)221 }222 }223}224225#[cfg(test)]226mod tests {227 use super::*;228229 #[test]230 fn verify_command() {231 use clap::CommandFactory;232 RootOpts::command().debug_assert();233 }234}