summaryrefslogtreecommitdiff
path: root/vendor/github.com/vbauerster/mpb/v7/bar.go
diff options
context:
space:
mode:
Diffstat (limited to 'vendor/github.com/vbauerster/mpb/v7/bar.go')
-rw-r--r--vendor/github.com/vbauerster/mpb/v7/bar.go492
1 files changed, 492 insertions, 0 deletions
diff --git a/vendor/github.com/vbauerster/mpb/v7/bar.go b/vendor/github.com/vbauerster/mpb/v7/bar.go
new file mode 100644
index 000000000..ed6c73eda
--- /dev/null
+++ b/vendor/github.com/vbauerster/mpb/v7/bar.go
@@ -0,0 +1,492 @@
+package mpb
+
+import (
+ "bytes"
+ "context"
+ "fmt"
+ "io"
+ "log"
+ "runtime/debug"
+ "strings"
+ "time"
+
+ "github.com/acarl005/stripansi"
+ "github.com/mattn/go-runewidth"
+ "github.com/vbauerster/mpb/v7/decor"
+)
+
+// Bar represents a progress bar.
+type Bar struct {
+ priority int // used by heap
+ index int // used by heap
+
+ extendedLines int
+ toShutdown bool
+ toDrop bool
+ noPop bool
+ hasEwmaDecorators bool
+ operateState chan func(*bState)
+ frameCh chan io.Reader
+ syncTableCh chan [][]chan int
+ completed chan bool
+
+ // cancel is called either by user or on complete event
+ cancel func()
+ // done is closed after cacheState is assigned
+ done chan struct{}
+ // cacheState is populated, right after close(shutdown)
+ cacheState *bState
+
+ container *Progress
+ dlogger *log.Logger
+ recoveredPanic interface{}
+}
+
+type extenderFunc func(in io.Reader, reqWidth int, st decor.Statistics) (out io.Reader, lines int)
+
+// bState is actual bar state. It gets passed to *Bar.serve(...) monitor
+// goroutine.
+type bState struct {
+ id int
+ priority int
+ reqWidth int
+ total int64
+ current int64
+ refill int64
+ lastN int64
+ iterated bool
+ trimSpace bool
+ completed bool
+ completeFlushed bool
+ triggerComplete bool
+ dropOnComplete bool
+ noPop bool
+ aDecorators []decor.Decorator
+ pDecorators []decor.Decorator
+ averageDecorators []decor.AverageDecorator
+ ewmaDecorators []decor.EwmaDecorator
+ shutdownListeners []decor.ShutdownListener
+ bufP, bufB, bufA *bytes.Buffer
+ filler BarFiller
+ middleware func(BarFiller) BarFiller
+ extender extenderFunc
+
+ // runningBar is a key for *pState.parkedBars
+ runningBar *Bar
+
+ debugOut io.Writer
+}
+
+func newBar(container *Progress, bs *bState) *Bar {
+ logPrefix := fmt.Sprintf("%sbar#%02d ", container.dlogger.Prefix(), bs.id)
+ ctx, cancel := context.WithCancel(container.ctx)
+
+ bar := &Bar{
+ container: container,
+ priority: bs.priority,
+ toDrop: bs.dropOnComplete,
+ noPop: bs.noPop,
+ operateState: make(chan func(*bState)),
+ frameCh: make(chan io.Reader, 1),
+ syncTableCh: make(chan [][]chan int, 1),
+ completed: make(chan bool, 1),
+ done: make(chan struct{}),
+ cancel: cancel,
+ dlogger: log.New(bs.debugOut, logPrefix, log.Lshortfile),
+ }
+
+ go bar.serve(ctx, bs)
+ return bar
+}
+
+// ProxyReader wraps r with metrics required for progress tracking.
+// Panics if r is nil.
+func (b *Bar) ProxyReader(r io.Reader) io.ReadCloser {
+ if r == nil {
+ panic("expected non nil io.Reader")
+ }
+ return newProxyReader(r, b)
+}
+
+// ID returs id of the bar.
+func (b *Bar) ID() int {
+ result := make(chan int)
+ select {
+ case b.operateState <- func(s *bState) { result <- s.id }:
+ return <-result
+ case <-b.done:
+ return b.cacheState.id
+ }
+}
+
+// Current returns bar's current number, in other words sum of all increments.
+func (b *Bar) Current() int64 {
+ result := make(chan int64)
+ select {
+ case b.operateState <- func(s *bState) { result <- s.current }:
+ return <-result
+ case <-b.done:
+ return b.cacheState.current
+ }
+}
+
+// SetRefill sets refill flag with specified amount.
+// The underlying BarFiller will change its visual representation, to
+// indicate refill event. Refill event may be referred to some retry
+// operation for example.
+func (b *Bar) SetRefill(amount int64) {
+ select {
+ case b.operateState <- func(s *bState) {
+ s.refill = amount
+ }:
+ case <-b.done:
+ }
+}
+
+// TraverseDecorators traverses all available decorators and calls cb func on each.
+func (b *Bar) TraverseDecorators(cb func(decor.Decorator)) {
+ select {
+ case b.operateState <- func(s *bState) {
+ for _, decorators := range [...][]decor.Decorator{
+ s.pDecorators,
+ s.aDecorators,
+ } {
+ for _, d := range decorators {
+ cb(extractBaseDecorator(d))
+ }
+ }
+ }:
+ case <-b.done:
+ }
+}
+
+// SetTotal sets total dynamically.
+// If total is less than or equal to zero it takes progress' current value.
+func (b *Bar) SetTotal(total int64, triggerComplete bool) {
+ select {
+ case b.operateState <- func(s *bState) {
+ s.triggerComplete = triggerComplete
+ if total <= 0 {
+ s.total = s.current
+ } else {
+ s.total = total
+ }
+ if s.triggerComplete && !s.completed {
+ s.current = s.total
+ s.completed = true
+ go b.refreshTillShutdown()
+ }
+ }:
+ case <-b.done:
+ }
+}
+
+// SetCurrent sets progress' current to an arbitrary value.
+// Setting a negative value will cause a panic.
+func (b *Bar) SetCurrent(current int64) {
+ select {
+ case b.operateState <- func(s *bState) {
+ s.iterated = true
+ s.lastN = current - s.current
+ s.current = current
+ if s.triggerComplete && s.current >= s.total {
+ s.current = s.total
+ s.completed = true
+ go b.refreshTillShutdown()
+ }
+ }:
+ case <-b.done:
+ }
+}
+
+// Increment is a shorthand for b.IncrInt64(1).
+func (b *Bar) Increment() {
+ b.IncrInt64(1)
+}
+
+// IncrBy is a shorthand for b.IncrInt64(int64(n)).
+func (b *Bar) IncrBy(n int) {
+ b.IncrInt64(int64(n))
+}
+
+// IncrInt64 increments progress by amount of n.
+func (b *Bar) IncrInt64(n int64) {
+ select {
+ case b.operateState <- func(s *bState) {
+ s.iterated = true
+ s.lastN = n
+ s.current += n
+ if s.triggerComplete && s.current >= s.total {
+ s.current = s.total
+ s.completed = true
+ go b.refreshTillShutdown()
+ }
+ }:
+ case <-b.done:
+ }
+}
+
+// DecoratorEwmaUpdate updates all EWMA based decorators. Should be
+// called on each iteration, because EWMA's unit of measure is an
+// iteration's duration. Panics if called before *Bar.Incr... family
+// methods.
+func (b *Bar) DecoratorEwmaUpdate(dur time.Duration) {
+ select {
+ case b.operateState <- func(s *bState) {
+ ewmaIterationUpdate(false, s, dur)
+ }:
+ case <-b.done:
+ ewmaIterationUpdate(true, b.cacheState, dur)
+ }
+}
+
+// DecoratorAverageAdjust adjusts all average based decorators. Call
+// if you need to adjust start time of all average based decorators
+// or after progress resume.
+func (b *Bar) DecoratorAverageAdjust(start time.Time) {
+ select {
+ case b.operateState <- func(s *bState) {
+ for _, d := range s.averageDecorators {
+ d.AverageAdjust(start)
+ }
+ }:
+ case <-b.done:
+ }
+}
+
+// SetPriority changes bar's order among multiple bars. Zero is highest
+// priority, i.e. bar will be on top. If you don't need to set priority
+// dynamically, better use BarPriority option.
+func (b *Bar) SetPriority(priority int) {
+ select {
+ case <-b.done:
+ default:
+ b.container.setBarPriority(b, priority)
+ }
+}
+
+// Abort interrupts bar's running goroutine. Call this, if you'd like
+// to stop/remove bar before completion event. It has no effect after
+// completion event. If drop is true bar will be removed as well.
+func (b *Bar) Abort(drop bool) {
+ select {
+ case <-b.done:
+ default:
+ if drop {
+ b.container.dropBar(b)
+ }
+ b.cancel()
+ }
+}
+
+// Completed reports whether the bar is in completed state.
+func (b *Bar) Completed() bool {
+ select {
+ case b.operateState <- func(s *bState) { b.completed <- s.completed }:
+ return <-b.completed
+ case <-b.done:
+ return true
+ }
+}
+
+func (b *Bar) serve(ctx context.Context, s *bState) {
+ defer b.container.bwg.Done()
+ for {
+ select {
+ case op := <-b.operateState:
+ op(s)
+ case <-ctx.Done():
+ b.cacheState = s
+ close(b.done)
+ // Notifying decorators about shutdown event
+ for _, sl := range s.shutdownListeners {
+ sl.Shutdown()
+ }
+ return
+ }
+ }
+}
+
+func (b *Bar) render(tw int) {
+ select {
+ case b.operateState <- func(s *bState) {
+ stat := newStatistics(tw, s)
+ defer func() {
+ // recovering if user defined decorator panics for example
+ if p := recover(); p != nil {
+ if b.recoveredPanic == nil {
+ s.extender = makePanicExtender(p)
+ b.toShutdown = !b.toShutdown
+ b.recoveredPanic = p
+ }
+ frame, lines := s.extender(nil, s.reqWidth, stat)
+ b.extendedLines = lines
+ b.frameCh <- frame
+ b.dlogger.Println(p)
+ }
+ s.completeFlushed = s.completed
+ }()
+ frame, lines := s.extender(s.draw(stat), s.reqWidth, stat)
+ b.extendedLines = lines
+ b.toShutdown = s.completed && !s.completeFlushed
+ b.frameCh <- frame
+ }:
+ case <-b.done:
+ s := b.cacheState
+ stat := newStatistics(tw, s)
+ var r io.Reader
+ if b.recoveredPanic == nil {
+ r = s.draw(stat)
+ }
+ frame, lines := s.extender(r, s.reqWidth, stat)
+ b.extendedLines = lines
+ b.frameCh <- frame
+ }
+}
+
+func (b *Bar) subscribeDecorators() {
+ var averageDecorators []decor.AverageDecorator
+ var ewmaDecorators []decor.EwmaDecorator
+ var shutdownListeners []decor.ShutdownListener
+ b.TraverseDecorators(func(d decor.Decorator) {
+ if d, ok := d.(decor.AverageDecorator); ok {
+ averageDecorators = append(averageDecorators, d)
+ }
+ if d, ok := d.(decor.EwmaDecorator); ok {
+ ewmaDecorators = append(ewmaDecorators, d)
+ }
+ if d, ok := d.(decor.ShutdownListener); ok {
+ shutdownListeners = append(shutdownListeners, d)
+ }
+ })
+ select {
+ case b.operateState <- func(s *bState) {
+ s.averageDecorators = averageDecorators
+ s.ewmaDecorators = ewmaDecorators
+ s.shutdownListeners = shutdownListeners
+ }:
+ b.hasEwmaDecorators = len(ewmaDecorators) != 0
+ case <-b.done:
+ }
+}
+
+func (b *Bar) refreshTillShutdown() {
+ for {
+ select {
+ case b.container.refreshCh <- time.Now():
+ case <-b.done:
+ return
+ }
+ }
+}
+
+func (b *Bar) wSyncTable() [][]chan int {
+ select {
+ case b.operateState <- func(s *bState) { b.syncTableCh <- s.wSyncTable() }:
+ return <-b.syncTableCh
+ case <-b.done:
+ return b.cacheState.wSyncTable()
+ }
+}
+
+func (s *bState) draw(stat decor.Statistics) io.Reader {
+ nlr := strings.NewReader("\n")
+ tw := stat.AvailableWidth
+ for _, d := range s.pDecorators {
+ str := d.Decor(stat)
+ stat.AvailableWidth -= runewidth.StringWidth(stripansi.Strip(str))
+ s.bufP.WriteString(str)
+ }
+ if stat.AvailableWidth < 1 {
+ trunc := strings.NewReader(runewidth.Truncate(stripansi.Strip(s.bufP.String()), tw, "…"))
+ s.bufP.Reset()
+ return io.MultiReader(trunc, nlr)
+ }
+
+ if !s.trimSpace && stat.AvailableWidth > 1 {
+ stat.AvailableWidth -= 2
+ s.bufB.WriteByte(' ')
+ defer s.bufB.WriteByte(' ')
+ }
+
+ tw = stat.AvailableWidth
+ for _, d := range s.aDecorators {
+ str := d.Decor(stat)
+ stat.AvailableWidth -= runewidth.StringWidth(stripansi.Strip(str))
+ s.bufA.WriteString(str)
+ }
+ if stat.AvailableWidth < 1 {
+ trunc := strings.NewReader(runewidth.Truncate(stripansi.Strip(s.bufA.String()), tw, "…"))
+ s.bufA.Reset()
+ return io.MultiReader(s.bufP, s.bufB, trunc, nlr)
+ }
+
+ s.filler.Fill(s.bufB, s.reqWidth, stat)
+
+ return io.MultiReader(s.bufP, s.bufB, s.bufA, nlr)
+}
+
+func (s *bState) wSyncTable() [][]chan int {
+ columns := make([]chan int, 0, len(s.pDecorators)+len(s.aDecorators))
+ var pCount int
+ for _, d := range s.pDecorators {
+ if ch, ok := d.Sync(); ok {
+ columns = append(columns, ch)
+ pCount++
+ }
+ }
+ var aCount int
+ for _, d := range s.aDecorators {
+ if ch, ok := d.Sync(); ok {
+ columns = append(columns, ch)
+ aCount++
+ }
+ }
+ table := make([][]chan int, 2)
+ table[0] = columns[0:pCount]
+ table[1] = columns[pCount : pCount+aCount : pCount+aCount]
+ return table
+}
+
+func newStatistics(tw int, s *bState) decor.Statistics {
+ return decor.Statistics{
+ ID: s.id,
+ AvailableWidth: tw,
+ Total: s.total,
+ Current: s.current,
+ Refill: s.refill,
+ Completed: s.completeFlushed,
+ }
+}
+
+func extractBaseDecorator(d decor.Decorator) decor.Decorator {
+ if d, ok := d.(decor.Wrapper); ok {
+ return extractBaseDecorator(d.Base())
+ }
+ return d
+}
+
+func ewmaIterationUpdate(done bool, s *bState, dur time.Duration) {
+ if !done && !s.iterated {
+ panic("increment required before ewma iteration update")
+ } else {
+ s.iterated = false
+ }
+ for _, d := range s.ewmaDecorators {
+ d.EwmaUpdate(s.lastN, dur)
+ }
+}
+
+func makePanicExtender(p interface{}) extenderFunc {
+ pstr := fmt.Sprint(p)
+ stack := debug.Stack()
+ stackLines := bytes.Count(stack, []byte("\n"))
+ return func(_ io.Reader, _ int, st decor.Statistics) (io.Reader, int) {
+ mr := io.MultiReader(
+ strings.NewReader(runewidth.Truncate(pstr, st.AvailableWidth, "…")),
+ strings.NewReader(fmt.Sprintf("\n%#v\n", st)),
+ bytes.NewReader(stack),
+ )
+ return mr, stackLines + 1
+ }
+}