sfdupes/progress.go

package main

import (
	"fmt"
	"io"
	"math"
	"os"
	"slices"
	"strings"
	"sync"
	"time"
	"unicode/utf8"

	"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

// rateWindow is the time the rate a display shows covers: the items
// completed over the last rateWindow, which is the time since the line
// before when stderr is not a TTY.
const rateWindow = 5 * time.Second

// stallTime is how long a phase goes with no item completed before its
// line says it has stalled, and where it waits, in place of the rate.
const stallTime = 10 * time.Second

// spinnerType selects the progressbar library's spinner style.
const spinnerType = 14

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

// barCell is one column of the filled part of a bar.
const barCell = "█"

// barColumns is the width of what percentBar returns: the percent in
// four columns and its sign, a space, and the bar with its two ends.
const barColumns = barWidth + 8

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

// countingRecords is what the display of a phase that reads the records
// says while it counts them, before it has its total.
const countingRecords = "counting the records"

// openPhase is the name of the phase in which report, trees and info
// open the database, and openingDatabase what its display says while
// it does.
const (
	openPhase       = "open"
	openingDatabase = "opening the database"
)

// 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/4] 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
	stalledOn  string // where the phase waits, as "disk/db"; see stallTime
}

// drawnCount is the count of items completed as of one draw, and when
// that draw was.
type drawnCount struct {
	at    time.Time
	count int64
}

// progress renders one phase's progress on stderr. On a TTY the
// progressbar library draws its spinner and then line, fitted to the
// terminal's width. A phase with an exact total ends the line with its
// percent and bar, at the terminal's right edge, when the terminal has
// room for them. 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, and its writer writes
// with no display to clear.
type progress struct {
	now   func() time.Time // reads the time
	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
	widest     int                      // the widest line bar has drawn, in columns
	label      string
	unit       string
	total      int64
	estimate   bool
	waitingFor string
	stalledOn  string
	count      int64
	start      time.Time // when the phase started, for the elapsed time
	rateStart  time.Time // when the count started, for the time remaining
	lastItem   time.Time // when the last item was completed, or the count started
	// drawn holds the count at each draw over the last rateWindow, oldest
	// first; the rate counts from the first. It keeps the draw before the
	// latest even when that is older, as the line before is when lines
	// are plainInterval apart and a tick comes late.
	drawn []drawnCount
}

// newProgress starts the display of one phase on stderr.
func newProgress(params progressParams) *progress {
	if stderrIsTTY() {
		return startProgress(params, terminalWidth, redrawInterval, time.Now)
	}

	return startProgress(params, nil, plainInterval, time.Now)
}

// 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. The display reads the time by calling now.
func startProgress(params progressParams, width func() int,
	interval time.Duration, now func() time.Time,
) *progress {
	start := now()
	p := &progress{
		now: now, 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,
		stalledOn: params.stalledOn, start: start, rateStart: start,
		lastItem: start, drawn: []drawnCount{{at: start}},
	}

	// 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. It first adds the count to
// drawn, then drops from drawn the draws older than rateWindow, except
// the one before this draw.
func (p *progress) draw() {
	p.mu.Lock()
	defer p.mu.Unlock()

	now := p.now()

	p.drawn = append(p.drawn, drawnCount{at: now, count: p.count})
	for len(p.drawn) > 2 && now.Sub(p.drawn[0].at) > rateWindow {
		p.drawn = p.drawn[1:]
	}

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

		return
	}

	width := p.width()
	room := width - spinnerMargin
	text := p.line(now, room, true)

	// A phase with an exact total ends its text with its percent and bar
	// when its whole line, without the percent the bar shows, fits beside
	// them: the bar gives way before any part of the line. The line is
	// padded to the room the bar leaves it, so the bar stays at the
	// terminal's right edge however long the line is. A total of 0 leaves
	// the bar nothing to fill.
	if p.total > 0 && !p.estimate {
		line := p.line(now, math.MaxInt, false)
		if len(line) <= room-barColumns {
			text = fmt.Sprintf("%-*s%s", room-barColumns, line, p.percentBar())
		}
	}

	// Before each draw the library erases its line with as many spaces as
	// the widest line it has drawn. A display gets a new one 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.widest >= width {
		p.newBar(width)
	}

	p.bar.Describe(text)

	// The library drew the spinner, the text and the spaces around it,
	// less the terminal's last column, which spinnerMargin keeps clear. A
	// bar's cells take one column each but more than one byte.
	p.widest = max(p.widest, utf8.RuneCountInString(text)+spinnerMargin-1)
}

// newBar has the library make what draw draws: its spinner, followed by
// the text draw gives it and nothing of the library's own. p.mu must be
// held.
func (p *progress) newBar(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")
	}

	// A total of -1 gives the library's spinner.
	p.bar = 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 and the time remaining.
		progressbar.OptionSetElapsedTime(false),
		progressbar.OptionSetPredictTime(false),
	)
	p.widest = 0
}

