sfdupes/progress.go

package main

import (
	"fmt"
	"math"
	"os"
	"slices"
	"strings"
	"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

// barWidth is the width of the progressbar library's bar between its
// two ends, as in the example line in README "Progress".
const barWidth = 10

// percentScale converts a fraction to a percentage.
const percentScale = 100

// fallbackWidth is the width taken for a terminal whose width cannot be
// read, as the progressbar library takes it.
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 every line stays clear of so
// that no terminal wraps it.
const spinnerMargin = 5

// barMargin is the columns of a bar's line besides its text: the
// percent and the bar the library draws after the text, the spaces
// around them, and the terminal's last column.
const barMargin = barWidth + 11

// 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
	unit       string // what the phase counts, as "files"
	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 line, fitted to the terminal's width: after
// the library's spinner, or, for a phase with an exact total, before the
// library's bar when the terminal has room for the bar beside it. It 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. 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 {
	width func() int    // reads the terminal's width; nil when not a TTY
	stop  chan struct{} // closed by finish to end the draws
	done  chan struct{} // closed once the draws have ended

	mu         sync.Mutex
	bar        *progressbar.ProgressBar // what the library draws; nil until then
	withBar    bool                     // bar is the library's bar, not its spinner
	widest     int                      // the widest line bar has drawn, in columns
	label      string
	unit       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, terminalWidth, redrawInterval)
	}

	return startProgress(params, nil, plainInterval)
}

// startProgress shows the phase params describes at once, and again
// every interval until finish: with the progressbar library on a
// terminal, whose width each draw reads by calling width, or as a plain
// line when width is nil.
func startProgress(params progressParams, width func() int,
	interval time.Duration,
) *progress {
	now := time.Now()
	p := &progress{
		width: width, stop: make(chan struct{}), done: make(chan struct{}),
		label: params.label, unit: params.unit, 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
}

// 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()

	if p.width == nil {
		text, _ := p.line(math.MaxInt, true)
		fmt.Fprintln(os.Stderr, text)

		return
	}

	width := p.width()

	// A phase with an exact total gets the library's bar when its line,
	// without the percent the bar shows, fits whole beside the bar. A
	// total of 0 leaves the bar nothing to fill.
	text, withBar := "", false
	if p.total > 0 && !p.estimate {
		text, withBar = p.line(width-barMargin, false)
	}

	margin := barMargin
	if !withBar {
		margin = spinnerMargin
		text, _ = p.line(width-margin, true)
	}

	// The library makes a bar or a spinner, and before each draw erases
	// its line with as many spaces as the widest line it has drawn. A
	// display gets a new one when it turns from one to the other, as its
	// line grows or the terminal is resized, and when the terminal has
	// been made too narrow for that erase, which would wrap and leave a
	// line behind at every draw.
	if p.bar == nil || p.withBar != withBar || p.widest >= width {
		p.newBar(withBar, width)
	}

	if withBar {
		_ = p.bar.Set64(p.count)
	}

	p.bar.Describe(text)

	// The library drew the text and its margin, less the terminal's last
	// column, which the margin keeps clear.
	p.widest = max(p.widest, len(text)+margin-1)
}

// newBar has the library make what draw draws: a bar for the phase's
// total when withBar is set, and a spinner when not. Either shows the
// text draw gives it, and nothing of its own besides the spinner, or
// the bar and its percent. p.mu must be held.
func (p *progress) newBar(withBar bool, width int) {
	if p.bar != nil {
		// Erase what was drawn before, short of the last column of a
		// terminal that may have been made narrower since: the library's
		// own erase is as wide as the widest line it drew.
		blank := strings.Repeat(" ", min(p.widest, width-1))
		fmt.Fprint(os.Stderr, "\r"+blank+"\r")
	}

	total := int64(-1) // the library's spinner
	if withBar {
		total = p.total
	}

	p.bar = progressbar.NewOptions64(total,
		progressbar.OptionSetWriter(os.Stderr),
		progressbar.OptionSetWidth(barWidth),
		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 and the ETA.
		progressbar.OptionSetElapsedTime(false),
		progressbar.OptionSetPredictTime(false),
	)
	p.withBar, p.widest = withBar, 0
}

// 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.width != 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 in at most room columns, as
//
//	[3/5] hash: [1234/5000] 25%, 300 files/s, elapsed 4s, eta 13s
//
// where "files" is the phase's unit, without the percent when percent
// is false, as before the library's bar, which shows its own. A phase
// with no total shows the items done instead, or what it waits for.
//
// The phase and its place, the progress, the elapsed time and the ETA
// are kept whole. A line wider than room leaves out the rate, or what
// the phase waits for; only if it is still wider is it cut at room, and
// line reports that it does not fit. Every part is ASCII, so its length
// is its width in columns. p.mu must be held.
func (p *progress) line(room int, percent bool) (string, bool) {
	elapsed := time.Since(p.start).Round(time.Second)

	// head and tail hold the parts kept whole; extra, between them, is
	// the part left out.
	var head, extra, tail string

	if p.total < 0 && p.waitingFor != "" {
		head, extra = p.label+": ", p.waitingFor
		tail = fmt.Sprintf("elapsed %s, eta unknown", elapsed)
	} else {
		rate := 0.0
		if s := time.Since(p.rateStart).Seconds(); s > 0 {
			rate = float64(p.count) / s
		}

		head = p.label + ": " + p.itemsDone(percent) + ", "
		extra = fmt.Sprintf("%.0f %s/s", rate, p.unit)
		tail = fmt.Sprintf("elapsed %s, eta %s", elapsed, p.eta(rate))
	}

	full, kept := head+extra+", "+tail, head+tail

	switch {
	case len(full) <= room:
		return full, true
	case len(kept) <= room:
		return kept, true
	}

	return kept[:max(room, 0)], false
}

// itemsDone formats the items done, as "[1234/5000] 25%", with "~"
// before a total that is an estimate and without the percent when
// percent is false, or as "1234 files", in the phase's unit, when the
// phase has no total. p.mu must be held.
func (p *progress) itemsDone(percent bool) string {
	if p.total < 0 {
		return fmt.Sprintf("%d %s", p.count, p.unit)
	}

	approx := ""
	if p.estimate {
		approx = "~"
	}

	done := fmt.Sprintf("[%d/%s%d]", p.count, approx, p.total)
	if !percent {
		return done
	}

	pct := float64(percentScale)
	if p.total > 0 {
		pct = percentScale * float64(p.count) / float64(p.total)
	}

	return fmt.Sprintf("%s %.0f%%", done, pct)
}

// 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()
}