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::time::Duration;15use std::{ffi::OsString, process::ExitCode};1617use anyhow::{bail, Result};18use clap::Parser;1920use cmds::{21 build_systems::{BuildSystems, Deploy},22 info::Info,23 secrets::Secret,24};25use futures::future::LocalBoxFuture;26use futures::stream::FuturesUnordered;27use futures::TryStreamExt;28use host::{Config, FleetOpts};29use human_repr::HumanCount;30use indicatif::{ProgressState, ProgressStyle};31use tracing::{error, info};32use tracing::{info_span, Instrument};33use tracing_indicatif::IndicatifLayer;34use tracing_subscriber::{prelude::*, EnvFilter};3536use crate::command::MyCommand;3738#[derive(Parser)]39struct Prefetch {}40impl Prefetch {41 async fn run(&self, config: &Config) -> Result<()> {42 let mut prefetch_dir = config.directory.to_path_buf();43 prefetch_dir.push("prefetch");44 if !prefetch_dir.is_dir() {45 info!("nothing to prefetch: no prefetch directory");46 return Ok(());47 }48 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();49 for entry in std::fs::read_dir(&prefetch_dir)? {50 tasks.push(Box::pin(async {51 let entry = entry?;52 if !entry.metadata()?.is_file() {53 bail!("only files should exist in prefetch directory");54 }55 let span = info_span!(56 "prefetching",57 name = entry.file_name().to_string_lossy().as_ref()58 );59 let mut path = OsString::new();60 path.push("file://");61 path.push(entry.path());6263 let mut status = MyCommand::new("nix");64 status.args(&config.nix_args);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 if pos > len {122 let _ = write!(writer, "{}", pos.human_count_bare());123 } else {124 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());125 }126 })127 .with_key(128 "color_start",129 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {130 let elapsed = state.elapsed();131132 if elapsed > Duration::from_secs(60) {133 134 let _ = write!(writer, "\x1b[{}m", 1 + 30);135 } else if elapsed > Duration::from_secs(30) {136 137 let _ = write!(writer, "\x1b[{}m", 3 + 30);138 }139 },140 )141 .with_key(142 "color_end",143 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {144 if state.elapsed() > Duration::from_secs(30) {145 let _ = write!(writer, "\x1b[0m");146 }147 },148 ),149 );150151 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));152153 tracing_subscriber::registry()154 .with(155 tracing_subscriber::fmt::layer()156 .without_time()157 .with_target(true)158 .with_writer(indicatif_layer.get_stdout_writer())159 .with_filter(filter), 160 )161 .with(indicatif_layer)162 .init();163}164165#[tokio::main]166async fn main() -> ExitCode {167 setup_logging();168 if let Err(e) = main_real().await {169 170 info!("fixme: this line gets eaten by tracing-indicatif on levels info+");171 error!("{e:#}");172 return ExitCode::FAILURE;173 }174 ExitCode::SUCCESS175}176177async fn main_real() -> Result<()> {178 let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());179180 let nix_args = std::env::var_os("NIX_ARGS")181 .map(|a| extra_args::parse_os(&a))182 .transpose()?183 .unwrap_or_default();184 let opts = RootOpts::parse();185 let config = opts.fleet_opts.build(nix_args).await?;186187 match run_command(&config, opts.command).await {188 Ok(()) => {189 config.save()?;190 Ok(())191 }192 Err(e) => {193 let _ = config.save();194 Err(e)195 }196 }197}