git.delta.rocks / jrsonnet / refs/commits / b7ea1908bd32

1use std::collections::HashMap;2use std::marker::PhantomData;3use std::mem::forget;4use std::sync::{Arc, RwLock};5use std::time::Duration;67use futures::StreamExt as _;8use futures::stream::empty;9use futures_core::stream;10use prodash::messages::MessageLevel;11use prodash::render::tui::{self, ticker};12use prodash::tree::{Item, Root};13use prodash::unit::{Kind, label};14use prodash::{Progress, Unit};15use tokio::join;16use tokio::time::sleep;17use tracing::{Event, Instrument, Level, Span, Subscriber, info, info_span, span, trace};18use tracing_subscriber::EnvFilter;19use tracing_subscriber::layer::{self, SubscriberExt as _};20use tracing_subscriber::registry::LookupSpan;21use tracing_subscriber::util::SubscriberInitExt;2223pub(crate) struct WithContext(fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut Item)));2425pub struct ProdashLayer<S> {26	root: Arc<Root>,27	root_logs: Item,28	with_context: WithContext,2930	_marker: PhantomData<S>,31}3233impl<S> ProdashLayer<S>34where35	S: Subscriber + for<'a> LookupSpan<'a>,36{37	pub fn new(root: Arc<Root>) -> Self {38		let root_logs = root.add_child("root");39		Self {40			root,41			root_logs,42			with_context: WithContext(Self::get_context),43			_marker: PhantomData,44		}45	}4647	fn get_context(dispatch: &tracing::Dispatch, id: &span::Id, f: &mut dyn FnMut(&mut Item)) {48		let subscriber = dispatch49			.downcast_ref::<S>()50			.expect("subscriber should downcast to expected type");51		let span = subscriber.span(id).expect("span should be in context");52		let span = span.as_ref().map(|s| s.extensions());53		let span = span54			.as_ref()55			.map(|e| e.get::<Item>().expect("existence checked"));56	}57}5859impl<S> layer::Layer<S> for ProdashLayer<S>60where61	S: Subscriber + for<'a> LookupSpan<'a>,62{63	fn on_event(&self, event: &Event<'_>, ctx: layer::Context<'_, S>) {64		let span = ctx.current_span();65		let cur_span = span66			.id()67			.and_then(|s| ctx.span_scope(s))68			.and_then(|mut scope| {69				scope.find(|span| {70					let ext = span.extensions();71					ext.get::<Item>().is_some()72				})73			});74		let cur_span = cur_span.as_ref().map(|s| s.extensions());75		let cur_span = cur_span76			.as_ref()77			.map(|e| e.get::<Item>().expect("existence checked"));7879		let cur_span = cur_span.unwrap_or(&self.root_logs);80		cur_span.info("hello".to_owned());81	}82	fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: layer::Context<'_, S>) {83		let span = ctx.span(id).expect("span should be in context");84		let mut ext = span.extensions_mut();85		ext.insert(self.with_context);8687		let mut spans = self.spans.write().expect("not poisoned");88		let name = attrs.metadata().name();89		if let Some(parent) = attrs90			.parent()91			.cloned()92			.or_else(|| {93				if attrs.is_contextual() {94					ctx.current_span().id().cloned()95				} else {96					None97				}98			})99			.and_then(|v| spans.get_mut(&v))100		{101			let child = parent.add_child(name);102			spans.insert(id.clone(), child);103		} else {104			let child = self.root.add_child(name);105			spans.insert(id.clone(), child);106		};107	}108	fn on_close(&self, id: span::Id, _ctx: layer::Context<'_, S>) {109		let mut spans = self.spans.write().expect("not poisoned");110		spans.remove(&id);111	}112}113114// #[tokio::test]115// async fn test() {116// 	let root: Arc<Root> = prodash::tree::root::Options {117// 		message_buffer_capacity: 1000,118// 		..Default::default()119// 	}120// 	.create()121// 	.into();122//123// 	let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info"));124// 	let reg = tracing_subscriber::registry()125// 		.with(filter)126// 		.with(ProdashLayer::new(root.clone()));127// 	reg.init();128//129// 	let render = tui::render(130// 		std::io::stdout(),131// 		Arc::downgrade(&root),132// 		tui::Options {133// 			frames_per_second: 10.0,134// 			..tui::Options::default()135// 		},136// 	)137// 	.expect("render");138//139// 	let render = tokio::task::spawn(render);140//141// 	loop {142// 		info!("Hello, world!");143// 		sleep(Duration::from_secs(3))144// 			.instrument(info_span!("sleeping root"))145// 			.await;146// 		async {147// 			sleep(Duration::from_secs(3))148// 				.instrument(info_span!("sleeping 3"))149// 				.await;150// 			sleep(Duration::from_secs(2))151// 				.instrument(info_span!("sleeping 2"))152// 				.await;153// 		}154// 		.instrument(info_span!("sleep parent"))155// 		.await;156// 	}157//158// 	// loop {159// 	// 	sleep(Duration::from_secs(3));160// 	render.await;161// }