peridot/vendor/github.com/vbauerster/mpb/v7/bar.go

493 lines
12 KiB
Go
Raw Normal View History

2022-07-07 20:11:50 +00:00
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
}
}