1#![recursion_limit = "512"]2#![feature(try_blocks)]34pub(crate) mod cmds;56pub(crate) mod extra_args;78use std::{ffi::OsString, process::ExitCode};910use anyhow::{bail, Result};11use clap::{CommandFactory, Parser};12use cmds::{13 build_systems::{BuildSystems, Deploy},14 complete::Complete,15 info::Info,16 secrets::Secret,17 tf::Tf,18};19use fleet_base::{host::Config, opts::FleetOpts};20use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};2122#[cfg(feature = "indicatif")]23use human_repr::HumanCount;24#[cfg(feature = "indicatif")]25use indicatif::{ProgressState, ProgressStyle};26use tracing::{error, info, info_span, Instrument};27#[cfg(feature = "indicatif")]28use tracing_indicatif::IndicatifLayer;29use tracing_subscriber::{prelude::*, EnvFilter};3031#[derive(Parser)]32struct Prefetch {}33impl Prefetch {34 async fn run(&self, config: &Config) -> Result<()> {35 let mut prefetch_dir = config.directory.to_path_buf();36 prefetch_dir.push("prefetch");37 if !prefetch_dir.is_dir() {38 info!("nothing to prefetch: no prefetch directory");39 return Ok(());40 }41 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();42 for entry in std::fs::read_dir(&prefetch_dir)? {43 tasks.push(Box::pin(async {44 let entry = entry?;45 if !entry.metadata()?.is_file() {46 bail!("only files should exist in prefetch directory");47 }48 let span = info_span!(49 "prefetching",50 name = entry.file_name().to_string_lossy().as_ref()51 );52 let mut path = OsString::new();53 path.push("file://");54 path.push(entry.path());5556 let mut status = config.local_host().cmd("nix").await?;57 status.args(&config.nix_args);58 status.arg("store").arg("prefetch-file").arg(path);59 status.run_nix_string().instrument(span).await?;60 Ok(())61 }));62 }63 tasks.try_collect::<Vec<()>>().await?;64 Ok(())65 }66}6768#[derive(Parser)]69enum Opts {70 71 BuildSystems(BuildSystems),7273 Deploy(Deploy),74 75 #[clap(subcommand)]76 Secret(Secret),77 78 Prefetch(Prefetch),79 80 Info(Info),81 82 #[clap(hide(true))]83 Complete(Complete),84 85 Tf(Tf),86}8788#[derive(Parser)]89#[clap(version, author)]90struct RootOpts {91 #[clap(flatten)]92 fleet_opts: FleetOpts,93 #[clap(subcommand)]94 command: Opts,95}9697async fn run_command(config: &Config, opts: FleetOpts, command: Opts) -> Result<()> {98 match command {99 Opts::BuildSystems(c) => c.run(config, &opts).await?,100 Opts::Deploy(d) => d.run(config, &opts).await?,101 Opts::Secret(s) => s.run(config, &opts).await?,102 Opts::Info(i) => i.run(config).await?,103 Opts::Prefetch(p) => p.run(config).await?,104 Opts::Tf(t) => t.run(config).await?,105 106 Opts::Complete(c) => {107 tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?108 }109 };110 Ok(())111}112113fn setup_logging() {114 #[cfg(feature = "indicatif")]115 let indicatif_layer = {116 use std::time::Duration;117118 IndicatifLayer::new().with_progress_style(119 ProgressStyle::with_template(120 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",121 )122 .unwrap()123 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {124 let Some(len) = state.len() else {125 return;126 };127 let pos = state.pos();128 if pos > len {129 let _ = write!(writer, "{}", pos.human_count_bare());130 } else {131 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());132 }133 })134 .with_key(135 "color_start",136 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {137 let elapsed = state.elapsed();138139 if elapsed > Duration::from_secs(60) {140 141 let _ = write!(writer, "\x1b[{}m", 1 + 30);142 } else if elapsed > Duration::from_secs(30) {143 144 let _ = write!(writer, "\x1b[{}m", 3 + 30);145 }146 },147 )148 .with_key(149 "color_end",150 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {151 if state.elapsed() > Duration::from_secs(30) {152 let _ = write!(writer, "\x1b[0m");153 }154 },155 ),156 )157 };158159 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));160161 let reg = tracing_subscriber::registry().with({162 let sub = tracing_subscriber::fmt::layer()163 .without_time()164 .with_target(false);165 #[cfg(feature = "indicatif")]166 let sub = sub.with_writer(indicatif_layer.get_stdout_writer());167 sub.with_filter(filter) 168 });169 170 #[cfg(feature = "indicatif")]171 let reg = reg.with(indicatif_layer);172 reg.init();173}174175fn main() -> ExitCode {176 let opts = RootOpts::parse();177 if let Opts::Complete(c) = &opts.command {178 c.run(RootOpts::command());179 return ExitCode::SUCCESS;180 }181182 setup_logging();183 async_main(opts)184}185186#[tokio::main]187async fn async_main(opts: RootOpts) -> ExitCode {188 if let Err(e) = main_real(opts).await {189 190 191 #[cfg(feature = "indicatif")]192 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");193 error!("{e:#}");194 return ExitCode::FAILURE;195 }196 ExitCode::SUCCESS197}198199async fn main_real(opts: RootOpts) -> Result<()> {200 nix_eval::init_tokio();201202 let nix_args = std::env::var_os("NIX_ARGS")203 .map(|a| extra_args::parse_os(&a))204 .transpose()?205 .unwrap_or_default();206 let config = opts207 .fleet_opts208 .build(209 nix_args,210 matches!(opts.command, Opts::Deploy(_) | Opts::BuildSystems(_)),211 )212 .await?;213214 match run_command(&config, opts.fleet_opts, 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}