VictoriaMetrics/app/vmctl/barpool/pool.go
Dmytro Kozlov 8b46bb0c41
vmctl: disable progress bar for prometheus snapshot migrations (#6385)
* deprecate `--vm-disable-progress-bar` in favour of `--disable-progress-bar`
* new `--disable-progress-bar` consistently disables usage of progress bar
for all migration modes.

https://github.com/VictoriaMetrics/VictoriaMetrics/issues/6367

---------

Signed-off-by: hagen1778 <roman@victoriametrics.com>
Co-authored-by: hagen1778 <roman@victoriametrics.com>
2024-06-10 12:20:52 +02:00

99 lines
2.3 KiB
Go

// Package barpool provides access to the global
// pool of progress bars, so they could be rendered
// altogether.
package barpool
import (
"fmt"
"io"
"os"
"github.com/cheggaaa/pb/v3"
"github.com/VictoriaMetrics/VictoriaMetrics/app/vmctl/terminal"
)
var isDisabled bool
// Disable sets progress bar to be no-op if v==true
func Disable(v bool) {
isDisabled = v
}
var pool = pb.NewPool()
// Bar is an interface for progress bar
type Bar interface {
Add(value int)
Increment()
Start()
Finish()
NewProxyReader(r io.Reader) *pb.Reader
}
type progressBar struct {
*pb.ProgressBar
}
func (pb *progressBar) Finish() { pb.ProgressBar.Finish() }
func (pb *progressBar) Start() { pb.ProgressBar.Start() }
func (pb *progressBar) Add(value int) { pb.ProgressBar.Add(value) }
func (pb *progressBar) Increment() { pb.ProgressBar.Increment() }
func (pb *progressBar) NewProxyReader(r io.Reader) *pb.Reader {
return pb.ProgressBar.NewProxyReader(r)
}
type progressBarNoOp struct{}
func (pbno *progressBarNoOp) Finish() {}
func (pbno *progressBarNoOp) Start() {}
func (pbno *progressBarNoOp) Add(int) {}
func (pbno *progressBarNoOp) Increment() {}
func (pbno *progressBarNoOp) NewProxyReader(_ io.Reader) *pb.Reader { return nil }
// Start starts the global pool
// Must be called after all progress bars were added
func Start() error {
if isDisabled {
return nil
}
return pool.Start()
}
// Stop stops the global pool
func Stop() {
if isDisabled {
return
}
_ = pool.Stop()
}
// AddWithTemplate adds bar with the given template
// to the global pool
func AddWithTemplate(format string, total int) Bar {
if isDisabled {
return &progressBarNoOp{}
}
tpl := getTemplate(format)
bar := pb.ProgressBarTemplate(tpl).New(total)
pool.Add(bar)
return &progressBar{bar}
}
// NewSingleProgress returns progress bar with given template
func NewSingleProgress(format string, total int) Bar {
if isDisabled {
return &progressBarNoOp{}
}
tpl := getTemplate(format)
return &progressBar{pb.ProgressBarTemplate(tpl).New(total)}
}
func getTemplate(format string) string {
isTerminal := terminal.IsTerminal(int(os.Stdout.Fd()))
if !isTerminal {
format = fmt.Sprintf("%s\n", format)
}
return format
}