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

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