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}113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161