git.delta.rocks / jrsonnet / refs/commits / 754b45cdacd0

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};29#[cfg(feature = "indicatif")]30use human_repr::HumanCount;31#[cfg(feature = "indicatif")]32use indicatif::{ProgressState, ProgressStyle};33use tracing::{error, info};34use tracing::{info_span, Instrument};35#[cfg(feature = "indicatif")]36use tracing_indicatif::IndicatifLayer;37use tracing_subscriber::{prelude::*, EnvFilter};3839use crate::command::MyCommand;4041#[derive(Parser)]42struct Prefetch {}43impl Prefetch {44	async fn run(&self, config: &Config) -> Result<()> {45		let mut prefetch_dir = config.directory.to_path_buf();46		prefetch_dir.push("prefetch");47		if !prefetch_dir.is_dir() {48			info!("nothing to prefetch: no prefetch directory");49			return Ok(());50		}51		let tasks = <FuturesUnordered<LocalBoxFuture<Result<()>>>>::new();52		for entry in std::fs::read_dir(&prefetch_dir)? {53			tasks.push(Box::pin(async {54				let entry = entry?;55				if !entry.metadata()?.is_file() {56					bail!("only files should exist in prefetch directory");57				}58				let span = info_span!(59					"prefetching",60					name = entry.file_name().to_string_lossy().as_ref()61				);62				let mut path = OsString::new();63				path.push("file://");64				path.push(entry.path());6566				let mut status = MyCommand::new("nix");67				status.args(&config.nix_args);68				status.arg("store").arg("prefetch-file").arg(path);69				status.run_nix_string().instrument(span).await?;70				Ok(())71			}));72		}73		tasks.try_collect::<Vec<()>>().await?;74		Ok(())75	}76}7778#[derive(Parser)]79enum Opts {80	/// Prepare systems for deployments81	BuildSystems(BuildSystems),8283	Deploy(Deploy),84	/// Secret management85	#[clap(subcommand)]86	Secret(Secret),87	/// Upload prefetch directory to the nix store88	Prefetch(Prefetch),89	/// Config parsing90	Info(Info),91}9293#[derive(Parser)]94#[clap(version, author)]95struct RootOpts {96	#[clap(flatten)]97	fleet_opts: FleetOpts,98	#[clap(subcommand)]99	command: Opts,100}101102async fn run_command(config: &Config, command: Opts) -> Result<()> {103	match command {104		Opts::BuildSystems(c) => c.run(config).await?,105		Opts::Deploy(d) => d.run(config).await?,106		Opts::Secret(s) => s.run(config).await?,107		Opts::Info(i) => i.run(config).await?,108		Opts::Prefetch(p) => p.run(config).await?,109	};110	Ok(())111}112113fn setup_logging() {114	#[cfg(feature = "indicatif")]115	let indicatif_layer = IndicatifLayer::new().with_progress_style(116		ProgressStyle::with_template(117			"{color_start}{span_child_prefix} {span_name}{{{span_fields}}}{color_end} {wide_msg} {color_start}{download_progress} {elapsed}{color_end}",118		)119		.unwrap()120		.with_key("download_progress", |state: &ProgressState, writer: &mut dyn std::fmt::Write| {121			let Some(len) = state.len() else {122				return;123			};124			let pos = state.pos();125			if pos > len {126				let _ = write!(writer, "{}", pos.human_count_bare());127			} else {128				let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());129			}130		})131		.with_key(132			"color_start",133			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {134				use std::time::Duration;135				let elapsed = state.elapsed();136137				if elapsed > Duration::from_secs(60) {138					// Red139					let _ = write!(writer, "\x1b[{}m", 1 + 30);140				} else if elapsed > Duration::from_secs(30) {141					// Yellow142					let _ = write!(writer, "\x1b[{}m", 3 + 30);143				}144			},145		)146		.with_key(147			"color_end",148			|state: &ProgressState, writer: &mut dyn std::fmt::Write| {149				if state.elapsed() > Duration::from_secs(30) {150					let _ = write!(writer, "\x1b[0m");151				}152			},153		),154	);155156	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));157158	let reg = tracing_subscriber::registry().with({159		let sub = tracing_subscriber::fmt::layer()160			.without_time()161			.with_target(true);162		#[cfg(feature = "indicatif")]163		let sub = sub.with_writer(indicatif_layer.get_stdout_writer());164		sub.with_filter(filter) // .withou,165	});166	// #[cfg(feature = "indicatif")]167	#[cfg(feature = "indicatif")]168	let reg = reg.with(indicatif_layer);169	reg.init();170}171172#[tokio::main]173async fn main() -> ExitCode {174	setup_logging();175	if let Err(e) = main_real().await {176		// If I remove this line, the next error!() line gets eaten.177		info!("fixme: this line gets eaten by tracing-indicatif on levels info+");178		error!("{e:#}");179		return ExitCode::FAILURE;180	}181	ExitCode::SUCCESS182}183184async fn main_real() -> Result<()> {185	let _ = better_nix_eval::TOKIO_RUNTIME.set(tokio::runtime::Handle::current());186187	let nix_args = std::env::var_os("NIX_ARGS")188		.map(|a| extra_args::parse_os(&a))189		.transpose()?190		.unwrap_or_default();191	let opts = RootOpts::parse();192	let config = opts.fleet_opts.build(nix_args).await?;193194	match run_command(&config, opts.command).await {195		Ok(()) => {196			config.save()?;197			Ok(())198		}199		Err(e) => {200			let _ = config.save();201			Err(e)202		}203	}204}