// percentBar returns the percent and the bar that end the text of a
// phase with an exact total, as "  25% |██        |": the bar's
// barWidth columns fill in proportion to the items done. A count past
// the total, as when a scan adds records while trees reads them, shows
// as 100%. p.mu must be held.
func (p *progress) percentBar() string {
	done := min(float64(p.count)/float64(p.total), 1)
	filled := strings.Repeat(barCell, int(done*barWidth))

	return fmt.Sprintf("%4.0f%% |%-*s|", done*percentScale, barWidth, filled)
}

// setTotal gives a phase that started without a total its total, once
// it is known. The rate, the time remaining and the time with no item
// completed 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
	}

	now := p.now()

	p.mu.Lock()
	defer p.mu.Unlock()

	p.total = total
	p.rateStart, p.lastItem = now, now
	p.drawn = []drawnCount{{at: now, count: p.count}}
}

// increment records one completed item, which the display shows at its
// next draw.
func (p *progress) increment() {
	if p == nil {
		return
	}

	now := p.now()

	p.mu.Lock()
	defer p.mu.Unlock()

	p.count++
	p.lastItem = now
}

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

// progressWriter writes a phase's output to out, which may be the file
// the phase's display is drawn on, as stdout is in a plain "sfdupes
// trees" on a terminal. It holds what is written to it and writes that
// to out in pieces of about ioBufSize bytes. When out is the file
// stderr is, each piece is written under the display's lock with the
// display cleared first, as warnf does for a warning: no draw lands in
// the middle of a piece, and the next draw shows the display again, on
// the line below. Every write to it must end with a newline, as a row
// written with one fmt.Fprintf does, so that each piece ends a line.
// Flush writes what it holds once the last row is written.
type progressWriter struct {
	// p is the display to clear, nil unless out is the file stderr is.
	p   *progress
	out io.Writer
	buf []byte
}

// writer returns a progressWriter writing to out around p's display.
// It clears the display and takes its lock only when out is the file
// stderr is: the same device and inode, as os.SameFile compares them,
// which is one terminal, or one pipe or file that both are sent to.
// Otherwise it writes to out without the lock, so the display keeps
// being drawn while a write to out waits for a slow reader.
func (p *progress) writer(out io.Writer) *progressWriter {
	w := &progressWriter{out: out}

	if f, ok := out.(*os.File); ok {
		outInfo, outErr := f.Stat()
		errInfo, errErr := os.Stderr.Stat()

		if outErr == nil && errErr == nil && os.SameFile(outInfo, errInfo) {
			w.p = p
		}
	}

	return w
}

// Write holds b, and writes what it holds to out once that comes to
// ioBufSize bytes.
func (w *progressWriter) Write(b []byte) (int, error) {
	w.buf = append(w.buf, b...)
	if len(w.buf) < ioBufSize {
		return len(b), nil
	}

	return len(b), w.Flush()
}

// Flush writes what w holds to out. When out is the file stderr is, the
// display is cleared first and not drawn again until it is written.
func (w *progressWriter) Flush() error {
	if w.p != nil {
		w.p.mu.Lock()
		defer w.p.mu.Unlock()

		if w.p.bar != nil {
			_ = w.p.bar.Clear()
		}
	}

	_, err := w.out.Write(w.buf)
	w.buf = w.buf[:0]

	return err
}

// 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 as of now in at most room columns,
// as
//
//	[3/4] hash: [1234/5000] 25%, 300 files/s, elapsed 4s, remain 13s
//
// where "files" is the phase's unit, without the percent when percent
// is false, as before the bar, which shows its own. A phase with no
// total shows the items done instead, or what it waits for. The rate is
// that since the first draw in drawn, and the time remaining comes from
// the rate since the count started. Once no item has been completed for
// stallTime, the rate, which would read 0, gives way to the word
// "stalled" and where the phase waits, as "stalled disk/db".
//
// The phase and its place, the progress, the elapsed time and the time
// remaining are kept whole. A line wider than room leaves out the rate,
// the words of a stalled phase, or what the phase waits for; only if it
// is still wider is it cut at room. Every part is ASCII, so its length
// is its width in columns. p.mu must be held.
func (p *progress) line(now time.Time, room int, percent bool) string {
	elapsed := now.Sub(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, remain unknown", elapsed)
	} else {
		since := p.drawn[0]
		rate := perSecond(p.count-since.count, now.Sub(since.at))
		average := perSecond(p.count, now.Sub(p.rateStart))

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

		if now.Sub(p.lastItem) >= stallTime {
			extra = "stalled " + p.stalledOn + ", "
		}
	}

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

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

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

// perSecond returns count items over d as items a second, or 0 when no
// time has passed.
func perSecond(count int64, d time.Duration) float64 {
	if d <= 0 {
		return 0
	}

	return float64(count) / d.Seconds()
}

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

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

	left := time.Duration(
		float64(p.total-p.count) / rate * float64(time.Second))

	return left.Round(time.Second).String()
}