git.delta.rocks / fleet / refs/commits / 9766610438d2

1//! Streaming, indented tree formatter.23use std::{4	fmt::Write as _,5	io::{IsTerminal as _, Write as _},6	sync::{7		Mutex,8		atomic::{AtomicBool, Ordering},9	},10	time::{Duration, Instant},11};1213use tracing::{Event, Level, Subscriber, field::Field, field::Visit, span::Attributes, span::Id};14use tracing_subscriber::{15	fmt::MakeWriter,16	layer::{Context, Layer},17	registry::{LookupSpan, Scope},18};1920/// Width of the right-aligned target column.21const NAME_LIMIT: usize = 18;2223#[derive(Default)]24struct Fields {25	message: String,26	rest: String,27}2829impl Visit for Fields {30	fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) {31		if field.name() == "message" {32			let _ = write!(self.message, "{value:?}");33		} else {34			let _ = write!(self.rest, " {}={:?}", field.name(), value);35		}36	}37}3839impl Fields {40	fn render(&self) -> String {41		format!("{}{}", self.message, self.rest).trim().to_owned()42	}43}4445struct SpanData {46	fields: String,47	started: Instant,48	opened: AtomicBool,49}5051fn level_code(level: &Level) -> &'static str {52	match *level {53		// Red54		Level::ERROR => "31m",55		// Yellow56		Level::WARN => "33m",57		// Blue58		Level::INFO => "34m",59		// Default fg60		Level::DEBUG => "39m",61		// Gray62		Level::TRACE => "90m",63	}64}6566fn gutter(depth: usize, sym: char) -> String {67	let mut s = "  ".repeat(depth);68	s.push(sym);69	s70}7172fn fmt_dur(d: Duration) -> String {73	if d.as_secs_f64() >= 1.0 {74		format!("{:.2}s", d.as_secs_f64())75	} else if d.as_millis() >= 1 {76		format!("{}ms", d.as_millis())77	} else {78		format!("{}µs", d.as_micros())79	}80}8182pub struct TreeLayer<W> {83	writer: W,84	color: bool,85	last: Mutex<Vec<Id>>,86}8788impl<W> TreeLayer<W> {89	pub fn new(writer: W) -> Self {90		Self {91			writer,92			color: std::io::stderr().is_terminal() && std::env::var_os("NO_COLOR").is_none(),93			last: Mutex::new(vec![]),94		}95	}9697	fn sgr(&self, code: &str) -> String {98		if self.color {99			format!("\x1b[{code}")100		} else {101			String::new()102		}103	}104105	fn name_col(&self, mut target: &str, code: &str) -> String {106		let w = NAME_LIMIT;107		if target.len() > w {108			target = &target[..w];109		}110		if self.color {111			format!("\x1b[{code}\x1b[1m{target:>w$}\x1b[0m")112		} else {113			format!("{target:>w$}")114		}115	}116117	fn print_ancestors<S>(&self, scope: Option<Scope<'_, S>>, out: &mut String) -> usize118	where119		S: Subscriber + for<'a> LookupSpan<'a>,120	{121		let scope: Vec<_> = scope.map(|s| s.from_root().collect()).unwrap_or_default();122123		let mut _last = self.last.lock().expect("not poisoned");124125		let skip = scope126			.iter()127			.enumerate()128			.take_while(|(i, s)| _last.get(*i) == Some(&s.id()))129			.count();130131		for (depth, span) in scope.iter().enumerate().skip(skip) {132			let ext = span.extensions();133			let Some(data) = ext.get::<SpanData>() else {134				continue;135			};136			let newly_open = !data.opened.swap(true, Ordering::Relaxed);137			let _ = writeln!(138				out,139				"{} {}{}{} {}{}{}",140				self.name_col(span.metadata().target(), "44m"),141				self.sgr("35m"),142				gutter(depth, if newly_open { '>' } else { '.' }),143				self.sgr("1m"),144				span.name(),145				data.fields,146				self.sgr("0m"),147			);148		}149150		*_last = scope.iter().map(|v| v.id()).collect();151152		scope.len()153	}154}155156impl<S, W> Layer<S> for TreeLayer<W>157where158	S: Subscriber + for<'a> LookupSpan<'a>,159	W: for<'a> MakeWriter<'a> + 'static,160{161	fn on_new_span(&self, attrs: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {162		let span = ctx.span(id).expect("new span exists");163		let mut fields = Fields::default();164		attrs.record(&mut fields);165		span.extensions_mut().insert(SpanData {166			fields: fields.rest,167			started: Instant::now(),168			opened: AtomicBool::new(false),169		});170	}171172	fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {173		let mut out = String::new();174175		let scope = ctx.event_scope(event);176		let pad = self.print_ancestors(scope, &mut out);177178		let mut fields = Fields::default();179		event.record(&mut fields);180		let message = fields.render();181		let meta = event.metadata();182		for (i, message) in message.split('\n').enumerate() {183			// Event line: black bg + colored target column, then indent + message.184			let _ = writeln!(185				out,186				"{}{}{}{}{}",187				self.sgr("40m"),188				self.name_col(189					if i == 0 { meta.target() } else { "|" },190					level_code(meta.level())191				),192				self.sgr("0m"),193				gutter(pad, ' '),194				message,195			);196		}197198		let _ = self.writer.make_writer().write_all(out.as_bytes());199	}200201	fn on_close(&self, id: Id, ctx: Context<'_, S>) {202		let Some(span) = ctx.span(&id) else { return };203		let (opened, dur) = span204			.extensions()205			.get::<SpanData>()206			.map_or((false, Duration::ZERO), |d| {207				(d.opened.load(Ordering::Relaxed), d.started.elapsed())208			});209		if !opened {210			return;211		}212213		let mut out = String::new();214		let _pad = self.print_ancestors(ctx.span_scope(&id), &mut out);215216		let depth = span.scope().skip(1).count(); // ancestors == depth217		let _ = writeln!(218			out,219			"{} {}{}{} {}{} (Done in {}){}",220			self.name_col(span.metadata().target(), "44m"),221			self.sgr("35m"),222			gutter(depth, '<'),223			self.sgr("1m"),224			span.name(),225			self.sgr("22m"),226			fmt_dur(dur),227			self.sgr("0m"),228		);229230		let _guard = self.last.lock().expect("not poisoned");231		let _ = self.writer.make_writer().write_all(out.as_bytes());232	}233}