corso/src/internal/observe/observe.go

package observe

import (
	"context"
	"fmt"
	"io"
	"os"
	"sync"

	"github.com/spf13/cobra"
	"github.com/spf13/pflag"
	"github.com/vbauerster/mpb/v8"
	"github.com/vbauerster/mpb/v8/decor"
)

const (
	noProgressBarsFN = "no-progress-bars"
	progressBarWidth = 32
)

var (
	wg sync.WaitGroup
	// TODO: Revisit this being a global nd make it a parameter to the progress methods
	// so that each bar can be initialized with different contexts if needed.
	con      context.Context
	writer   io.Writer
	progress *mpb.Progress
	cfg      *config
)

func init() {
	cfg = &config{}

	makeSpinFrames(progressBarWidth)
}

// adds the persistent boolean flag --no-progress-bars to the provided command.
// This is a hack for help displays.  Due to seeding the context, we also
// need to parse the configuration before we execute the command.
func AddProgressBarFlags(parent *cobra.Command) {
	fs := parent.PersistentFlags()
	fs.Bool(noProgressBarsFN, false, "turn off the progress bar displays")
}

// Due to races between the lazy evaluation of flags in cobra and the need to init observer
// behavior in a ctx, these options get pre-processed manually here using pflags.  The canonical
// AddProgressBarFlag() ensures the flags are displayed as part of the help/usage output.
func PreloadFlags() bool {
	fs := pflag.NewFlagSet("seed-observer", pflag.ContinueOnError)
	fs.ParseErrorsWhitelist.UnknownFlags = true
	fs.Bool(noProgressBarsFN, false, "turn off the progress bar displays")
	// prevents overriding the corso/cobra help processor
	fs.BoolP("help", "h", false, "")

	// parse the os args list to find the log level flag
	if err := fs.Parse(os.Args[1:]); err != nil {
		return false
	}

	// retrieve the user's preferred display
	// automatically defaults to "info"
	shouldHide, err := fs.GetBool(noProgressBarsFN)
	if err != nil {
		return false
	}

	return shouldHide
}

// ---------------------------------------------------------------------------
// configuration
// ---------------------------------------------------------------------------

// config handles observer configuration
type config struct {
	doNotDisplay bool
}

// SeedWriter adds default writer to the observe package.
// Uses a noop writer until seeded.
func SeedWriter(ctx context.Context, w io.Writer, hide bool) {
	writer = w
	con = ctx

	if con == nil {
		con = context.Background()
	}

	cfg = &config{
		doNotDisplay: hide,
	}

	progress = mpb.NewWithContext(
		con,
		mpb.WithWidth(progressBarWidth),
		mpb.WithWaitGroup(&wg),
		mpb.WithOutput(writer),
	)
}

// Complete blocks until the progress finishes writing out all data.
// Afterwards, the progress instance is reset.
func Complete() {
	if progress != nil {
		progress.Wait()
	}

	SeedWriter(con, writer, cfg.doNotDisplay)
}

const (
	ItemBackupMsg  = "Backing up item:"
	ItemRestoreMsg = "Restoring item:"
	ItemQueueMsg   = "Queuing items:"
)

// Progress Updates

// Message is used to display a progress message
func Message(message string) {
	if writer == nil {
		return
	}

	wg.Add(1)

	bar := progress.New(
		-1,
		mpb.NopStyle(),
		mpb.PrependDecorators(
			decor.Name(message, decor.WC{W: len(message) + 1, C: decor.DidentRight}),
		),
	)

	// Complete the bar immediately
	bar.SetTotal(-1, true)

	waitAndCloseBar(bar)()
}

// MessageWithCompletion is used to display progress with a spinner
// that switches to "done" when the completion channel is signalled
func MessageWithCompletion(message string) (chan<- struct{}, func()) {
	completionCh := make(chan struct{}, 1)

	if writer == nil {
		return completionCh, func() {}
	}

	wg.Add(1)

	frames := []string{"∙∙∙", "●∙∙", "∙●∙", "∙∙●", "∙∙∙"}

	bar := progress.New(
		-1,
		mpb.SpinnerStyle(frames...).PositionLeft(),
		mpb.PrependDecorators(
			decor.Name(message),
			decor.Elapsed(decor.ET_STYLE_GO, decor.WC{W: 4}),
		),
		mpb.BarFillerOnComplete("done"),
	)

	go func(ci <-chan struct{}) {
		for {
			select {
			case <-con.Done():
				bar.SetTotal(-1, true)
			case <-ci:
				// We don't care whether the channel was signalled or closed
				// Use either one as an indication that the bar is done
				bar.SetTotal(-1, true)
			}
		}
	}(completionCh)

	return completionCh, waitAndCloseBar(bar)
}

