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;15use std::process::exit;16use std::time::Duration;1718use anyhow::{bail, Result};19use clap::Parser;2021use cmds::{22 build_systems::{BuildSystems, Deploy},23 info::Info,24 secrets::Secret,25};26use futures::future::LocalBoxFuture;27use futures::stream::FuturesUnordered;28use futures::TryStreamExt;29use host::{Config, FleetOpts};30use human_repr::HumanCount;31use indicatif::{ProgressState, ProgressStyle};32use tracing::{error, info};33use tracing::{info_span, Instrument};34use tracing_indicatif::IndicatifLayer;35use tracing_subscriber::{prelude::*, EnvFilter};3637use crate::command::MyCommand;3839#[derive(Parser)]40struct Prefetch {}41impl Prefetch {42 async fn run(&self, config: &Config) -> Result<()> {43 let mut prefetch_dir = config.directory.to_path_buf();44 prefetch_dir.push("prefetch");45 if !prefetch_dir.is_dir() {46 info!("nothing to prefetch: no prefetch directory");47 return Ok(());48 }49 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();50 for entry in std::fs::read_dir(&prefetch_dir)? {51 tasks.push(Box::pin(async {52 let entry = entry?;53 if !entry.metadata()?.is_file() {54 bail!("only files should exist in prefetch directory");55 }56 let span = info_span!(57 "prefetching",58 name = entry.file_name().to_string_lossy().as_ref()59 );60 let mut path = OsString::new();61 path.push("file://");62 path.push(entry.path());6364 let mut status = MyCommand::new("nix");65 status.arg("store").arg("prefetch-file").arg(path);66 status.run_nix_string().instrument(span).await?;67 Ok(())68 }));69 }70 tasks.try_collect::<Vec<()>>().await?;71 Ok(())72 }73}7475#[derive(Parser)]76enum Opts {77 78 BuildSystems(BuildSystems),7980 Deploy(Deploy),81 82 #[clap(subcommand)]83 Secret(Secret),84 85 Prefetch(Prefetch),86 87 Info(Info),88}8990#[derive(Parser)]91#[clap(version, author)]92struct RootOpts {93 #[clap(flatten)]94 fleet_opts: FleetOpts,95 #[clap(subcommand)]96 command: Opts,97}9899async fn run_command(config: &Config, command: Opts) -> Result<()> {100 match command {101 Opts::BuildSystems(c) => c.run(config).await?,102 Opts::Deploy(d) => d.run(config).await?,103 Opts::Secret(s) => s.run(config).await?,104 Opts::Info(i) => i.run(config).await?,105 Opts::Prefetch(p) => p.run(config).await?,106 };107 Ok(())108}109110fn setup_logging() {111 let indicatif_layer = IndicatifLayer::new().with_progress_style(112 ProgressStyle::with_template(113 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",114 )115 .unwrap()116 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {117 let Some(len) = state.len() else {118 return;119 };120 let pos = state.pos();121 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());122 })123 .with_key(124 "color_start",125 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {126 let elapsed = state.elapsed();127128 if elapsed > Duration::from_secs(60) {129 130 let _ = write!(writer, "\x1b[{}m", 1 + 30);131 } else if elapsed > Duration::from_secs(30) {132 133 let _ = write!(writer, "\x1b[{}m", 3 + 30);134 }135 },136 )137 .with_key(138 "color_end",139 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {140 if state.elapsed() > Duration::from_secs(30) {141 let _ = write!(writer, "\x1b[0m");142 }143 },144 ),145 );146147 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));148149 tracing_subscriber::registry()150 .with(151 tracing_subscriber::fmt::layer()152 .without_time()153 .with_target(true)154 .with_writer(indicatif_layer.get_stderr_writer())155 .with_filter(filter), 156 )157 .with(indicatif_layer)158 .init();159}160161#[tokio::main]162async fn main() {163 setup_logging();164 if let Err(e) = main_real().await {165 error!("{e:#}");166 exit(1);167 }168}169170async fn main_real() -> Result<()> {171 let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());172173 let nix_args = std::env::var_os("NIX_ARGS")174 .map(|a| extra_args::parse_os(&a))175 .transpose()?176 .unwrap_or_default();177 let opts = RootOpts::parse();178 let config = opts.fleet_opts.build(nix_args).await?;179180 match run_command(&config, opts.command).await {181 Ok(()) => {182 config.save()?;183 Ok(())184 }185 Err(e) => {186 let _ = config.save();187 Err(e)188 }189 }190}