Skip to main content

cli/data/
progress.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Minimal internal progress abstraction for Export/Import V2.
16//!
17//! This is intentionally small and log/internal oriented. It does not touch
18//! stdout and is safe for non-interactive runs. A TTY progress bar (e.g.
19//! `indicatif`) and a `--progress` CLI flag are deliberately out of scope; they
20//! can layer on top of this trait later without changing call sites.
21
22/// Receives progress events from long-running Export/Import V2 work.
23///
24/// The trait is object-safe so callers can take `&dyn ProgressReporter` and stay
25/// agnostic about the concrete implementation (no-op in production today, a
26/// recording fake in tests).
27pub(crate) trait ProgressReporter: Send + Sync {
28    /// Begins a phase with an optional known total number of units.
29    fn start_phase(&self, name: &str, total: Option<u64>);
30
31    /// Advances the current phase by `delta` units.
32    fn inc(&self, delta: u64);
33
34    /// Marks the current phase as finished.
35    fn finish_phase(&self);
36}
37
38/// A reporter that discards every event. Used as the production default and in
39/// tests that do not care about progress.
40pub(crate) struct NoopProgress;
41
42impl ProgressReporter for NoopProgress {
43    fn start_phase(&self, _name: &str, _total: Option<u64>) {}
44    fn inc(&self, _delta: u64) {}
45    fn finish_phase(&self) {}
46}
47
48/// RAII guard for a started progress phase.
49///
50/// This keeps future stateful reporters safe on every early-return path after a
51/// phase starts. Call [`Self::finish`] to end the phase at a deliberate point;
52/// otherwise the guard finishes it when dropped.
53#[must_use = "dropping the guard immediately finishes the phase"]
54pub(crate) struct ProgressPhase<'a> {
55    reporter: &'a dyn ProgressReporter,
56    finished: bool,
57}
58
59impl<'a> ProgressPhase<'a> {
60    pub(crate) fn start(
61        reporter: &'a dyn ProgressReporter,
62        name: &str,
63        total: Option<u64>,
64    ) -> Self {
65        reporter.start_phase(name, total);
66        Self {
67            reporter,
68            finished: false,
69        }
70    }
71
72    pub(crate) fn finish(mut self) {
73        self.finish_once();
74    }
75
76    fn finish_once(&mut self) {
77        if !self.finished {
78            self.reporter.finish_phase();
79            self.finished = true;
80        }
81    }
82}
83
84impl Drop for ProgressPhase<'_> {
85    fn drop(&mut self) {
86        self.finish_once();
87    }
88}
89
90#[cfg(test)]
91pub(crate) mod test_util {
92    use std::sync::Mutex;
93
94    use super::ProgressReporter;
95
96    /// A single recorded progress event, used to assert progress behavior in
97    /// unit tests.
98    #[derive(Debug, Clone, PartialEq, Eq)]
99    pub(crate) enum ProgressEvent {
100        StartPhase { name: String, total: Option<u64> },
101        Inc { delta: u64 },
102        FinishPhase,
103    }
104
105    /// A reporter that records every event in order for later assertions.
106    #[derive(Default)]
107    pub(crate) struct RecordingProgress {
108        events: Mutex<Vec<ProgressEvent>>,
109    }
110
111    impl RecordingProgress {
112        pub(crate) fn events(&self) -> Vec<ProgressEvent> {
113            self.events.lock().unwrap().clone()
114        }
115
116        /// Sum of all `inc` deltas recorded so far.
117        pub(crate) fn total_inc(&self) -> u64 {
118            self.events
119                .lock()
120                .unwrap()
121                .iter()
122                .filter_map(|event| match event {
123                    ProgressEvent::Inc { delta } => Some(*delta),
124                    _ => None,
125                })
126                .sum()
127        }
128
129        fn push(&self, event: ProgressEvent) {
130            self.events.lock().unwrap().push(event);
131        }
132    }
133
134    impl ProgressReporter for RecordingProgress {
135        fn start_phase(&self, name: &str, total: Option<u64>) {
136            self.push(ProgressEvent::StartPhase {
137                name: name.to_string(),
138                total,
139            });
140        }
141
142        fn inc(&self, delta: u64) {
143            self.push(ProgressEvent::Inc { delta });
144        }
145
146        fn finish_phase(&self) {
147            self.push(ProgressEvent::FinishPhase);
148        }
149    }
150}