sfdupes/progress.go

package main

import (
	"fmt"
	"io"
	"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

// 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 what it is doing, 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 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

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

// reading is what the display of a phase that reads files or records
// says it is doing once it has stalled; see stallTime.
const reading = "reading"

// 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
	doing      string // what the phase is doing, as "reading"; 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 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, 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
	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
	doing      string
	count      int64
	start      time.Time // when the phase started, for the elapsed time
	rateStart  time.Time // when the count started, for the ETA
	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,
		doing: params.doing, 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 {
		text, _ := p.line(now, 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(now, width-barMargin, false)
	}

	margin := barMargin
	if !withBar {
		margin = spinnerMargin
		text, _ = p.line(now, 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, the ETA 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/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
// rate is that since the first draw in drawn, and the ETA 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 what the phase is doing, as "stalled reading".
//
// The phase and its place, the progress, what a stalled phase is doing,
// 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(now time.Time, room int,
	percent bool,
) (string, bool) {
	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, eta 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, eta %s", elapsed, p.eta(average))

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

	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
}

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

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