view rust/hg-core/src/progress.rs @ 52156:e6a44bc91bc2 stable

rust-update: make `update_from_null` respect `worker.numcpu` config option This was overlooked in the original series. This is important for tests (because we run many at once), and for the occasional end user that wants to keep their CPU usage in check. A future series should clean up this `worker` parameter tunelling business by rewriting the config in Rust, but doing so on stable would be a very bad idea.
author Rapha?l Gom?s <rgomes@octobus.net>
date Tue, 05 Nov 2024 15:21:09 +0100
parents 92e23ba257d1
children a876ab6c3fd5
line wrap: on
line source

//! Progress-bar related things

use std::{
    sync::atomic::{AtomicBool, Ordering},
    time::Duration,
};

use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};

/// A generic determinate progress bar trait
pub trait Progress: Send + Sync + 'static {
    /// Set the current position and optionally the total
    fn update(&self, pos: u64, total: Option<u64>);
    /// Increment the current position and optionally the total
    fn increment(&self, step: u64, total: Option<u64>);
    /// Declare that progress is over and the progress bar should be deleted
    fn complete(self);
}

const PROGRESS_DELAY: Duration = Duration::from_secs(1);

/// A generic (determinate) progress bar. Stays hidden until [`PROGRESS_DELAY`]
/// to prevent flickering a progress bar for super fast operations.
pub struct HgProgressBar {
    progress: ProgressBar,
    has_been_shown: AtomicBool,
}

impl HgProgressBar {
    // TODO pass config to check progress.disable/assume-tty/delay/etc.
    /// Return a new progress bar with `topic` as the prefix.
    /// The progress and total are both set to 0, and it is hidden until the
    /// next call to `update` given that more than a second has elapsed.
    pub fn new(topic: &str) -> Self {
        let template =
            format!("{} {{wide_bar}} {{pos}}/{{len}} {{eta}} ", topic);
        let style = ProgressStyle::with_template(&template).unwrap();
        let progress_bar = ProgressBar::new(0).with_style(style);
        // Hide the progress bar and only show it if we've elapsed more
        // than a second
        progress_bar.set_draw_target(ProgressDrawTarget::hidden());
        Self {
            progress: progress_bar,
            has_been_shown: false.into(),
        }
    }

    /// Called whenever the progress changes to determine whether to start
    /// showing the progress bar
    fn maybe_show(&self) {
        if self.progress.is_hidden()
            && self.progress.elapsed() > PROGRESS_DELAY
        {
            // Catch a race condition whereby we check if it's hidden, then
            // set the draw target from another thread, then do it again from
            // this thread, which results in multiple progress bar lines being
            // left drawn.
            let has_been_shown =
                self.has_been_shown.fetch_or(true, Ordering::Relaxed);
            if !has_been_shown {
                // Here we are certain that we're the only thread that has
                // set `has_been_shown` and we can change the draw target
                self.progress.set_draw_target(ProgressDrawTarget::stderr());
                self.progress.tick();
            }
        }
    }
}

impl Progress for HgProgressBar {
    fn update(&self, pos: u64, total: Option<u64>) {
        self.progress.update(|state| {
            state.set_pos(pos);
            if let Some(t) = total {
                state.set_len(t)
            }
        });
        self.maybe_show();
    }

    fn increment(&self, step: u64, total: Option<u64>) {
        self.progress.inc(step);
        if let Some(t) = total {
            self.progress.set_length(t)
        }
        self.maybe_show();
    }

    fn complete(self) {
        self.progress.finish_and_clear();
    }
}