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::io::{stderr, stdout, Write};16use std::process::exit;17use std::time::Duration;1819use anyhow::{bail, Result};20use clap::Parser;2122use cmds::{build_systems::BuildSystems, info::Info, secrets::Secret};23use futures::future::LocalBoxFuture;24use futures::stream::FuturesUnordered;25use futures::TryStreamExt;26use host::{Config, FleetOpts};27use human_repr::HumanCount;28use indicatif::{ProgressState, ProgressStyle};29use tracing::{error, info};30use tracing::{info_span, Instrument};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.arg("store").arg("prefetch-file").arg(path);63 status.run_nix_string().instrument(span).await?;64 Ok(())65 }));66 }67 tasks.try_collect::<Vec<()>>().await?;68 Ok(())69 }70}7172#[derive(Parser)]73enum Opts {74 75 BuildSystems(BuildSystems),76 77 #[clap(subcommand)]78 Secret(Secret),79 80 Prefetch(Prefetch),81 82 Info(Info),83}8485#[derive(Parser)]86#[clap(version, author)]87struct RootOpts {88 #[clap(flatten)]89 fleet_opts: FleetOpts,90 #[clap(subcommand)]91 command: Opts,92}9394async fn run_command(config: &Config, command: Opts) -> Result<()> {95 match command {96 Opts::BuildSystems(c) => c.run(config).await?,97 Opts::Secret(s) => s.run(config).await?,98 Opts::Info(i) => i.run(config).await?,99 Opts::Prefetch(p) => p.run(config).await?,100 };101 Ok(())102}103104fn setup_logging() {105 let indicatif_layer = IndicatifLayer::new().with_progress_style(106 ProgressStyle::with_template(107 "{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",108 )109 .unwrap()110 .with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {111 let Some(len) = state.len() else {112 return;113 };114 let pos = state.pos();115 let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());116 })117 .with_key(118 "color_start",119 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {120 let elapsed = state.elapsed();121122 if elapsed > Duration::from_secs(60) {123 124 let _ = write!(writer, "\x1b[{}m", 1 + 30);125 } else if elapsed > Duration::from_secs(30) {126 127 let _ = write!(writer, "\x1b[{}m", 3 + 30);128 }129 },130 )131 .with_key(132 "color_end",133 |state: &ProgressState, writer: &mut dyn std::fmt::Write| {134 if state.elapsed() > Duration::from_secs(30) {135 let _ = write!(writer, "\x1b[0m");136 }137 },138 ),139 );140141 let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));142143 tracing_subscriber::registry()144 .with(145 tracing_subscriber::fmt::layer()146 .without_time()147 .with_target(true)148 .with_writer(indicatif_layer.get_stderr_writer())149 .with_filter(filter), 150 )151 .with(indicatif_layer)152 .init();153}154155#[tokio::main]156async fn main() {157 setup_logging();158 if let Err(e) = main_real().await {159 error!("{e:#}");160 exit(1);161 }162}163164async fn main_real() -> Result<()> {165 let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());166167 let nix_args = std::env::var_os("NIX_ARGS")168 .map(|a| extra_args::parse_os(&a))169 .transpose()?170 .unwrap_or_default();171 let opts = RootOpts::parse();172 let config = opts.fleet_opts.build(nix_args).await?;173174 match run_command(&config, opts.command).await {175 Ok(()) => {176 config.save()?;177 Ok(())178 }179 Err(e) => {180 let _ = config.save();181 Err(e)182 }183 }184}