// ---------------------------------------------------------------------------
// Progress for Known Quantities
// ---------------------------------------------------------------------------

// ItemProgress tracks the display of an item in a folder by counting the bytes
// read through the provided readcloser, up until the byte count matches
// the totalBytes.
func ItemProgress(rc io.ReadCloser, header, iname string, totalBytes int64) (io.ReadCloser, func()) {
	if cfg.doNotDisplay || writer == nil || rc == nil || totalBytes == 0 {
		return rc, func() {}
	}

	wg.Add(1)

	bar := progress.New(
		totalBytes,
		mpb.NopStyle(),
		mpb.BarRemoveOnComplete(),
		mpb.PrependDecorators(
			decor.Name(header, decor.WCSyncSpaceR),
			decor.Name(iname, decor.WCSyncSpaceR),
			decor.CountersKibiByte(" %.1f/%.1f ", decor.WC{W: 8}),
			decor.NewPercentage("%d ", decor.WC{W: 4}),
		),
	)

	return bar.ProxyReader(rc), waitAndCloseBar(bar)
}

// ProgressWithCount tracks the display of a bar that tracks the completion
// of the specified count.
// Each write to the provided channel counts as a single increment.
// The caller is expected to close the channel.
func ProgressWithCount(header, message string, count int64) (chan<- struct{}, func()) {
	progressCh := make(chan struct{})

	if cfg.doNotDisplay || writer == nil {
		go func(ci <-chan struct{}) {
			for {
				_, ok := <-ci
				if !ok {
					return
				}
			}
		}(progressCh)

		return progressCh, func() {}
	}

	wg.Add(1)

	bar := progress.New(
		count,
		mpb.NopStyle(),
		mpb.BarRemoveOnComplete(),
		mpb.PrependDecorators(
			decor.Name(header, decor.WCSyncSpaceR),
			decor.Counters(0, " %d/%d "),
			decor.Name(message),
		),
	)

	ch := make(chan struct{})

	go func(ci <-chan struct{}) {
		for {
			select {
			case <-con.Done():
				bar.Abort(true)
				return

			case _, ok := <-ci:
				if !ok {
					bar.Abort(true)
					return
				}

				bar.Increment()
			}
		}
	}(ch)

	return ch, waitAndCloseBar(bar)
}

// ---------------------------------------------------------------------------
// Progress for Unknown Quantities
// ---------------------------------------------------------------------------

var spinFrames []string

// The bar width is set to a static 32 characters.  The default spinner is only
// one char wide, which puts a lot of white space between it and the useful text.
// This builds a custom spinner animation to fill up that whitespace for a cleaner
// display.
func makeSpinFrames(barWidth int) {
	s, l := rune('∙'), rune('●')

	line := []rune{}
	for i := 0; i < barWidth; i++ {
		line = append(line, s)
	}

	sl := make([]string, 0, barWidth+1)
	sl = append(sl, string(line))

	for i := 1; i < barWidth; i++ {
		l2 := make([]rune, len(line))
		copy(l2, line)
		l2[i] = l

		sl = append(sl, string(l2))
	}

	spinFrames = sl
}

// CollectionProgress tracks the display a spinner that idles while the collection
// incrementing the count of items handled.  Each write to the provided channel
// counts as a single increment.  The caller is expected to close the channel.
func CollectionProgress(user, category, dirName string) (chan<- struct{}, func()) {
	if cfg.doNotDisplay || writer == nil || len(user) == 0 || len(dirName) == 0 {
		ch := make(chan struct{})

		go func(ci <-chan struct{}) {
			for {
				_, ok := <-ci
				if !ok {
					return
				}
			}
		}(ch)

		return ch, func() {}
	}

	wg.Add(1)

	bar := progress.New(
		-1, // -1 to indicate an unbounded count
		mpb.SpinnerStyle(spinFrames...),
		mpb.BarRemoveOnComplete(),
		mpb.PrependDecorators(
			decor.Name(category),
		),
		mpb.AppendDecorators(
			decor.CurrentNoUnit("%d - ", decor.WCSyncSpace),
			decor.Name(fmt.Sprintf("%s - %s", user, dirName)),
		),
	)

	ch := make(chan struct{})

	go func(ci <-chan struct{}) {
		for {
			select {
			case <-con.Done():
				bar.SetTotal(-1, true)
				return

			case _, ok := <-ci:
				if !ok {
					bar.SetTotal(-1, true)
					return
				}

				bar.Increment()
			}
		}
	}(ch)

	return ch, waitAndCloseBar(bar)
}

func waitAndCloseBar(bar *mpb.Bar) func() {
	return func() {
		bar.Wait()
		wg.Done()
	}
}