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::time::Duration;1617use anyhow::{bail, Result};18use clap::Parser;1920use cmds::{build_systems::BuildSystems, info::Info, secrets::Secret};21use futures::future::LocalBoxFuture;22use futures::stream::FuturesUnordered;23use futures::TryStreamExt;24use host::{Config, FleetOpts};25use human_repr::HumanCount;26use indicatif::{ProgressState, ProgressStyle};27use tracing::info;28use tracing::{info_span, Instrument};29use tracing_indicatif::IndicatifLayer;30use tracing_subscriber::{prelude::*, EnvFilter};3132use crate::command::MyCommand;3334#[derive(Parser)]35struct Prefetch {}36impl Prefetch {37 async fn run(&self, config: &Config) -> Result<()> {38 let mut prefetch_dir = config.directory.to_path_buf();39 prefetch_dir.push("prefetch");40 if !prefetch_dir.is_dir() {41 info!("nothing to prefetch: no prefetch directory");42 return Ok(());43 }44 let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();45 for entry in std::fs::read_dir(&prefetch_dir)? {46 tasks.push(Box::pin(async {47 let entry = entry?;48 if !entry.metadata()?.is_file() {49 bail!("only files should exist in prefetch directory");50 }51 let span = info_span!(52 "prefetching",53 name = entry.file_name().to_string_lossy().as_ref()54 );55 let mut path = OsString::new();56 path.push("file://");57 path.push(entry.path());5859 let mut status = MyCommand::new("nix");60 status.arg("store").arg("prefetch-file").arg(path);61 status.run_nix_string().instrument(span).await?;62 Ok(())63 }));64 }65 tasks.try_collect::<Vec<()>>().await?;66 Ok(())67 }68}6970#[derive(Parser)]71enum Opts {72 73 BuildSystems(BuildSystems),74 75 #[clap(subcommand)]76 Secret(Secret),77 78 Prefetch(Prefetch),79 80 Info(Info),81}8283#[derive(Parser)]84#[clap(version = "1.0", author)]85struct RootOpts {86 #[clap(flatten)]87 fleet_opts: FleetOpts,88 #[clap(subcommand)]89 command: Opts,90}9192async fn run_command(config: &Config, command: Opts) -> Result<()> {93 match command {94 Opts::BuildSystems(c) => c.run(config).await?,95 Opts::Secret(s) => s.run(config).await?,96 Opts::Info(i) => i.run(config).await?,97 Opts::Prefetch(p) => p.run(config).await?,98 };99 Ok(())100}101102fn setup_logging() {103 let indicatif_layer = IndicatifLayer::new().with_progress_style(104 ProgressStyle::with_template(105 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",106 )107 .unwrap()108 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {109 let Some(len) = state.len() else {110 return;111 };112 let pos = state.pos();113 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());114 })115 .with_key(116 "color_start",117 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {118 let elapsed = state.elapsed();119120 if elapsed > Duration::from_secs(60) {121 122 let _ = write!(writer, "\x1b[{}m", 1 + 30);123 } else if elapsed > Duration::from_secs(30) {124 125 let _ = write!(writer, "\x1b[{}m", 3 + 30);126 }127 },128 )129 .with_key(130 "color_end",131 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {132 if state.elapsed() > Duration::from_secs(30) {133 let _ = write!(writer, "\x1b[0m");134 }135 },136 ),137 );138139 let filter = EnvFilter::from_default_env();140141 tracing_subscriber::registry()142 .with(143 tracing_subscriber::fmt::layer()144 .without_time()145 .with_target(false)146 .with_writer(indicatif_layer.get_stderr_writer())147 .with_filter(filter), 148 )149 .with(indicatif_layer)150 .init();151}152153#[tokio::main]154async fn main() -> Result<()> {155 setup_logging();156 let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());157158 let nix_args = std::env::var_os("NIX_ARGS")159 .map(|a| extra_args::parse_os(&a))160 .transpose()?161 .unwrap_or_default();162 let opts = RootOpts::parse();163 let config = opts.fleet_opts.build(nix_args).await?;164165 match run_command(&config, opts.command).await {166 Ok(()) => {167 config.save()?;168 Ok(())169 }170 Err(e) => {171 let _ = config.save();172 Err(e)173 }174 }175}