package main import ( "fmt" "os" "slices" "sync" "time" "github.com/schollz/progressbar/v3" "golang.org/x/term" ) // plainInterval is the time between progress lines while a phase runs, // when stderr is not a TTY. const plainInterval = 5 * time.Second // redrawInterval is the time between redraws of the progress display on // a TTY. const redrawInterval = 100 * time.Millisecond // spinnerType selects the progressbar library's spinner style. const spinnerType = 14 // percentScale converts a fraction to a percentage. const percentScale = 100 // stderrIsTTY reports whether stderr is attached to a terminal. func stderrIsTTY() bool { return term.IsTerminal(int(os.Stderr.Fd())) } // phaseLabel names phase with its place among phases, as "[2/5] walk". func phaseLabel(phases []string, phase string) string { return fmt.Sprintf("[%d/%d] %s", slices.Index(phases, phase)+1, len(phases), phase) } // progressParams describe one phase's progress display as it starts. type progressParams struct { label string // the phase and its place, from phaseLabel total int64 // the items the phase will complete; -1 when not known estimate bool // total is an estimate, which the count may pass waitingFor string // what the phase waits for until setTotal is called } // progress renders one phase's progress on stderr as one line (see // line): the phase and its place among the phases, the items done out // of the total, the rate, the elapsed time and the estimated time // remaining. On a TTY the progressbar library draws the line behind a // spinner, again every redrawInterval. When stderr is not a TTY it // emits no ANSI redraws: it prints the line as the phase starts, every // plainInterval while it runs, and as it ends. Either way the display // shows the phase is alive while it waits for its next item. // // The redraws come from a goroutine of their own, so mu guards what // they read. Every method must be called from the main goroutine. A // nil *progress is a valid no-display receiver: every method is a // no-op, so batched database flushes during the hash phase can reuse // the update phase's helpers without rendering anything. type progress struct { bar *progressbar.ProgressBar // nil when stderr is not a TTY stop chan struct{} // closed by finish to end the redraws done chan struct{} // closed once the redraws have ended mu sync.Mutex label string total int64 estimate bool waitingFor string count int64 start time.Time // when the phase started, for the elapsed time rateStart time.Time // when the count started, for the rate } // newProgress starts the display of one phase on stderr. func newProgress(params progressParams) *progress { if stderrIsTTY() { return startProgress(params, newSpinner(), redrawInterval) } return startProgress(params, nil, plainInterval) } // startProgress shows the phase params describes at once, on bar, or as // a plain line when bar is nil, and again every interval until finish. func startProgress(params progressParams, bar *progressbar.ProgressBar, interval time.Duration, ) *progress { now := time.Now() p := &progress{ bar: bar, stop: make(chan struct{}), done: make(chan struct{}), label: params.label, total: params.total, estimate: params.estimate, waitingFor: params.waitingFor, start: now, rateStart: now, } // Show the zero state at once: the first item may take minutes, // and a phase must never look hung. p.draw() go func() { defer close(p.done) ticker := time.NewTicker(interval) defer ticker.Stop() for { select { case <-p.stop: return case <-ticker.C: p.draw() } } }() return p } // newSpinner builds the TTY display for newProgress: the library's // spinner, followed by the text draw gives it. func newSpinner() *progressbar.ProgressBar { return progressbar.NewOptions64(-1, progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSpinnerType(spinnerType), // The spinner turns at each draw, not on a timer of the // library's own. progressbar.OptionSetSpinnerChangeInterval(0), // The text carries the elapsed time. progressbar.OptionSetElapsedTime(false), ) } // draw shows the phase's line: on a TTY in place of the one before, and // otherwise as a plain line. func (p *progress) draw() { p.mu.Lock() defer p.mu.Unlock() if p.bar == nil { fmt.Fprintln(os.Stderr, p.line()) return } p.bar.Describe(p.line()) } // setTotal gives a phase that started without a total its total, once // it is known. The rate, and so the ETA, count from then on: the time // spent waiting for the total says nothing about how fast items go. func (p *progress) setTotal(total int64) { if p == nil { return } p.mu.Lock() defer p.mu.Unlock() p.total = total p.rateStart = time.Now() } // increment records one completed item, which the display shows at its // next draw. func (p *progress) increment() { if p == nil { return } p.mu.Lock() defer p.mu.Unlock() p.count++ } // warnf prints a one-line warning to stderr without corrupting the // display. The whole message is escaped like a report's path columns, // so a path holding a newline cannot split the warning. func (p *progress) warnf(format string, args ...any) { if p == nil { return } msg := escapePath(fmt.Sprintf(format, args...)) p.mu.Lock() defer p.mu.Unlock() if p.bar == nil { fmt.Fprintln(os.Stderr, msg) return } // A direct write could land inside the line on the TTY. The bar // prints the warning on a line of its own, just before its next // draw. _, _ = progressbar.Bprintln(p.bar, msg) } // finish ends the phase's display with its line as the phase ends. A // phase stopped short of its total, as an interrupted one is, shows the // count it stopped at. func (p *progress) finish() { if p == nil { return } close(p.stop) <-p.done p.draw() if p.bar != nil { // Finish draws the line once more, below any warning it holds, // and leaves it on the TTY. _ = p.bar.Finish() fmt.Fprintln(os.Stderr) } } // line formats the phase's progress, as // // [3/5] hash: [1234/5000] 25%, 300 files/s, elapsed 4s, eta 13s // // with "~" before a total that is an estimate. A phase with no total // shows the items done instead, or what it waits for. p.mu must be // held. func (p *progress) line() string { elapsed := time.Since(p.start).Round(time.Second) if p.total < 0 && p.waitingFor != "" { return fmt.Sprintf("%s: %s, elapsed %s, eta unknown", p.label, p.waitingFor, elapsed) } rate := 0.0 if s := time.Since(p.rateStart).Seconds(); s > 0 { rate = float64(p.count) / s } done := fmt.Sprintf("%d files", p.count) if p.total >= 0 { approx := "" if p.estimate { approx = "~" } pct := float64(percentScale) if p.total > 0 { pct = percentScale * float64(p.count) / float64(p.total) } done = fmt.Sprintf("[%d/%s%d] %.0f%%", p.count, approx, p.total, pct) } return fmt.Sprintf("%s: %s, %.0f files/s, elapsed %s, eta %s", p.label, done, rate, elapsed, p.eta(rate)) } // eta returns the time the phase has left at rate items a second, or // "unknown": with no total, before the first item, and once the count // passes an estimate. p.mu must be held. func (p *progress) eta(rate float64) string { switch { case p.total < 0, p.estimate && p.count > p.total: return "unknown" case p.count >= p.total: return "0s" case rate <= 0: return "unknown" } remaining := time.Duration( float64(p.total-p.count) / rate * float64(time.Second)) return remaining.Round(time.Second).String() }