sfdupes/progress.go

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