Skip to main content

grit_lib/
progress.rs

1//! Progress reporting and cancellation across the library/CLI boundary.
2//!
3//! Library operations never touch the terminal. Long-running work reports
4//! progress through a [`ProgressSink`] supplied by the caller and checks a
5//! [`Cancel`] signal at loop boundaries. The CLI provides a concrete sink that
6//! draws to stderr (gated on `isatty`) and decides colour, redraw rate, and
7//! whether a pager is attached; the library makes none of those decisions.
8//!
9//! Sinks are invoked from the **calling thread only**, at **coarse** granularity
10//! (phase boundaries and periodic counts, not per-object spam). Totals passed
11//! to [`ProgressSink::start`] may be [`None`] when the operation cannot know
12//! the final count up front.
13//!
14//! Tests, `grit-cli`, and any call site that does not want output use the
15//! no-op [`NullProgress`] / [`NeverCancel`].
16//!
17//! Library guide: [progress](https://grit-scm.com/docs/library/parallelism)
18//!
19//! ```rust
20//! use grit_lib::progress::{NeverCancel, NullProgress};
21//! # fn main() {
22//! let _sink = NullProgress;
23//! let _cancel = NeverCancel;
24//! # }
25//! ```
26
27mod estimate;
28
29pub use estimate::{OverallEstimate, OverallEstimator, PhaseRates};
30
31/// A coarse phase of a long-running Git operation.
32///
33/// New variants may be added over time; match exhaustively only in the CLI or
34/// with a wildcard arm when forwarding unknown phases.
35#[non_exhaustive]
36#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
37pub enum ProgressPhase {
38    /// Receiving a pack or loose objects from a remote or bundle stream.
39    ReceivingObjects,
40    /// Walking reachability to decide which objects a local fetch must copy.
41    EnumeratingObjects,
42    /// Building or reading a pack index after objects are on disk.
43    IndexingObjects,
44    /// Applying delta chains inside a pack.
45    ResolvingDeltas,
46    /// Copying objects between local repositories (e.g. local fetch).
47    CopyingObjects,
48    /// Updating the work tree to match a tree (checkout, reset --hard, clone).
49    CheckingOutFiles,
50    /// Hashing work-tree files into the index (add, auto-staging commit, stash push).
51    StagingFiles,
52    /// Downloading or uploading a bundle that can resume where it stopped
53    /// (clone's `bundle-uri` bundle, a large push). Counted in bytes, including
54    /// any bytes an earlier run already moved; the detailed events go to
55    /// [`crate::fetch::Progress::resumable_transfer`].
56    TransferringBundle,
57    /// Streaming a pack to a remote during push (counted in bytes).
58    SendingObjects,
59}
60
61impl ProgressPhase {
62    /// Returns the unit callers should use when reporting counts for `self`.
63    #[must_use]
64    pub fn default_unit(self) -> ProgressUnit {
65        match self {
66            Self::ReceivingObjects
67            | Self::EnumeratingObjects
68            | Self::IndexingObjects
69            | Self::ResolvingDeltas
70            | Self::CopyingObjects => ProgressUnit::Objects,
71            Self::CheckingOutFiles | Self::StagingFiles => ProgressUnit::Files,
72            Self::TransferringBundle | Self::SendingObjects => ProgressUnit::Bytes,
73        }
74    }
75}
76
77/// Unit for [`ProgressSink::inc`] / [`ProgressSink::set`] during a phase.
78#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
79pub enum ProgressUnit {
80    /// Git objects (commits, trees, blobs, tags).
81    Objects,
82    /// Raw bytes (e.g. pack stream length when known).
83    Bytes,
84    /// Work-tree paths.
85    Files,
86}
87
88/// How far a pack transfer has got, counted in objects rather than bytes.
89///
90/// Reported through [`crate::fetch::Progress::received_objects`] while a pack
91/// streams in, when the sink asks for it with
92/// [`crate::fetch::Progress::wants_received_objects`].
93#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
94pub struct ReceivedObjects {
95    /// Objects fully received so far.
96    pub received: u64,
97    /// Objects in the pack, from its header.
98    pub total: u64,
99    /// How many of the received objects are deltas.
100    pub deltas: u64,
101}
102
103/// A sink for progress updates emitted by a library operation.
104///
105/// Every method has a no-op default, so a sink implements only what it needs.
106/// The library calls these from the caller's thread at coarse intervals; it
107/// never decides whether output is a tty, what colour to use, or how often to
108/// redraw — that is the CLI's responsibility.
109pub trait ProgressSink {
110    /// Begin a new progress `phase`, optionally with a known total count of
111    /// [`ProgressPhase::default_unit`] units for that phase.
112    ///
113    /// `total` may be [`None`] when the operation cannot estimate the final
114    /// count (common for network receives).
115    ///
116    fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
117        let _ = (phase, total);
118    }
119
120    /// Begin a phase with an explicit count unit (for example [`ProgressUnit::Bytes`]
121    /// when streaming a pack whose size is known in bytes).
122    ///
123    /// Defaults to [`Self::start`]; override when the sink renders counts differently
124    /// per [`ProgressUnit`].
125    fn start_with_unit(&mut self, phase: ProgressPhase, unit: ProgressUnit, total: Option<u64>) {
126        let _ = unit;
127        self.start(phase, total);
128    }
129
130    /// Advance the current phase by `units` in the active phase unit.
131    fn inc(&mut self, units: u64) {
132        let _ = units;
133    }
134
135    /// Set the absolute progress of the current phase to `current` units.
136    fn set(&mut self, current: u64) {
137        let _ = current;
138    }
139
140    /// Finish the current phase.
141    fn finish(&mut self) {}
142
143    /// When `true`, the operation should stop at the next [`poll_cancel`] check.
144    ///
145    /// CLI progress sinks may set this when the user interrupts (Ctrl+C). Most
146    /// call sites rely on a separate [`Cancel`] argument instead; either signal
147    /// produces [`crate::error::Error::Cancelled`].
148    fn should_cancel(&self) -> bool {
149        false
150    }
151}
152
153/// Forwarding impl so a `&mut P` can be threaded into sub-operations that also
154/// take `impl ProgressSink` without re-borrowing gymnastics at every call site.
155impl<T: ProgressSink + ?Sized> ProgressSink for &mut T {
156    fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
157        (**self).start(phase, total);
158    }
159    fn start_with_unit(&mut self, phase: ProgressPhase, unit: ProgressUnit, total: Option<u64>) {
160        (**self).start_with_unit(phase, unit, total);
161    }
162    fn inc(&mut self, units: u64) {
163        (**self).inc(units);
164    }
165    fn set(&mut self, current: u64) {
166        (**self).set(current);
167    }
168    fn finish(&mut self) {
169        (**self).finish();
170    }
171    fn should_cancel(&self) -> bool {
172        (**self).should_cancel()
173    }
174}
175
176/// A [`ProgressSink`] that discards every update. The default for tests,
177/// `grit-cli`, and any call site that does not want progress output.
178#[derive(Debug, Default, Clone, Copy)]
179pub struct NullProgress;
180
181impl ProgressSink for NullProgress {}
182
183/// A cancellation signal checked by long-running library operations at loop
184/// boundaries. Implementations must be cheap to query.
185pub trait Cancel {
186    /// Returns `true` once the operation should stop as soon as possible.
187    fn is_cancelled(&self) -> bool;
188
189    /// When cancellation is backed by a shared flag (for example [`FlagCancel`]), returns it
190    /// so worker threads can observe the same signal.
191    fn shared_watch(&self) -> Option<std::sync::Arc<std::sync::atomic::AtomicBool>> {
192        None
193    }
194
195    /// Mirror a stack-scoped cancel into a shared watch before blocking HTTP waits.
196    fn poll_before_wait(&self) {}
197}
198
199/// A [`Cancel`] that never signals cancellation. The default for call sites
200/// that do not support interruption.
201#[derive(Debug, Default, Clone, Copy)]
202pub struct NeverCancel;
203
204impl Cancel for NeverCancel {
205    fn is_cancelled(&self) -> bool {
206        false
207    }
208}
209
210/// How often long object walks and pack builds report [`ProgressSink::set`].
211pub(crate) const OBJECT_PROGRESS_STRIDE: u64 = 256;
212
213/// How often worktree blob stat scans and similar loops poll [`poll_watch`].
214pub(crate) const CANCEL_POLL_STRIDE: usize = 256;
215
216/// Wraps a [`Write`] and reports bytes written to a [`ProgressSink`] as
217/// [`ProgressPhase::SendingObjects`].
218pub(crate) struct ProgressWrite<'a, W: ?Sized> {
219    inner: &'a mut W,
220    sink: &'a mut dyn ProgressSink,
221    started: bool,
222}
223
224impl<'a, W: ?Sized + std::io::Write> ProgressWrite<'a, W> {
225    /// Create a wrapper that forwards writes to `inner` and counts them on `sink`.
226    pub fn new(inner: &'a mut W, sink: &'a mut dyn ProgressSink) -> Self {
227        Self {
228            inner,
229            sink,
230            started: false,
231        }
232    }
233
234    /// End the sending phase when one was started.
235    pub fn finish(&mut self) {
236        if self.started {
237            self.sink.finish();
238            self.started = false;
239        }
240    }
241}
242
243impl<W: ?Sized + std::io::Write> std::io::Write for ProgressWrite<'_, W> {
244    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
245        let n = self.inner.write(buf)?;
246        if n > 0 {
247            if !self.started {
248                self.sink
249                    .start_with_unit(ProgressPhase::SendingObjects, ProgressUnit::Bytes, None);
250                self.started = true;
251            }
252            self.sink.inc(n as u64);
253        }
254        Ok(n)
255    }
256
257    fn flush(&mut self) -> std::io::Result<()> {
258        self.inner.flush()
259    }
260}
261
262/// Report monotonic object progress at [`OBJECT_PROGRESS_STRIDE`] boundaries.
263///
264/// `done` is the absolute number of objects processed so far; `last_reported` is
265/// updated to match each emitted `set`.
266pub(crate) fn report_object_progress_stride(
267    sink: &mut dyn ProgressSink,
268    done: u64,
269    last_reported: &mut u64,
270    total: Option<u64>,
271) {
272    let stride = OBJECT_PROGRESS_STRIDE;
273    while done.saturating_sub(*last_reported) >= stride {
274        *last_reported += stride;
275        sink.set(*last_reported);
276    }
277    if total == Some(done) && *last_reported < done {
278        sink.set(done);
279        *last_reported = done;
280    }
281}
282
283/// Shared flag for tests and embedders that flip cancellation mid-operation.
284#[derive(Debug, Clone)]
285pub struct FlagCancel(pub std::sync::Arc<std::sync::atomic::AtomicBool>);
286
287impl FlagCancel {
288    /// Create a cancel signal that starts unset.
289    #[must_use]
290    pub fn new() -> Self {
291        Self(std::sync::Arc::new(std::sync::atomic::AtomicBool::new(
292            false,
293        )))
294    }
295
296    /// Signal cancellation to every consumer polling this handle.
297    pub fn cancel(&self) {
298        self.0.store(true, std::sync::atomic::Ordering::Relaxed);
299    }
300}
301
302impl Default for FlagCancel {
303    fn default() -> Self {
304        Self::new()
305    }
306}
307
308impl Cancel for FlagCancel {
309    fn is_cancelled(&self) -> bool {
310        self.0.load(std::sync::atomic::Ordering::Relaxed)
311    }
312
313    fn shared_watch(&self) -> Option<std::sync::Arc<std::sync::atomic::AtomicBool>> {
314        Some(std::sync::Arc::clone(&self.0))
315    }
316}
317
318/// Copy a stack [`Cancel`] into `watch` for worker threads (call on the owning thread).
319#[inline]
320pub(crate) fn mirror_stack_cancel(cancel: &dyn Cancel, watch: &std::sync::atomic::AtomicBool) {
321    cancel.poll_before_wait();
322    if cancel.is_cancelled() {
323        watch.store(true, std::sync::atomic::Ordering::Relaxed);
324    }
325}
326
327/// Shared flag for smart-HTTP waits; the owning thread mirrors a stack [`Cancel`] into it.
328pub(crate) struct HttpStackWatch(pub(crate) std::sync::Arc<std::sync::atomic::AtomicBool>);
329
330impl HttpStackWatch {
331    #[must_use]
332    pub fn new(cancel: &dyn Cancel) -> Self {
333        Self(cancel.shared_watch().unwrap_or_else(|| {
334            std::sync::Arc::new(std::sync::atomic::AtomicBool::new(cancel.is_cancelled()))
335        }))
336    }
337}
338
339/// Keeps [`HttpStackWatch`] alive for a fetch/push scope.
340pub(crate) struct HttpWaitGuard {
341    _watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
342}
343
344impl HttpWaitGuard {
345    pub(crate) fn bind(
346        cancel: &dyn Cancel,
347    ) -> (std::sync::Arc<std::sync::atomic::AtomicBool>, Self) {
348        let watch = HttpStackWatch::new(cancel).0;
349        (std::sync::Arc::clone(&watch), Self { _watch: watch })
350    }
351}
352
353/// Share a [`Cancel`] handle with smart-HTTP connect options for worker threads.
354///
355/// When [`Cancel::shared_watch`] is set, workers observe the same flag as the caller.
356/// Otherwise `stack_watch` must be the fetch/push scope flag; the owning thread mirrors
357/// the stack [`Cancel`] into it during HTTP waits (see [`mirror_stack_cancel`]).
358#[must_use]
359pub(crate) fn connect_cancel_arc(
360    cancel: &dyn Cancel,
361    stack_watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
362) -> std::sync::Arc<dyn Cancel + Send + Sync> {
363    mirror_stack_cancel(cancel, &stack_watch);
364    if let Some(flag) = cancel.shared_watch() {
365        return std::sync::Arc::new(FlagCancel(flag));
366    }
367    std::sync::Arc::new(ConnectOperationCancel { watch: stack_watch })
368}
369
370/// Shared watch for HTTP worker threads when the caller's [`Cancel`] has no [`Cancel::shared_watch`].
371struct ConnectOperationCancel {
372    watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
373}
374
375impl Cancel for ConnectOperationCancel {
376    fn is_cancelled(&self) -> bool {
377        self.watch.load(std::sync::atomic::Ordering::Relaxed)
378    }
379
380    fn shared_watch(&self) -> Option<std::sync::Arc<std::sync::atomic::AtomicBool>> {
381        Some(std::sync::Arc::clone(&self.watch))
382    }
383}
384
385/// Shared [`AtomicBool`](std::sync::atomic::AtomicBool) for cancellation visible across threads.
386///
387/// For [`FlagCancel`], this is the same flag as the handle (so embedders can cancel mid-scan).
388/// For other [`Cancel`] implementations, a private flag is seeded from [`Cancel::is_cancelled`]
389/// at creation time; combine with [`poll_cancel_watch`] on the main thread for progress signals.
390#[must_use]
391pub fn cancel_watch_arc(cancel: &dyn Cancel) -> std::sync::Arc<std::sync::atomic::AtomicBool> {
392    cancel.shared_watch().unwrap_or_else(|| {
393        std::sync::Arc::new(std::sync::atomic::AtomicBool::new(cancel.is_cancelled()))
394    })
395}
396
397/// Returns [`crate::error::Error::Cancelled`] when `watch` was set.
398#[inline]
399pub fn poll_watch(watch: &std::sync::atomic::AtomicBool) -> crate::error::Result<()> {
400    if watch.load(std::sync::atomic::Ordering::Relaxed) {
401        Err(crate::error::Error::Cancelled)
402    } else {
403        Ok(())
404    }
405}
406
407/// Poll hook for long [`crate::diff::diff_index_to_worktree_with_options`] scans on the stack
408/// frame that owns `progress`.
409///
410/// [`Self::poll_owner`] must run only on the thread that owns the captured `WatchProgress`.
411pub(crate) struct StackCancelPoll {
412    progress: usize,
413    watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
414}
415
416impl StackCancelPoll {
417    /// Capture `progress` and `watch` for [`Self::poll_owner`].
418    #[must_use]
419    pub fn new(
420        progress: &mut WatchProgress<'_>,
421        watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
422    ) -> Self {
423        Self {
424            progress: std::ptr::from_mut(progress).addr(),
425            watch,
426        }
427    }
428
429    /// Mirror [`ProgressSink::should_cancel`] into `watch`, then [`poll_watch`].
430    pub(crate) fn poll_owner(&self) -> crate::error::Result<()> {
431        // SAFETY: only invoked on the owning thread while the captured `WatchProgress` is live.
432        let progress = unsafe { &mut *(self.progress as *mut WatchProgress<'_>) };
433        if progress.should_cancel() {
434            self.watch.store(true, std::sync::atomic::Ordering::Relaxed);
435            return Err(crate::error::Error::Cancelled);
436        }
437        poll_watch(&self.watch)
438    }
439}
440
441/// Cancellation inputs for worktree stat scans: worker-visible `watch` plus an owner-thread poll hook.
442pub(crate) struct StatScanCancel {
443    pub watch: Option<std::sync::Arc<std::sync::atomic::AtomicBool>>,
444    pub owner_poll: Option<std::sync::Arc<StackCancelPoll>>,
445}
446
447impl StatScanCancel {
448    /// Poll owner progress (when present), otherwise the shared watch flag.
449    pub(crate) fn poll(&self) -> crate::error::Result<()> {
450        if let Some(poll) = self.owner_poll.as_ref() {
451            poll.poll_owner()
452        } else if let Some(w) = self.watch.as_ref() {
453            poll_watch(w)
454        } else {
455            Ok(())
456        }
457    }
458}
459
460/// Like [`poll_cancel`], but also records cancellation in `watch` for worker threads.
461#[inline]
462pub fn poll_cancel_watch(
463    progress: &mut dyn ProgressSink,
464    cancel: &dyn Cancel,
465    watch: &std::sync::atomic::AtomicBool,
466) -> crate::error::Result<()> {
467    if cancel.is_cancelled() || progress.should_cancel() {
468        watch.store(true, std::sync::atomic::Ordering::Relaxed);
469        return Err(crate::error::Error::Cancelled);
470    }
471    poll_watch(watch)
472}
473
474/// A [`Cancel`] polled on the calling thread, plus the flag its worker threads read.
475///
476/// [`Self::poll`] records a cancel in [`Self::flag`], so workers also stop for
477/// [`Cancel`] implementations that are not backed by a shared flag (once the calling
478/// thread next polls). For a [`FlagCancel`] the flag is the handle's own, so workers
479/// see a cancel from another thread right away.
480pub(crate) struct CancelSignal<'a> {
481    cancel: &'a dyn Cancel,
482    flag: std::sync::Arc<std::sync::atomic::AtomicBool>,
483}
484
485impl<'a> CancelSignal<'a> {
486    pub fn new(cancel: &'a dyn Cancel) -> Self {
487        Self {
488            cancel,
489            flag: cancel_watch_arc(cancel),
490        }
491    }
492
493    /// A signal that never fires, for callers that cannot be interrupted.
494    pub(crate) fn never() -> CancelSignal<'static> {
495        CancelSignal::new(&NeverCancel)
496    }
497
498    /// Returns [`crate::error::Error::Cancelled`] once the signal fired.
499    #[inline]
500    pub(crate) fn poll(&self) -> crate::error::Result<()> {
501        self.cancel.poll_before_wait();
502        if self.cancel.is_cancelled() {
503            self.flag.store(true, std::sync::atomic::Ordering::Relaxed);
504        }
505        poll_watch(&self.flag)
506    }
507
508    /// The flag worker threads check; set once [`Self::poll`] observed a cancel.
509    pub(crate) fn flag(&self) -> &std::sync::atomic::AtomicBool {
510        &self.flag
511    }
512
513    /// Shared flag for long walks on other threads or modules.
514    pub(crate) fn watch_arc(&self) -> std::sync::Arc<std::sync::atomic::AtomicBool> {
515        std::sync::Arc::clone(&self.flag)
516    }
517}
518
519/// [`ProgressSink`] adapter that mirrors [`ProgressSink::should_cancel`] into `watch`.
520pub struct WatchProgress<'a> {
521    inner: &'a mut dyn ProgressSink,
522    watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
523}
524
525impl<'a> WatchProgress<'a> {
526    /// Wrap `inner` so any cancel request on the sink sets `watch`.
527    #[must_use]
528    pub fn new(
529        inner: &'a mut dyn ProgressSink,
530        watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
531    ) -> Self {
532        Self { inner, watch }
533    }
534
535    fn note_cancel(&self) {
536        if self.inner.should_cancel() {
537            self.watch.store(true, std::sync::atomic::Ordering::Relaxed);
538        }
539    }
540}
541
542impl ProgressSink for WatchProgress<'_> {
543    fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
544        self.note_cancel();
545        self.inner.start(phase, total);
546    }
547    fn start_with_unit(&mut self, phase: ProgressPhase, unit: ProgressUnit, total: Option<u64>) {
548        self.note_cancel();
549        self.inner.start_with_unit(phase, unit, total);
550    }
551    fn inc(&mut self, units: u64) {
552        self.note_cancel();
553        self.inner.inc(units);
554    }
555    fn set(&mut self, current: u64) {
556        self.note_cancel();
557        self.inner.set(current);
558    }
559    fn finish(&mut self) {
560        self.note_cancel();
561        self.inner.finish();
562    }
563    fn should_cancel(&self) -> bool {
564        self.inner.should_cancel()
565    }
566}
567
568/// The progress sink, cancel signal and shared watch flag one operation threads through its
569/// helpers, so worker threads and the calling thread stop on the same signal.
570pub(crate) struct ProgressScope<'a> {
571    sink: WatchProgress<'a>,
572    cancel: &'a dyn Cancel,
573    watch: std::sync::Arc<std::sync::atomic::AtomicBool>,
574}
575
576impl<'a> ProgressScope<'a> {
577    pub fn new(sink: &'a mut dyn ProgressSink, cancel: &'a dyn Cancel) -> Self {
578        let watch = cancel_watch_arc(cancel);
579        Self {
580            sink: WatchProgress::new(sink, std::sync::Arc::clone(&watch)),
581            cancel,
582            watch,
583        }
584    }
585
586    /// Returns [`crate::error::Error::Cancelled`] once either signal fired.
587    pub(crate) fn poll(&mut self) -> crate::error::Result<()> {
588        poll_cancel_watch(&mut self.sink, self.cancel, &self.watch)
589    }
590
591    /// Flag for worker threads; set once [`Self::poll`] observed a cancel.
592    pub(crate) fn watch(&self) -> &std::sync::atomic::AtomicBool {
593        &self.watch
594    }
595
596    pub(crate) fn sink(&mut self) -> &mut dyn ProgressSink {
597        &mut self.sink
598    }
599}
600
601/// [`ProgressSink`] adapter that forwards every update but never asks to stop, for steps that
602/// must run to completion once an operation has published its state.
603pub(crate) struct UncancellableProgress<'a> {
604    inner: &'a mut dyn ProgressSink,
605}
606
607impl<'a> UncancellableProgress<'a> {
608    pub fn new(inner: &'a mut dyn ProgressSink) -> Self {
609        Self { inner }
610    }
611}
612
613impl ProgressSink for UncancellableProgress<'_> {
614    fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
615        self.inner.start(phase, total);
616    }
617    fn start_with_unit(&mut self, phase: ProgressPhase, unit: ProgressUnit, total: Option<u64>) {
618        self.inner.start_with_unit(phase, unit, total);
619    }
620    fn inc(&mut self, units: u64) {
621        self.inner.inc(units);
622    }
623    fn set(&mut self, current: u64) {
624        self.inner.set(current);
625    }
626    fn finish(&mut self) {
627        self.inner.finish();
628    }
629}
630
631/// Returns [`crate::error::Error::Cancelled`] when `cancel` or `progress` request stop.
632#[inline]
633pub fn poll_cancel(
634    progress: &mut dyn ProgressSink,
635    cancel: &dyn Cancel,
636) -> crate::error::Result<()> {
637    cancel.poll_before_wait();
638    if cancel.is_cancelled() || progress.should_cancel() {
639        Err(crate::error::Error::Cancelled)
640    } else {
641        Ok(())
642    }
643}
644
645/// [`ProgressSink`] that stops forwarding once [`Cancel`] or [`ProgressSink::should_cancel`] fires.
646///
647/// Long checkout loops call [`Self::finish_poll`] after apply so callers get [`crate::error::Error::Cancelled`]
648/// even when cancellation happened on an [`ProgressSink::inc`] boundary.
649pub struct CancelProgress<'a> {
650    inner: &'a mut dyn ProgressSink,
651    cancel: &'a dyn Cancel,
652    stopped: std::cell::Cell<bool>,
653}
654
655impl<'a> CancelProgress<'a> {
656    /// Wrap `inner` so every progress update observes `cancel`.
657    pub fn new(inner: &'a mut dyn ProgressSink, cancel: &'a dyn Cancel) -> Self {
658        Self {
659            inner,
660            cancel,
661            stopped: std::cell::Cell::new(false),
662        }
663    }
664
665    /// Whether a cancel signal was observed.
666    #[must_use]
667    pub fn stopped(&self) -> bool {
668        self.stopped.get()
669    }
670
671    /// Returns [`crate::error::Error::Cancelled`] when [`Self::stopped`].
672    pub fn finish_poll(&self) -> crate::error::Result<()> {
673        if self.stopped() {
674            Err(crate::error::Error::Cancelled)
675        } else {
676            Ok(())
677        }
678    }
679
680    fn note_stop(&self) {
681        if self.cancel.is_cancelled() || self.inner.should_cancel() {
682            self.stopped.set(true);
683        }
684    }
685}
686
687impl ProgressSink for CancelProgress<'_> {
688    fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
689        self.note_stop();
690        if !self.stopped.get() {
691            self.inner.start(phase, total);
692        }
693    }
694
695    fn start_with_unit(&mut self, phase: ProgressPhase, unit: ProgressUnit, total: Option<u64>) {
696        self.note_stop();
697        if !self.stopped.get() {
698            self.inner.start_with_unit(phase, unit, total);
699        }
700    }
701
702    fn inc(&mut self, units: u64) {
703        self.note_stop();
704        if !self.stopped.get() {
705            self.inner.inc(units);
706        }
707    }
708
709    fn set(&mut self, current: u64) {
710        self.note_stop();
711        if !self.stopped.get() {
712            self.inner.set(current);
713        }
714    }
715
716    fn finish(&mut self) {
717        self.note_stop();
718        if !self.stopped.get() {
719            self.inner.finish();
720        }
721    }
722
723    fn should_cancel(&self) -> bool {
724        self.stopped.get() || self.cancel.is_cancelled() || self.inner.should_cancel()
725    }
726}
727
728#[cfg(test)]
729mod tests {
730    use super::*;
731    use crate::fetch::{NoProgress, Progress};
732
733    /// A sink that records typed phase updates, to prove the trait + forwarding
734    /// impl route calls correctly.
735    #[derive(Default)]
736    struct PhaseRecorder {
737        started: Vec<(ProgressPhase, Option<u64>)>,
738        incs: u64,
739        finished: usize,
740    }
741
742    impl ProgressSink for PhaseRecorder {
743        fn start(&mut self, phase: ProgressPhase, total: Option<u64>) {
744            self.started.push((phase, total));
745        }
746        fn inc(&mut self, units: u64) {
747            self.incs += units;
748        }
749        fn finish(&mut self) {
750            self.finished += 1;
751        }
752    }
753
754    fn drive_typed(sink: &mut impl ProgressSink) {
755        sink.start(ProgressPhase::ReceivingObjects, Some(3));
756        sink.inc(1);
757        sink.inc(2);
758        sink.finish();
759    }
760
761    #[test]
762    fn null_progress_is_inert() {
763        drive_typed(&mut NullProgress);
764    }
765
766    #[test]
767    fn no_progress_is_inert() {
768        drive_typed(&mut NoProgress);
769    }
770
771    #[test]
772    fn typed_phase_recorder_roundtrip() {
773        let mut rec = PhaseRecorder::default();
774        drive_typed(&mut &mut rec);
775        assert_eq!(
776            rec.started,
777            vec![(ProgressPhase::ReceivingObjects, Some(3))]
778        );
779        assert_eq!(rec.incs, 3);
780        assert_eq!(rec.finished, 1);
781    }
782
783    #[test]
784    fn progress_supertrait_upcasts() {
785        struct SidebandCap(Vec<u8>);
786
787        impl ProgressSink for SidebandCap {}
788
789        impl Progress for SidebandCap {
790            fn message(&mut self, bytes: &[u8]) {
791                self.0.extend_from_slice(bytes);
792            }
793        }
794
795        fn takes_sink(sink: &mut dyn ProgressSink) {
796            sink.start(ProgressPhase::IndexingObjects, None);
797            sink.finish();
798        }
799
800        let mut cap = SidebandCap(Vec::new());
801        let progress: &mut dyn Progress = &mut cap;
802        takes_sink(progress);
803    }
804
805    #[test]
806    fn never_cancel_never_cancels() {
807        assert!(!NeverCancel.is_cancelled());
808    }
809
810    #[test]
811    fn poll_cancel_returns_typed_error() {
812        let flag = FlagCancel::new();
813        flag.cancel();
814        let mut sink = NullProgress;
815        let err = poll_cancel(&mut sink, &flag).expect_err("cancelled");
816        assert!(matches!(err, crate::error::Error::Cancelled));
817    }
818
819    #[test]
820    fn phase_default_unit() {
821        assert_eq!(
822            ProgressPhase::ReceivingObjects.default_unit(),
823            ProgressUnit::Objects
824        );
825        assert_eq!(
826            ProgressPhase::CheckingOutFiles.default_unit(),
827            ProgressUnit::Files
828        );
829        assert_eq!(
830            ProgressPhase::StagingFiles.default_unit(),
831            ProgressUnit::Files
832        );
833    }
834}