From ea42e1231e4b191c8c6b20882598823591c8d9cb Mon Sep 17 00:00:00 2001 From: Gabe Cook Date: Thu, 19 Oct 2023 17:46:30 -0500 Subject: [PATCH] fix(progress): Make progress bars more responsive --- internal/actions/dump/dump.go | 1 + internal/actions/restore/restore.go | 1 + internal/progressbar/logger.go | 37 ++++++++++++ internal/progressbar/new.go | 78 ------------------------ internal/progressbar/progressbar.go | 92 +++++++++++++++++++++++++++++ 5 files changed, 131 insertions(+), 78 deletions(-) create mode 100644 internal/progressbar/logger.go delete mode 100644 internal/progressbar/new.go create mode 100644 internal/progressbar/progressbar.go diff --git a/internal/actions/dump/dump.go b/internal/actions/dump/dump.go index 00804cdd..9a95c09c 100644 --- a/internal/actions/dump/dump.go +++ b/internal/actions/dump/dump.go @@ -68,6 +68,7 @@ func (action Dump) Run(ctx context.Context) (err error) { startTime := time.Now() bar := progressbar.New(-1, "downloading", action.Spinner) + defer bar.Close() plogger := progressbar.NewBarSafeLogger(os.Stderr, bar) log.SetOutput(plogger) diff --git a/internal/actions/restore/restore.go b/internal/actions/restore/restore.go index 98f9e452..5d7d33b5 100644 --- a/internal/actions/restore/restore.go +++ b/internal/actions/restore/restore.go @@ -46,6 +46,7 @@ func (action Restore) Run(ctx context.Context) (err error) { startTime := time.Now() bar := progressbar.New(-1, "uploading", action.Spinner) + defer bar.Close() errLog := progressbar.NewBarSafeLogger(os.Stderr, bar) outLog := progressbar.NewBarSafeLogger(os.Stdout, bar) log.SetOutput(errLog) diff --git a/internal/progressbar/logger.go b/internal/progressbar/logger.go new file mode 100644 index 00000000..2e47342b --- /dev/null +++ b/internal/progressbar/logger.go @@ -0,0 +1,37 @@ +package progressbar + +import ( + "bytes" + "io" +) + +func NewBarSafeLogger(w io.Writer, bar *ProgressBar) *BarSafeLogger { + return &BarSafeLogger{ + out: w, + bar: bar, + } +} + +type BarSafeLogger struct { + out io.Writer + bar *ProgressBar + buf bytes.Buffer +} + +func (l *BarSafeLogger) Write(p []byte) (int, error) { + if l.bar.IsFinished() { + return l.out.Write(p) + } + + l.buf.Write([]byte("\r\x1B[K")) + l.buf.Write(p) + if p[len(p)-1] == '\n' { + l.buf.WriteString(l.bar.String()) + } + l.bar.mu.Lock() + defer l.bar.mu.Unlock() + if n, err := io.Copy(l.out, &l.buf); err != nil { + return int(n), err + } + return len(p), nil +} diff --git a/internal/progressbar/new.go b/internal/progressbar/new.go deleted file mode 100644 index bcc80562..00000000 --- a/internal/progressbar/new.go +++ /dev/null @@ -1,78 +0,0 @@ -package progressbar - -import ( - "bytes" - "fmt" - "io" - "os" - "time" - - "github.com/clevyr/kubedb/internal/config/flags" - "github.com/gabe565/go-spinners" - "github.com/mattn/go-isatty" - "github.com/schollz/progressbar/v3" - log "github.com/sirupsen/logrus" -) - -func New(max int64, label string, spinnerKey string) *progressbar.ProgressBar { - s, ok := spinner.Map[spinnerKey] - if !ok { - log.WithField("spinner", spinnerKey).Warn("invalid spinner") - s = spinner.Map[flags.DefaultSpinner] - } - - options := []progressbar.Option{ - progressbar.OptionSetDescription(label), - progressbar.OptionSetWriter(os.Stderr), - progressbar.OptionShowBytes(true), - progressbar.OptionSetWidth(10), - progressbar.OptionShowCount(), - progressbar.OptionSpinnerCustom(s.Frames), - progressbar.OptionFullWidth(), - } - - if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()) { - options = append(options, - progressbar.OptionThrottle(65*time.Millisecond), - progressbar.OptionSetRenderBlankState(true), - progressbar.OptionOnCompletion(func() { - _, _ = fmt.Fprint(os.Stderr, "\r\x1B[K") - }), - ) - } else { - options = append(options, - progressbar.OptionThrottle(2*time.Second), - ) - } - - return progressbar.NewOptions64(max, options...) -} - -func NewBarSafeLogger(w io.Writer, bar *progressbar.ProgressBar) *BarSafeLogger { - return &BarSafeLogger{ - out: w, - bar: bar, - } -} - -type BarSafeLogger struct { - out io.Writer - bar *progressbar.ProgressBar - buf bytes.Buffer -} - -func (l *BarSafeLogger) Write(p []byte) (int, error) { - if l.bar.IsFinished() { - return l.out.Write(p) - } - - l.buf.Write([]byte("\r\x1B[K")) - l.buf.Write(p) - if p[len(p)-1] == '\n' { - l.buf.WriteString(l.bar.String()) - } - if n, err := io.Copy(l.out, &l.buf); err != nil { - return int(n), err - } - return len(p), nil -} diff --git a/internal/progressbar/progressbar.go b/internal/progressbar/progressbar.go new file mode 100644 index 00000000..245bc4ec --- /dev/null +++ b/internal/progressbar/progressbar.go @@ -0,0 +1,92 @@ +package progressbar + +import ( + "fmt" + "io" + "os" + "sync" + "time" + + "github.com/clevyr/kubedb/internal/config/flags" + "github.com/gabe565/go-spinners" + "github.com/mattn/go-isatty" + "github.com/schollz/progressbar/v3" + log "github.com/sirupsen/logrus" +) + +func New(max int64, label string, spinnerKey string) *ProgressBar { + s, ok := spinner.Map[spinnerKey] + if !ok { + log.WithField("spinner", spinnerKey).Warn("invalid spinner") + s = spinner.Map[flags.DefaultSpinner] + } + + options := []progressbar.Option{ + progressbar.OptionSetDescription(label), + progressbar.OptionSetWriter(io.Discard), + progressbar.OptionShowBytes(true), + progressbar.OptionSetWidth(10), + progressbar.OptionShowCount(), + progressbar.OptionSpinnerCustom(s.Frames), + progressbar.OptionFullWidth(), + } + + var throttle time.Duration + if isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsCygwinTerminal(os.Stdout.Fd()) { + throttle = 65 * time.Millisecond + options = append(options, + progressbar.OptionSetRenderBlankState(true), + progressbar.OptionOnCompletion(func() { + _, _ = fmt.Fprint(os.Stderr, "\r\x1B[K") + }), + ) + } else { + throttle = 2 * time.Second + } + options = append(options, + progressbar.OptionThrottle(throttle), + ) + + cancelChan := make(chan struct{}) + bar := &ProgressBar{ + ProgressBar: progressbar.NewOptions64(max, options...), + cancelChan: cancelChan, + } + go func() { + for { + select { + case <-cancelChan: + return + case <-time.After(throttle): + if bar.IsFinished() { + return + } + if bar.mu.TryLock() { + _ = bar.RenderBlank() + _, _ = os.Stderr.Write([]byte(bar.String())) + bar.mu.Unlock() + } + } + } + }() + + return bar +} + +type ProgressBar struct { + *progressbar.ProgressBar + mu sync.Mutex + cancelChan chan struct{} + cancelOnce sync.Once +} + +func (p *ProgressBar) Finish() error { + p.Close() + return p.ProgressBar.Finish() +} + +func (p *ProgressBar) Close() { + p.cancelOnce.Do(func() { + close(p.cancelChan) + }) +}