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 // fallbackWidth is the width the progressbar library takes a terminal // to have when it cannot read it, as when stderr is not a terminal. const fallbackWidth = 80 // spinnerMargin is the columns of a spinner's line besides its text: // the spinner and the three spaces the library draws around the text, // and the terminal's last column, which the line stays clear of, as the // library's bars do. const spinnerMargin = 5 // stderrIsTTY reports whether stderr is attached to a terminal. func stderrIsTTY() bool { return term.IsTerminal(int(os.Stderr.Fd())) } // terminalWidth returns the width of the terminal on stderr, or // fallbackWidth, as the progressbar library does, when it cannot read // one. func terminalWidth() int { width, _, err := term.GetSize(int(os.Stderr.Fd())) if err != nil || width <= 0 { return fallbackWidth } return width } // 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. On a TTY the // progressbar library draws it. A phase with an exact total from its // start gets the library's bar, with the phase and its place as the // bar's description and the library's counts, rate, elapsed time and // ETA beside it; the bar takes the room the rest leaves on the line. Any // other phase gets the library's spinner, with line as its text, cut to // the terminal's width. Either is drawn again every redrawInterval, so // its elapsed time keeps moving while the phase waits for its next // item. When stderr is not a TTY it emits no ANSI redraws: it prints // line as the phase starts, every plainInterval while it runs, and as it // ends. // // The draws come from a goroutine of their own, so mu guards the fields // below it and every call on bar. 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 spinner bool // bar is a spinner, with line as its text stop chan struct{} // closed by finish to end the draws done chan struct{} // closed once the draws have ended mu sync.Mutex label string total int64 estimate bool waitingFor string count int64 drawn int64 // the items bar has been given, unless a spinner 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, true, redrawInterval) } return startProgress(params, false, plainInterval) } // startProgress shows the phase params describes at once, with the // progressbar library when onTTY is set and as a plain line when not, // and again every interval until finish. func startProgress(params progressParams, onTTY bool, interval time.Duration, ) *progress { now := time.Now() p := &progress{ stop: make(chan struct{}), done: make(chan struct{}), label: params.label, total: params.total, estimate: params.estimate, waitingFor: params.waitingFor, start: now, rateStart: now, } switch { case onTTY && (params.total < 0 || params.estimate): // The library's bar needs an exact total. p.bar, p.spinner = newSpinner(), true case onTTY: p.bar = newBar(params.label, params.total) } // 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 } // newBar builds the library's bar for a phase with an exact total, with // label as its description and the library's counts, rate, elapsed time // and ETA beside it. The bar is as wide as the rest of the line leaves // room for on the terminal. func newBar(label string, total int64) *progressbar.ProgressBar { return progressbar.NewOptions64(total, progressbar.OptionSetWriter(os.Stderr), progressbar.OptionSetDescription(label), progressbar.OptionShowCount(), progressbar.OptionShowIts(), progressbar.OptionSetItsString("files"), progressbar.OptionSetElapsedTime(true), progressbar.OptionSetPredictTime(true), progressbar.OptionShowElapsedTimeOnFinish(), progressbar.OptionFullWidth(), ) } // newSpinner builds the library's spinner for a phase without an exact // total, 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 progress: on a TTY in place of the line // before, and otherwise as a plain line. func (p *progress) draw() { p.mu.Lock() defer p.mu.Unlock() switch { case p.bar == nil: fmt.Fprintln(os.Stderr, p.line()) case p.spinner: text := p.line() if room := terminalWidth() - spinnerMargin; len(text) > room { text = text[:max(room, 0)] } p.bar.Describe(text) case p.count > p.drawn: // The library measures its rate over the items it is given, // so it is given items only when there are new ones: an // empty add would count a wait as a rate of zero. _ = p.bar.Add64(p.count - p.drawn) p.drawn = p.count default: // No new items: the bar is drawn again all the same, with // its elapsed time as of now. p.bar.Describe(p.label) } } // 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 { // The warning takes the display's line; the next draw shows // the display again, on the line below. _ = p.bar.Clear() } fmt.Fprintln(os.Stderr, msg) } // finish ends the phase's display, drawn once more as the phase ends. // A phase stopped short of its total, as an interrupted one is, shows // the count it stopped at rather than a full bar. func (p *progress) finish() { if p == nil { return } close(p.stop) <-p.done p.draw() if p.bar != nil { // The display stays on the TTY as drawn; what follows starts // on a line of its own. 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() }