git.delta.rocks / fleet / refs/commits / b66112ef1437

1#![recursion_limit = "512"]23pub(crate) mod cmds;4mod log_tree;56use std::{process::ExitCode, sync::Arc};78use anyhow::{Context as _, Result, bail};9use camino::Utf8PathBuf;10use clap::{CommandFactory, Parser};11use cmds::{12	build_systems::{BuildSystems, Deploy},13	complete::Complete,14	info::Info,15	rollback::RollbackSingle,16	secrets::Secret,17	tf::Tf,18};19use fleet_base::{host::Config, opts::FleetOpts};20use futures::{TryStreamExt, stream::FuturesUnordered};21#[cfg(feature = "indicatif")]22use human_repr::HumanCount;23#[cfg(feature = "indicatif")]24use indicatif::{ProgressState, ProgressStyle};25use log_tree::TreeLayer;26use nix_eval::{27	eval_store, gc_register_my_thread, gc_unregister_my_thread, init_libraries, init_tokio_for_nix,28};29use opentelemetry::trace::TracerProvider;30use opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge;31use opentelemetry_exporter_env::{32	OtlpBaseSettings, OtlpLogsSettings, OtlpTracesSettings, ResolvedOtlpSettings,33};34use opentelemetry_sdk::{logs::SdkLoggerProvider, trace::SdkTracerProvider};35use tokio::task::spawn_blocking;36use tracing::{Instrument, error, info, info_span};37#[cfg(feature = "indicatif")]38use tracing_indicatif::IndicatifLayer;39use tracing_subscriber::{EnvFilter, prelude::*};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::new();52		for entry in std::fs::read_dir(&prefetch_dir)? {53			let entry = entry?;54			if !entry.metadata()?.is_file() {55				bail!("only files should exist in prefetch directory");56			}57			let name = entry.file_name().to_string_lossy().into_owned();58			let path =59				Utf8PathBuf::try_from(entry.path()).context("prefetch path should be utf8")?;60			let span = info_span!("prefetching", name = %name);61			tasks.push(async move {62				let store = eval_store();63				let added = spawn_blocking(move || store.add_file(&name, &path))64					.instrument(span.clone())65					.await??;66				let _g = span.enter();67				info!("{} -> {}", added.hash, added.store_path);68				anyhow::Ok(())69			});70		}71		tasks.try_collect::<Vec<()>>().await?;72		Ok(())73	}74}7576#[derive(Parser)]77enum Opts {78	/// Build system closures79	BuildSystems(BuildSystems),80	/// Upload and switch system closures81	Deploy(Deploy),82	/// Rollback remote machine by redeploying old generation as the new one83	RollbackSingle(RollbackSingle),84	/// Secret management85	#[clap(subcommand)]86	Secret(Secret),87	/// Upload prefetch directory to the nix store88	Prefetch(Prefetch),89	/// Config parsing90	Info(Info),91	/// Command completions92	#[clap(hide(true))]93	Complete(Complete),94	/// Compile and evaluate terranix configuration95	Tf(Tf),96}9798#[derive(Parser)]99#[clap(version, author)]100struct RootOpts {101	#[clap(flatten)]102	fleet_opts: FleetOpts,103	#[clap(subcommand)]104	command: Opts,105	#[clap(long, next_help_heading = "Telemetry", env = "OTEL_FLEET")]106	otel: bool,107	#[clap(flatten)]108	otlp_base: OtlpBaseSettings,109	#[clap(flatten)]110	otel_logs: OtlpLogsSettings,111	#[clap(flatten)]112	otel_traces: OtlpTracesSettings,113}114115async fn run_command(config: &Config, opts: FleetOpts, command: Opts) -> Result<()> {116	match command {117		Opts::BuildSystems(c) => c.run(config, &opts).await?,118		Opts::Deploy(d) => d.run(config, &opts).await?,119		Opts::RollbackSingle(r) => r.run(config, &opts).await?,120		Opts::Secret(s) => s.run(config, &opts).await?,121		Opts::Info(i) => i.run(config).await?,122		Opts::Prefetch(p) => p.run(config).await?,123		Opts::Tf(t) => t.run(config).await?,124		// TODO: actually parse commands before starting the async runtime125		Opts::Complete(c) => spawn_blocking(move || c.run(RootOpts::command())).await?,126	};127	Ok(())128}129130fn setup_logging(opts: &RootOpts) -> Result<()> {131	#[cfg(feature = "indicatif")]132	let indicatif_layer = {133		use std::fmt;134		use std::time::Duration;135136		IndicatifLayer::new().with_max_progress_bars(10, Some(ProgressStyle::default_spinner()))137			.with_span_child_prefix_indent("  ")138			.with_span_child_prefix_symbol("")139			.with_progress_style(140			ProgressStyle::with_template(141				"{span_child_prefix:.magenta}{chevron:.magenta}{span_name:<18.blue.bold} {wide_msg} {span_fields:.dim} {color_start}{download_progress} {elapsed}{color_end}",142			)143				.unwrap()144				.with_key("chevron", |_: &ProgressState, writer: &mut dyn fmt::Write| {145					let _ = write!(writer, "> ");146				})147				.with_key("download_progress", |state: &ProgressState, writer: &mut dyn fmt::Write| {148					let Some(len) = state.len() else {149						return;150					};151					let pos = state.pos();152					if pos > len {153						let _ = write!(writer, "{}", pos.human_count_bare());154					} else {155						let _ = write!(writer, "{} / {}", pos.human_count_bare(), len.human_count_bare());156					}157				})158				.with_key(159					"color_start",160					|state: &ProgressState, writer: &mut dyn fmt::Write| {161						let elapsed = state.elapsed();162163						if elapsed > Duration::from_secs(60) {164							// Red165							let _ = write!(writer, "\x1b[{}m", 1 + 30);166						} else if elapsed > Duration::from_secs(30) {167							// Yellow168							let _ = write!(writer, "\x1b[{}m", 3 + 30);169						}170					},171				)172				.with_key(173					"color_end",174					|state: &ProgressState, writer: &mut dyn fmt::Write| {175						if state.elapsed() > Duration::from_secs(30) {176							let _ = write!(writer, "\x1b[0m");177						}178					},179				),180		)181	};182183	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("noq_udp=warn,info"));184185	let tree = {186		#[cfg(feature = "indicatif")]187		let writer = indicatif_layer.get_stderr_writer();188		#[cfg(not(feature = "indicatif"))]189		let writer = || std::io::stderr();190		TreeLayer::new(writer)191	};192193	let reg = tracing_subscriber::registry().with(tree.with_filter(filter));194195	#[cfg(feature = "indicatif")]196	let reg = reg.with(indicatif_layer);197198	if opts.otel {199		let traces = ResolvedOtlpSettings::traces(&opts.otlp_base, &opts.otel_traces)?;200		let span_exporter = traces.span_exporter()?;201		let logs = ResolvedOtlpSettings::logs(&opts.otlp_base, &opts.otel_logs)?;202		let log_exporter = logs.log_exporter()?;203204		let span_provider = SdkTracerProvider::builder()205			.with_batch_exporter(span_exporter)206			.build();207		let log_provider = SdkLoggerProvider::builder()208			.with_batch_exporter(log_exporter)209			.build();210211		let logger = OpenTelemetryTracingBridge::new(&log_provider);212		let tracer = span_provider.tracer("fleet");213214		reg.with(tracing_opentelemetry::layer().with_tracer(tracer))215			.with(logger)216			.init();217	} else {218		reg.init();219	}220221	Ok(())222}223224fn main() -> ExitCode {225	let opts = RootOpts::parse();226	if let Opts::Complete(c) = &opts.command {227		c.run(RootOpts::command());228		return ExitCode::SUCCESS;229	}230231	if let Err(e) = setup_logging(&opts) {232		eprintln!("{e:#}");233		return ExitCode::FAILURE;234	}235236	init_libraries();237238	let runtime = tokio::runtime::Builder::new_multi_thread()239		.enable_all()240		.on_thread_start(|| {241			gc_register_my_thread();242		})243		.on_thread_stop(|| {244			gc_unregister_my_thread();245		})246		.build()247		.expect("failed to build runtime");248	let runtime = Arc::new(runtime);249250	init_tokio_for_nix(runtime.clone());251252	runtime.block_on(async {253		tokio::task::spawn(async move {254			if let Err(e) = main_real(opts).await {255				error!("{e:#}");256				ExitCode::FAILURE257			} else {258				ExitCode::SUCCESS259			}260		})261		.await262		.expect("primary task panicked")263	})264}265266async fn main_real(opts: RootOpts) -> Result<()> {267	let config = opts.fleet_opts.build(matches!(268		opts.command,269		Opts::Deploy(_) | Opts::BuildSystems(_)270	))?;271272	match run_command(&config, opts.fleet_opts, opts.command).await {273		Ok(()) => {274			config.save()?;275			Ok(())276		}277		Err(e) => {278			let _ = config.save();279			Err(e)280		}281	}282}283284#[cfg(test)]285mod tests {286	use super::*;287288	#[test]289	fn verify_command() {290		use clap::CommandFactory;291		RootOpts::command().debug_assert();292	}293}