tracing_indicatif/
lib.rs

1//! A [tracing](https://docs.rs/tracing/latest/tracing/) layer that automatically creates and manages [indicatif](https://docs.rs/indicatif/latest/indicatif/index.html) progress bars for active spans.
2//!
3//! Progress bars are a great way to make your CLIs feel more responsive. However,
4//! adding and managing progress bars in your libraries can be invasive, unergonomic,
5//! and difficult to keep track of.
6//!
7//! This library aims to make it easy to show progress bars for your CLI by tying
8//! progress bars to [tracing spans](https://docs.rs/tracing/latest/tracing/#spans).
9//! For CLIs/libraries already using tracing spans, this allow for a dead simple (3
10//! line) code change to enable a smooth progress bar experience for your program.
11//! This eliminates having to have code in your libraries to manually manage
12//! progress bar instances.
13//!
14//! This ends up working quite well as progress bars are fundamentally tracking the
15//! lifetime of some "span" (whether that "span" is defined explicitly or implicitly),
16//! so might as well make that relationship explicit.
17//!
18//! An easy quick start for this crate is:
19//! ```rust
20//! use tracing_indicatif::IndicatifLayer;
21//! use tracing_subscriber::layer::SubscriberExt;
22//! use tracing_subscriber::util::SubscriberInitExt;
23//!
24//! let indicatif_layer = IndicatifLayer::new();
25//!
26//! tracing_subscriber::registry()
27//!     .with(tracing_subscriber::fmt::layer().with_writer(indicatif_layer.get_stderr_writer()))
28//!     .with(indicatif_layer)
29//!     .init();
30//! ```
31//! See [`IndicatifLayer`] for additional documentation.
32//!
33//! See the [`examples`](https://github.com/emersonford/tracing-indicatif/tree/main/examples) folder for examples of how to customize the layer / progress bar
34//! appearance.
35//!
36//! Note: it is highly recommended you pass `indicatif_layer.get_stderr_writer()` or
37//! `indicatif_layer.get_stdout_writer()` to your `fmt::layer()` (depending on where you want to
38//! emit tracing logs) to prevent progress bars from clobbering any console logs.
39use std::any::TypeId;
40use std::marker::PhantomData;
41use std::sync::Mutex;
42use std::time::Duration;
43
44use indicatif::MultiProgress;
45use indicatif::ProgressBar;
46/// Re-export of [`indicatif`]'s style module for ease of use.
47pub use indicatif::style;
48use indicatif::style::ProgressStyle;
49use indicatif::style::ProgressTracker;
50use tracing_core::Subscriber;
51use tracing_core::span;
52use tracing_subscriber::fmt::FormatFields;
53use tracing_subscriber::fmt::FormattedFields;
54use tracing_subscriber::fmt::format::DefaultFields;
55use tracing_subscriber::layer;
56use tracing_subscriber::registry::LookupSpan;
57
58pub mod filter;
59mod pb_manager;
60pub mod span_ext;
61pub mod util;
62pub mod writer;
63
64use pb_manager::ProgressBarManager;
65pub use pb_manager::TickSettings;
66#[doc(inline)]
67pub use writer::IndicatifWriter;
68
69#[derive(Clone)]
70struct IndicatifProgressKey {
71    message: String,
72}
73
74impl ProgressTracker for IndicatifProgressKey {
75    fn clone_box(&self) -> Box<dyn ProgressTracker> {
76        Box::new(self.clone())
77    }
78
79    fn tick(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
80
81    fn reset(&mut self, _: &indicatif::ProgressState, _: std::time::Instant) {}
82
83    fn write(&self, _: &indicatif::ProgressState, w: &mut dyn std::fmt::Write) {
84        let _ = w.write_str(&self.message);
85    }
86}
87
88// Suppose we have a [Span] (maybe gotten via [Span::current]) and want access to our
89// [IndicatifLayer] instance from it. The way to do this would be something like
90// ```
91// span.with_subscriber(|(id, subscriber)| {
92//   let maybe_layer = subscriber.downcast_ref::<IndicatifLayer<S, F>>();
93//   ...
94// });
95// ```
96// but this has the problem that, because `IndicatifLayer` has generic params, we need to pass
97// a concrete type `S` and `F` to that `downcast_ref` call. And the callsite doesn't know what
98// those concrete types are.
99//
100// Therefore, we use this `WithContext` struct (along with the defined `downcast_raw` method) to do
101// a form of indirection to something that does already know (or "remembers") what those concrete
102// types `S` and `F` are, so the callsite doesn't need to care about it.
103//
104// This doesn't actually return a reference to our [IndicatifLayer] instance as we only care about
105// the associated span data, so we just pass that to the corresponding `fn`.
106//
107// See:
108// * https://github.com/tokio-rs/tracing/blob/a0126b2e2d465e8e6d514acdf128fcef5b863d27/tracing-error/src/subscriber.rs#L32
109// * https://github.com/tokio-rs/tracing/blob/a0126b2e2d465e8e6d514acdf128fcef5b863d27/tracing-opentelemetry/src/subscriber.rs#L74
110#[allow(clippy::type_complexity)]
111pub(crate) struct WithContext(
112    fn(&tracing::Dispatch, &span::Id, f: &mut dyn FnMut(&mut IndicatifSpanContext)),
113);
114
115#[allow(clippy::type_complexity)]
116pub(crate) struct WithStderrWriter(
117    fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>)),
118);
119
120#[allow(clippy::type_complexity)]
121pub(crate) struct WithStdoutWriter(
122    fn(&tracing::Dispatch, f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>)),
123);
124
125#[allow(clippy::type_complexity)]
126pub(crate) struct WithMultiProgress(fn(&tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)));
127
128impl WithContext {
129    pub(crate) fn with_context(
130        &self,
131        dispatch: &tracing::Dispatch,
132        id: &span::Id,
133        mut f: impl FnMut(&mut IndicatifSpanContext),
134    ) {
135        (self.0)(dispatch, id, &mut f)
136    }
137}
138
139impl WithStderrWriter {
140    pub(crate) fn with_context(
141        &self,
142        dispatch: &tracing::Dispatch,
143        mut f: impl FnMut(IndicatifWriter<writer::Stderr>),
144    ) {
145        (self.0)(dispatch, &mut f)
146    }
147}
148
149impl WithStdoutWriter {
150    pub(crate) fn with_context(
151        &self,
152        dispatch: &tracing::Dispatch,
153        mut f: impl FnMut(IndicatifWriter<writer::Stdout>),
154    ) {
155        (self.0)(dispatch, &mut f)
156    }
157}
158
159impl WithMultiProgress {
160    pub(crate) fn with_context(
161        &self,
162        dispatch: &tracing::Dispatch,
163        mut f: impl FnMut(MultiProgress),
164    ) {
165        (self.0)(dispatch, &mut f)
166    }
167}
168
169#[derive(Default)]
170struct ProgressBarInitSettings {
171    style: Option<ProgressStyle>,
172    len: Option<u64>,
173    pos: Option<u64>,
174    message: Option<String>,
175}
176
177struct IndicatifSpanContext {
178    // If this progress bar is `Some(pb)` and `pb.is_hidden`, it means the progress bar is queued.
179    // We start the progress bar in hidden mode so things like `elapsed` are accurate.
180    //
181    // If this progress bar is `None`, it means the span has not yet been entered.
182    progress_bar: Option<ProgressBar>,
183    // If `Some`, the progress bar will use this style when the span is entered for the first time.
184    pb_init_settings: ProgressBarInitSettings,
185    // Notes:
186    // * A parent span cannot close before its child spans, so if a parent span has a progress bar,
187    //   that parent progress bar's lifetime will be greater than this span's progress bar.
188    // * The ProgressBar is just a wrapper around `Arc`, so cloning and tracking it here is fine.
189    parent_progress_bar: Option<ProgressBar>,
190    // This is only `Some` if we have some parent with a progress bar.
191    parent_span: Option<span::Id>,
192    // Fields to be passed to the progress bar as keys.
193    span_fields_formatted: Option<String>,
194    span_name: String,
195    span_child_prefix: String,
196    // Used to quickly compute a child span's prefix without having to traverse up the entire span
197    // scope.
198    level: u16,
199    // If `Some`, this is the message that will be displayed when the progress bar is finished.
200    finish_message: Option<String>,
201}
202
203impl IndicatifSpanContext {
204    fn add_keys_to_style(&self, style: ProgressStyle) -> ProgressStyle {
205        style
206            .with_key(
207                "span_name",
208                IndicatifProgressKey {
209                    message: self.span_name.clone(),
210                },
211            )
212            .with_key(
213                "span_fields",
214                IndicatifProgressKey {
215                    message: self.span_fields_formatted.to_owned().unwrap_or_default(),
216                },
217            )
218            .with_key(
219                "span_child_prefix",
220                IndicatifProgressKey {
221                    message: self.span_child_prefix.clone(),
222                },
223            )
224    }
225
226    fn make_progress_bar(&mut self, default_style: &ProgressStyle) {
227        if self.progress_bar.is_none() {
228            let pb = ProgressBar::hidden().with_style(
229                self.pb_init_settings
230                    .style
231                    .take()
232                    .unwrap_or_else(|| self.add_keys_to_style(default_style.clone())),
233            );
234
235            if let Some(len) = self.pb_init_settings.len.take() {
236                pb.set_length(len);
237            }
238
239            if let Some(msg) = self.pb_init_settings.message.take() {
240                pb.set_message(msg);
241            }
242
243            if let Some(pos) = self.pb_init_settings.pos.take() {
244                pb.set_position(pos);
245            }
246
247            self.progress_bar = Some(pb);
248        }
249    }
250
251    fn set_progress_bar_style(&mut self, style: ProgressStyle) {
252        if let Some(ref pb) = self.progress_bar {
253            pb.set_style(self.add_keys_to_style(style));
254        } else {
255            self.pb_init_settings.style = Some(self.add_keys_to_style(style));
256        }
257    }
258
259    fn set_progress_bar_length(&mut self, len: u64) {
260        if let Some(ref pb) = self.progress_bar {
261            pb.set_length(len);
262        } else {
263            self.pb_init_settings.len = Some(len);
264        }
265    }
266
267    fn set_progress_bar_position(&mut self, pos: u64) {
268        if let Some(ref pb) = self.progress_bar {
269            pb.set_position(pos);
270        } else {
271            self.pb_init_settings.pos = Some(pos);
272        }
273    }
274
275    fn set_progress_bar_message(&mut self, msg: String) {
276        if let Some(ref pb) = self.progress_bar {
277            pb.set_message(msg);
278        } else {
279            self.pb_init_settings.message = Some(msg);
280        }
281    }
282
283    fn inc_progress_bar_position(&mut self, pos: u64) {
284        if let Some(ref pb) = self.progress_bar {
285            pb.inc(pos);
286        } else if let Some(ref mut pb_pos) = self.pb_init_settings.pos {
287            *pb_pos += pos;
288        } else {
289            // indicatif defaults position to 0, so copy that behavior.
290            self.pb_init_settings.pos = Some(pos);
291        }
292    }
293
294    fn inc_progress_bar_length(&mut self, len: u64) {
295        if let Some(ref pb) = self.progress_bar {
296            pb.inc_length(len);
297        } else if let Some(ref mut pb_len) = self.pb_init_settings.len {
298            *pb_len += len;
299        }
300    }
301
302    fn progress_bar_tick(&mut self) {
303        if let Some(ref pb) = self.progress_bar {
304            pb.tick()
305        }
306    }
307
308    fn reset_progress_bar(&mut self) {
309        if let Some(ref pb) = self.progress_bar {
310            pb.reset();
311        }
312    }
313
314    fn reset_progress_bar_elapsed(&mut self) {
315        if let Some(ref pb) = self.progress_bar {
316            pb.reset_elapsed();
317        }
318    }
319
320    fn reset_progress_bar_eta(&mut self) {
321        if let Some(ref pb) = self.progress_bar {
322            pb.reset_eta();
323        }
324    }
325
326    fn set_progress_bar_finish_message(&mut self, msg: String) {
327        self.finish_message = Some(msg);
328    }
329
330    fn eta(&self) -> Duration {
331        if let Some(ref pb) = self.progress_bar {
332            pb.eta()
333        } else {
334            Duration::new(0, 0)
335        }
336    }
337
338    fn elapsed(&self) -> Duration {
339        if let Some(ref pb) = self.progress_bar {
340            pb.elapsed()
341        } else {
342            Duration::new(0, 0)
343        }
344    }
345}
346
347/// The layer that handles creating and managing indicatif progress bars for active spans. This
348/// layer must be registered with your tracing subscriber to have any effect.
349///
350/// This layer performs no filtering on which spans to show progress bars for. It is expected one
351/// attaches [filters to this
352/// layer](https://docs.rs/tracing-subscriber/latest/tracing_subscriber/layer/index.html#filtering-with-layers)
353/// to control which spans actually have progress bars generated for them. See
354/// [`filter::IndicatifFilter`] for a rudimentary filter.
355///
356/// Progress bars will be started the very first time a span is [entered](tracing::Span::enter)
357/// or when one of its child spans is entered for the first time, and will finish when the span
358/// is [closed](tracing_subscriber::Layer::on_close) (including all child spans having closed).
359///
360/// Progress bars are emitted to stderr.
361///
362/// Under the hood, this just uses indicatif's [`MultiProgress`] struct to
363/// manage individual [`ProgressBar`] instances per span.
364pub struct IndicatifLayer<S, F = DefaultFields> {
365    pb_manager: Mutex<ProgressBarManager>,
366    // Allows us to fetch the `MultiProgress` without taking a lock.
367    // Do not mutate `mp` directly, always go through `pb_manager`.
368    mp: MultiProgress,
369    span_field_formatter: F,
370    progress_style: ProgressStyle,
371    span_child_prefix_indent: &'static str,
372    span_child_prefix_symbol: &'static str,
373    get_context: WithContext,
374    get_stderr_writer_context: WithStderrWriter,
375    get_stdout_writer_context: WithStdoutWriter,
376    get_multi_progress_context: WithMultiProgress,
377    inner: PhantomData<S>,
378}
379
380impl<S> IndicatifLayer<S>
381where
382    S: Subscriber + for<'a> LookupSpan<'a>,
383{
384    /// Spawns a progress bar for every tracing span that is received by this layer.
385    ///
386    /// The default settings for this layer are 7 progress bars maximum and progress bars in the
387    /// style of:
388    /// ```text
389    /// ⠄ do_work{val=0}
390    /// ⠄ do_work{val=1}
391    /// ⠄ do_work{val=2}
392    ///   ↳ ⠴ do_sub_work{val=2}
393    ///   ↳ ⠴ do_sub_work{val=2}
394    /// ⠄ do_work{val=3}
395    /// ⠄ do_work{val=4}
396    /// ...and 5 more not shown above.
397    /// ```
398    pub fn new() -> Self {
399        Self::default()
400    }
401}
402
403impl<S> Default for IndicatifLayer<S>
404where
405    S: Subscriber + for<'a> LookupSpan<'a>,
406{
407    fn default() -> Self {
408        let pb_manager = ProgressBarManager::new(
409            7,
410            Some(
411                ProgressStyle::with_template(
412                    "...and {pending_progress_bars} more not shown above.",
413                )
414                .expect("valid template"),
415            ),
416            TickSettings::default(),
417        );
418        let mp = pb_manager.mp.clone();
419
420        Self {
421            pb_manager: Mutex::new(pb_manager),
422            mp,
423            span_field_formatter: DefaultFields::new(),
424            progress_style: ProgressStyle::with_template(
425                "{span_child_prefix}{spinner} {span_name}{{{span_fields}}}",
426            )
427            .expect("valid template"),
428            span_child_prefix_indent: "  ",
429            span_child_prefix_symbol: "↳ ",
430            get_context: WithContext(Self::get_context),
431            get_stderr_writer_context: WithStderrWriter(Self::get_stderr_writer_context),
432            get_stdout_writer_context: WithStdoutWriter(Self::get_stdout_writer_context),
433            get_multi_progress_context: WithMultiProgress(Self::get_multi_progress_context),
434            inner: PhantomData,
435        }
436    }
437}
438
439// pub methods
440impl<S, F> IndicatifLayer<S, F>
441where
442    S: Subscriber + for<'a> LookupSpan<'a>,
443{
444    #[deprecated(since = "0.2.3", note = "use get_stderr_writer() instead")]
445    pub fn get_fmt_writer(&self) -> IndicatifWriter<writer::Stderr> {
446        self.get_stderr_writer()
447    }
448
449    /// Returns the a writer for [`std::io::Stderr`] that ensures its output will not be clobbered by
450    /// active progress bars.
451    ///
452    /// Instead of `eprintln!(...)` prefer `writeln!(indicatif_layer.get_stderr_writer(), ...)`
453    /// instead to ensure your output is not clobbered by active progress bars.
454    ///
455    /// If one wishes tracing logs to be output to stderr, this should be passed into
456    /// [`fmt::Layer::with_writer`](tracing_subscriber::fmt::Layer::with_writer).
457    pub fn get_stderr_writer(&self) -> IndicatifWriter<writer::Stderr> {
458        // `MultiProgress` is merely a wrapper over an `Arc`, so we can clone here.
459        IndicatifWriter::new(self.mp.clone())
460    }
461
462    /// Returns the a writer for [`std::io::Stdout`] that ensures its output will not be clobbered by
463    /// active progress bars.
464    ///
465    /// Instead of `println!(...)` prefer `writeln!(indicatif_layer.get_stdout_writer(), ...)`
466    /// instead to ensure your output is not clobbered by active progress bars.
467    ///
468    /// If one wishes tracing logs to be output to stdout, this should be passed into
469    /// [`fmt::Layer::with_writer`](tracing_subscriber::fmt::Layer::with_writer).
470    pub fn get_stdout_writer(&self) -> IndicatifWriter<writer::Stdout> {
471        // `MultiProgress` is merely a wrapper over an `Arc`, so we can clone here.
472        IndicatifWriter::new(self.mp.clone())
473    }
474
475    /// Set the formatter for span fields, the result of which will be available as the
476    /// progress bar template key `span_fields`.
477    ///
478    /// The default is the [`DefaultFields`] formatter.
479    pub fn with_span_field_formatter<F2>(self, formatter: F2) -> IndicatifLayer<S, F2>
480    where
481        F2: for<'writer> FormatFields<'writer> + 'static,
482    {
483        IndicatifLayer {
484            pb_manager: self.pb_manager,
485            mp: self.mp,
486            span_field_formatter: formatter,
487            progress_style: self.progress_style,
488            span_child_prefix_indent: self.span_child_prefix_indent,
489            span_child_prefix_symbol: self.span_child_prefix_symbol,
490            get_context: WithContext(IndicatifLayer::<S, F2>::get_context),
491            get_stderr_writer_context: WithStderrWriter(
492                IndicatifLayer::<S, F2>::get_stderr_writer_context,
493            ),
494            get_stdout_writer_context: WithStdoutWriter(
495                IndicatifLayer::<S, F2>::get_stdout_writer_context,
496            ),
497            get_multi_progress_context: WithMultiProgress(
498                IndicatifLayer::<S, F2>::get_multi_progress_context,
499            ),
500            inner: self.inner,
501        }
502    }
503
504    /// Override the style used for displayed progress bars.
505    ///
506    /// Two additional keys are available for the progress bar template:
507    /// * `span_fields` - the formatted string of this span's fields
508    /// * `span_name` - the name of the span
509    /// * `span_child_prefix` - a prefix that increase in size according to the number of parents
510    ///   the span has.
511    ///
512    /// The default template is `{span_child_prefix}{spinner} {span_name}{{{span_fields}}}`.
513    pub fn with_progress_style(mut self, style: ProgressStyle) -> Self {
514        self.progress_style = style;
515        self
516    }
517
518    /// Set the indent used to mark the "level" of a given child span's progress bar.
519    ///
520    /// For example, if the given span is two levels deep (iow has two parent spans with progress
521    /// bars), and this is " ", the `{span_child_prefix}` key for this span's progress bar will be
522    /// prefixed with "  ".
523    pub fn with_span_child_prefix_indent(mut self, indent: &'static str) -> Self {
524        self.span_child_prefix_indent = indent;
525        self
526    }
527
528    /// Set the symbol used to denote this is a progress bar from a child span.
529    ///
530    /// This is ultimately concatenated with the child prefix indent to make the
531    /// `span_child_prefix` progress bar key.
532    pub fn with_span_child_prefix_symbol(mut self, symbol: &'static str) -> Self {
533        self.span_child_prefix_symbol = symbol;
534        self
535    }
536
537    /// Set the maximum number of progress bars that will be displayed, and the possible footer
538    /// "progress bar" that displays when there are more progress bars than can be displayed.
539    ///
540    /// `footer_style` dictates the appearance of the footer, and the footer will only appear if
541    /// there are more progress bars than can be displayed. If it is `None`, no footer will be
542    /// displayed. `footer_style` has the following keys available to it:
543    /// * `pending_progress_bars` - the number of progress bars waiting to be shown
544    pub fn with_max_progress_bars(
545        mut self,
546        max_progress_bars: u64,
547        footer_style: Option<ProgressStyle>,
548    ) -> Self {
549        if let Ok(pb_manager) = self.pb_manager.get_mut() {
550            pb_manager.set_max_progress_bars(max_progress_bars, footer_style);
551        }
552
553        self
554    }
555
556    /// Configures how often progress bars are recalcuated and redrawn to the terminal.
557    pub fn with_tick_settings(mut self, tick_settings: TickSettings) -> Self {
558        if let Ok(pb_manager) = self.pb_manager.get_mut() {
559            pb_manager.set_tick_settings(tick_settings);
560        }
561
562        self
563    }
564}
565
566impl<S, F> IndicatifLayer<S, F>
567where
568    S: Subscriber + for<'a> LookupSpan<'a>,
569    F: for<'writer> FormatFields<'writer> + 'static,
570{
571    fn get_context(
572        dispatch: &tracing::Dispatch,
573        id: &span::Id,
574        f: &mut dyn FnMut(&mut IndicatifSpanContext),
575    ) {
576        // The only way `get_context` can be called is if we have an `IndicatifLayer` added to the
577        // expected subscriber, hence why we can `.expect` here.
578        let subscriber = dispatch
579            .downcast_ref::<S>()
580            .expect("subscriber should downcast to expected type; this is a bug!");
581        let span = subscriber
582            .span(id)
583            .expect("Span not found in context, this is a bug");
584
585        let mut ext = span.extensions_mut();
586
587        if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
588            f(indicatif_ctx);
589        }
590    }
591
592    fn get_stderr_writer_context(
593        dispatch: &tracing::Dispatch,
594        f: &mut dyn FnMut(IndicatifWriter<writer::Stderr>),
595    ) {
596        let layer = dispatch
597            .downcast_ref::<IndicatifLayer<S, F>>()
598            .expect("subscriber should downcast to expected type; this is a bug!");
599
600        f(layer.get_stderr_writer())
601    }
602
603    fn get_stdout_writer_context(
604        dispatch: &tracing::Dispatch,
605        f: &mut dyn FnMut(IndicatifWriter<writer::Stdout>),
606    ) {
607        let layer = dispatch
608            .downcast_ref::<IndicatifLayer<S, F>>()
609            .expect("subscriber should downcast to expected type; this is a bug!");
610
611        f(layer.get_stdout_writer())
612    }
613
614    fn get_multi_progress_context(dispatch: &tracing::Dispatch, f: &mut dyn FnMut(MultiProgress)) {
615        let layer = dispatch
616            .downcast_ref::<IndicatifLayer<S, F>>()
617            .expect("subscriber should downcast to expected type; this is a bug!");
618
619        f(layer.mp.clone())
620    }
621
622    fn handle_on_enter(
623        &self,
624        pb_manager: &mut ProgressBarManager,
625        id: &span::Id,
626        ctx: &layer::Context<'_, S>,
627    ) -> Option<ProgressBar> {
628        let span = ctx
629            .span(id)
630            .expect("Span not found in context, this is a bug");
631        let mut ext = span.extensions_mut();
632
633        if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
634            // Start the progress bar when we enter the span for the first time.
635            if indicatif_ctx.progress_bar.is_none() {
636                indicatif_ctx.make_progress_bar(&self.progress_style);
637
638                if let Some(ref parent_span_with_pb) = indicatif_ctx.parent_span {
639                    // Recursively start parent PBs if parent spans have not been entered yet.
640                    let parent_pb = self.handle_on_enter(pb_manager, parent_span_with_pb, ctx);
641
642                    indicatif_ctx.parent_progress_bar = parent_pb;
643                }
644
645                pb_manager.show_progress_bar(indicatif_ctx, id);
646            }
647
648            return indicatif_ctx.progress_bar.to_owned();
649        }
650
651        None
652    }
653}
654
655impl<S, F> layer::Layer<S> for IndicatifLayer<S, F>
656where
657    S: Subscriber + for<'a> LookupSpan<'a>,
658    F: for<'writer> FormatFields<'writer> + 'static,
659{
660    fn on_new_span(&self, attrs: &span::Attributes<'_>, id: &span::Id, ctx: layer::Context<'_, S>) {
661        let span = ctx
662            .span(id)
663            .expect("Span not found in context, this is a bug");
664        let mut ext = span.extensions_mut();
665
666        let mut fields = FormattedFields::<F>::new(String::new());
667        let _ = self
668            .span_field_formatter
669            .format_fields(fields.as_writer(), attrs);
670
671        // Get the next parent span with a progress bar.
672        let parent_span = ctx.span_scope(id).and_then(|scope| {
673            scope.skip(1).find(|span| {
674                let ext = span.extensions();
675
676                ext.get::<IndicatifSpanContext>().is_some()
677            })
678        });
679        let parent_span_id = parent_span.as_ref().map(|span| span.id());
680        let parent_span_ext = parent_span.as_ref().map(|span| span.extensions());
681        let parent_indicatif_ctx = parent_span_ext.as_ref().map(|ext| {
682            ext.get::<IndicatifSpanContext>()
683                .expect("validated it exists prior")
684        });
685
686        let (span_child_prefix, level) = match parent_indicatif_ctx {
687            Some(v) => {
688                let level = v.level + 1;
689
690                (
691                    format!(
692                        "{}{}",
693                        self.span_child_prefix_indent.repeat(level.into()),
694                        self.span_child_prefix_symbol
695                    ),
696                    level,
697                )
698            }
699            None => (String::new(), 0),
700        };
701
702        ext.insert(IndicatifSpanContext {
703            progress_bar: None,
704            pb_init_settings: ProgressBarInitSettings::default(),
705            parent_progress_bar: None,
706            parent_span: parent_span_id,
707            span_fields_formatted: Some(fields.fields),
708            span_name: span.name().to_string(),
709            span_child_prefix,
710            level,
711            finish_message: None,
712        });
713    }
714
715    fn on_enter(&self, id: &span::Id, ctx: layer::Context<'_, S>) {
716        if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
717            self.handle_on_enter(&mut pb_manager_lock, id, &ctx);
718        }
719    }
720
721    fn on_close(&self, id: span::Id, ctx: layer::Context<'_, S>) {
722        if let Ok(mut pb_manager_lock) = self.pb_manager.lock() {
723            let span = ctx
724                .span(&id)
725                .expect("Span not found in context, this is a bug");
726            let mut ext = span.extensions_mut();
727
728            // Clear the progress bar only when the span has closed completely.
729            if let Some(indicatif_ctx) = ext.get_mut::<IndicatifSpanContext>() {
730                pb_manager_lock.finish_progress_bar(indicatif_ctx, &ctx);
731            }
732        }
733    }
734
735    // See comments on [WithContext] for why we have this.
736    //
737    // SAFETY: this is safe because the `WithContext` function pointer is valid
738    // for the lifetime of `&self`.
739    unsafe fn downcast_raw(&self, id: TypeId) -> Option<*const ()> {
740        match id {
741            id if id == TypeId::of::<Self>() => Some(self as *const _ as *const ()),
742            id if id == TypeId::of::<WithContext>() => {
743                Some(&self.get_context as *const _ as *const ())
744            }
745            id if id == TypeId::of::<WithStderrWriter>() => {
746                Some(&self.get_stderr_writer_context as *const _ as *const ())
747            }
748            id if id == TypeId::of::<WithStdoutWriter>() => {
749                Some(&self.get_stdout_writer_context as *const _ as *const ())
750            }
751            id if id == TypeId::of::<WithMultiProgress>() => {
752                Some(&self.get_multi_progress_context as *const _ as *const ())
753            }
754            _ => None,
755        }
756    }
757}
758
759/// Hide all progress bars managed by [`IndicatifLayer`] (if it exists), executes `f`, then redraws
760/// the progress bars. Identical to [`indicatif::MultiProgress::suspend`].
761///
762/// Executes `f` even if there is no default tracing subscriber or if a `IndicatifLayer` has not
763/// been registered to that subscriber.
764///
765/// NOTE: this does not suspend stdout/stderr prints from other threads, including things like
766/// `tracing::info!`. This only suspends the drawing of progress bars.
767///
768/// WARNING: this holds an internal lock within `MultiProgress`. Calling methods like
769/// `writeln!(get_indicatif_stderr_writer(), "foobar")` or calling this method inside of `f` will
770/// result in a deadlock.
771pub fn suspend_tracing_indicatif<F: FnOnce() -> R, R>(f: F) -> R {
772    let mut mp: Option<MultiProgress> = None;
773
774    tracing::dispatcher::get_default(|dispatch| {
775        if let Some(ctx) = dispatch.downcast_ref::<WithMultiProgress>() {
776            ctx.with_context(dispatch, |fetched_mp| {
777                mp = Some(fetched_mp);
778            })
779        }
780    });
781
782    if let Some(mp) = mp {
783        mp.suspend(f)
784    } else {
785        f()
786    }
787}
788
789/// Helper macro that allows you to print to stdout without interfering with the progress bars
790/// created by tracing-indicatif.
791///
792/// Args are directly forwarded to `writeln!`. Do not call this macro inside of
793/// `suspend_tracing_indicatif` or you will trigger a deadlock.
794#[macro_export]
795macro_rules! indicatif_println {
796    ($($arg:tt)*) => {
797        {
798            use std::io::Write;
799
800            if let Some(mut writer) = $crate::writer::get_indicatif_stdout_writer() {
801                writeln!(writer, $($arg)*).unwrap();
802            } else {
803                #[allow(clippy::explicit_write)]
804                writeln!(std::io::stdout(), $($arg)*).unwrap();
805            }
806        }
807    }
808}
809
810/// Helper macro that allows you to print to stderr without interfering with the progress bars
811/// created by tracing-indicatif.
812///
813/// Args are directly forwarded to `writeln!`. Do not call this macro inside of
814/// `suspend_tracing_indicatif` or you will trigger a deadlock.
815#[macro_export]
816macro_rules! indicatif_eprintln {
817    ($($arg:tt)*) => {
818        {
819            use std::io::Write;
820
821            if let Some(mut writer) = $crate::writer::get_indicatif_stderr_writer() {
822                writeln!(writer, $($arg)*).unwrap();
823            } else {
824                #[allow(clippy::explicit_write)]
825                writeln!(std::io::stderr(), $($arg)*).unwrap();
826            }
827        }
828    }
829}
830
831#[cfg(test)]
832mod tests;