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};23use futures::{future::LocalBoxFuture, stream::FuturesUnordered, TryStreamExt};24use host::{Config, FleetOpts};25#[cfg(feature = "indicatif")]26use human_repr::HumanCount;27#[cfg(feature = "indicatif")]28use indicatif::{ProgressState, ProgressStyle};29use tracing::{error, info, info_span, Instrument};30#[cfg(feature = "indicatif")]31use tracing_indicatif::IndicatifLayer;32use tracing_subscriber::{prelude::*, EnvFilter};3334use crate::command::MyCommand;3536#[derive(Parser)]37struct Prefetch {}38impl Prefetch {39 async fn run(&self, config: &Config) -> Result<()> {40 let mut prefetch_dir = config.directory.to_path_buf();41 prefetch_dir.push("prefetch");42 if !prefetch_dir.is_dir() {43 info!("nothing to prefetch: no prefetch directory");44 return Ok(());45 }46 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();47 for entry in std::fs::read_dir(&prefetch_dir)? {48 tasks.push(Box::pin(async {49 let entry = entry?;50 if !entry.metadata()?.is_file() {51 bail!("only files should exist in prefetch directory");52 }53 let span = info_span!(54 "prefetching",55 name = entry.file_name().to_string_lossy().as_ref()56 );57 let mut path = OsString::new();58 path.push("file://");59 path.push(entry.path());6061 let mut status = MyCommand::new("nix");62 status.args(&config.nix_args);63 status.arg("store").arg("prefetch-file").arg(path);64 status.run_nix_string().instrument(span).await?;65 Ok(())66 }));67 }68 tasks.try_collect::<Vec<()>>().await?;69 Ok(())70 }71}7273#[derive(Parser)]74enum Opts {75 76 BuildSystems(BuildSystems),7778 Deploy(Deploy),79 80 #[clap(subcommand)]81 Secret(Secret),82 83 Prefetch(Prefetch),84 85 Info(Info),86 87 #[clap(hide(true))]88 Complete(Complete),89}9091#[derive(Parser)]92#[clap(version, author)]93struct RootOpts {94 #[clap(flatten)]95 fleet_opts: FleetOpts,96 #[clap(subcommand)]97 command: Opts,98}99100async fn run_command(config: &Config, command: Opts) -> Result<()> {101 match command {102 Opts::BuildSystems(c) => c.run(config).await?,103 Opts::Deploy(d) => d.run(config).await?,104 Opts::Secret(s) => s.run(config).await?,105 Opts::Info(i) => i.run(config).await?,106 Opts::Prefetch(p) => p.run(config).await?,107 108 Opts::Complete(c) => {109 tokio::task::spawn_blocking(move || c.run(RootOpts::command())).await?110 }111 };112 Ok(())113}114115fn setup_logging() {116 #[cfg(feature = "indicatif")]117 let indicatif_layer = {118 use std::time::Duration;119120 IndicatifLayer::new().with_progress_style(121 ProgressStyle::with_template(122 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",123 )124 .unwrap()125 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {126 let Some(len) = state.len() else {127 return;128 };129 let pos = state.pos();130 if pos > len {131 let _ = write!(writer, "{}", pos.human_count_bare());132 } else {133 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());134 }135 })136 .with_key(137 "color_start",138 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {139 let elapsed = state.elapsed();140141 if elapsed > Duration::from_secs(60) {142 143 let _ = write!(writer, "\x1b[{}m", 1 + 30);144 } else if elapsed > Duration::from_secs(30) {145 146 let _ = write!(writer, "\x1b[{}m", 3 + 30);147 }148 },149 )150 .with_key(151 "color_end",152 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {153 if state.elapsed() > Duration::from_secs(30) {154 let _ = write!(writer, "\x1b[0m");155 }156 },157 ),158 )159 };160161 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));162163 let reg = tracing_subscriber::registry().with({164 let sub = tracing_subscriber::fmt::layer()165 .without_time()166 .with_target(false);167 #[cfg(feature = "indicatif")]168 let sub = sub.with_writer(indicatif_layer.get_stdout_writer());169 sub.with_filter(filter) 170 });171 172 #[cfg(feature = "indicatif")]173 let reg = reg.with(indicatif_layer);174 reg.init();175}176177fn main() -> ExitCode {178 let opts = RootOpts::parse();179 if let Opts::Complete(c) = &opts.command {180 c.run(RootOpts::command());181 return ExitCode::SUCCESS;182 }183184 setup_logging();185 async_main(opts)186}187188#[tokio::main]189async fn async_main(opts: RootOpts) -> ExitCode {190 if let Err(e) = main_real(opts).await {191 192 193 #[cfg(feature = "indicatif")]194 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");195 error!("{e:#}");196 return ExitCode::FAILURE;197 }198 ExitCode::SUCCESS199}200201async fn main_real(opts: RootOpts) -> Result<()> {202 nix_eval::init_tokio();203204 let nix_args = std::env::var_os("NIX_ARGS")205 .map(|a| extra_args::parse_os(&a))206 .transpose()?207 .unwrap_or_default();208 let config = opts.fleet_opts.build(nix_args).await?;209210 match run_command(&config, opts.command).await {211 Ok(()) => {212 config.save()?;213 Ok(())214 }215 Err(e) => {216 let _ = config.save();217 Err(e)218 }219 }220}221222#[cfg(test)]223mod tests {224 use super::*;225226 #[test]227 fn verify_command() {228 use clap::CommandFactory;229 RootOpts::command().debug_assert();230 }231}