1#![recursion_limit = "512"]2#![feature(try_blocks, lint_reasons)]34pub(crate) mod cmds;5pub(crate) mod command;6pub(crate) mod host;7pub(crate) mod keys;89pub(crate) mod better_nix_eval;10pub(crate) mod extra_args;1112mod fleetdata;1314use std::{ffi::OsString, process::ExitCode};1516use anyhow::{bail, Result};17use clap::Parser;1819use cmds::{20 build_systems::{BuildSystems, Deploy},21 info::Info,22 secrets::Secret,23};24use futures::future::LocalBoxFuture;25use futures::stream::FuturesUnordered;26use futures::TryStreamExt;27use host::{Config, FleetOpts};28#[cfg(feature = "indicatif")]29use human_repr::HumanCount;30#[cfg(feature = "indicatif")]31use indicatif::{ProgressState, ProgressStyle};32use tracing::{error, info};33use tracing::{info_span, Instrument};34#[cfg(feature = "indicatif")]35use tracing_indicatif::IndicatifLayer;36use tracing_subscriber::{prelude::*, EnvFilter};3738use crate::command::MyCommand;3940#[derive(Parser)]41struct Prefetch {}42impl Prefetch {43 async fn run(&self, config: &Config) -> Result<()> {44 let mut prefetch_dir = config.directory.to_path_buf();45 prefetch_dir.push("prefetch");46 if !prefetch_dir.is_dir() {47 info!("nothing to prefetch: no prefetch directory");48 return Ok(());49 }50 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();51 for entry in std::fs::read_dir(&prefetch_dir)? {52 tasks.push(Box::pin(async {53 let entry = entry?;54 if !entry.metadata()?.is_file() {55 bail!("only files should exist in prefetch directory");56 }57 let span = info_span!(58 "prefetching",59 name = entry.file_name().to_string_lossy().as_ref()60 );61 let mut path = OsString::new();62 path.push("file://");63 path.push(entry.path());6465 let mut status = MyCommand::new("nix");66 status.args(&config.nix_args);67 status.arg("store").arg("prefetch-file").arg(path);68 status.run_nix_string().instrument(span).await?;69 Ok(())70 }));71 }72 tasks.try_collect::<Vec<()>>().await?;73 Ok(())74 }75}7677#[derive(Parser)]78enum Opts {79 80 BuildSystems(BuildSystems),8182 Deploy(Deploy),83 84 #[clap(subcommand)]85 Secret(Secret),86 87 Prefetch(Prefetch),88 89 Info(Info),90}9192#[derive(Parser)]93#[clap(version, author)]94struct RootOpts {95 #[clap(flatten)]96 fleet_opts: FleetOpts,97 #[clap(subcommand)]98 command: Opts,99}100101async fn run_command(config: &Config, command: Opts) -> Result<()> {102 match command {103 Opts::BuildSystems(c) => c.run(config).await?,104 Opts::Deploy(d) => d.run(config).await?,105 Opts::Secret(s) => s.run(config).await?,106 Opts::Info(i) => i.run(config).await?,107 Opts::Prefetch(p) => p.run(config).await?,108 };109 Ok(())110}111112fn setup_logging() {113 #[cfg(feature = "indicatif")]114 let indicatif_layer = IndicatifLayer::new().with_progress_style(115 ProgressStyle::with_template(116 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",117 )118 .unwrap()119 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {120 let Some(len) = state.len() else {121 return;122 };123 let pos = state.pos();124 if pos > len {125 let _ = write!(writer, "{}", pos.human_count_bare());126 } else {127 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());128 }129 })130 .with_key(131 "color_start",132 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {133 use std::time::Duration;134 let elapsed = state.elapsed();135136 if elapsed > Duration::from_secs(60) {137 138 let _ = write!(writer, "\x1b[{}m", 1 + 30);139 } else if elapsed > Duration::from_secs(30) {140 141 let _ = write!(writer, "\x1b[{}m", 3 + 30);142 }143 },144 )145 .with_key(146 "color_end",147 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {148 if state.elapsed() > Duration::from_secs(30) {149 let _ = write!(writer, "\x1b[0m");150 }151 },152 ),153 );154155 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));156157 let reg = tracing_subscriber::registry().with({158 let sub = tracing_subscriber::fmt::layer()159 .without_time()160 .with_target(false);161 #[cfg(feature = "indicatif")]162 let sub = sub.with_writer(indicatif_layer.get_stdout_writer());163 sub.with_filter(filter) 164 });165 166 #[cfg(feature = "indicatif")]167 let reg = reg.with(indicatif_layer);168 reg.init();169}170171#[tokio::main]172async fn main() -> ExitCode {173 setup_logging();174 if let Err(e) = main_real().await {175 176 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");177 error!("{e:#}");178 return ExitCode::FAILURE;179 }180 ExitCode::SUCCESS181}182183async fn main_real() -> Result<()> {184 let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());185186 let nix_args = std::env::var_os("NIX_ARGS")187 .map(|a| extra_args::parse_os(&a))188 .transpose()?189 .unwrap_or_default();190 let opts = RootOpts::parse();191 let config = opts.fleet_opts.build(nix_args).await?;192193 match run_command(&config, opts.command).await {194 Ok(()) => {195 config.save()?;196 Ok(())197 }198 Err(e) => {199 let _ = config.save();200 Err(e)201 }202 }203}