package mpb
import (
"context"
"errors"
"io"
"log"
"os"
"sync"
"time"
"unicode/utf8"
"github.com/vbauerster/mpb/cwriter"
)
var logger = log.New(os.Stderr, "mpb: ", log.LstdFlags|log.Lshortfile)
// ErrCallAfterStop thrown by panic, if Progress methods like AddBar() are called
// after Stop() has been called
var ErrCallAfterStop = errors.New("method call on stopped Progress instance")
type (
// BeforeRender is a func, which gets called before render process
BeforeRender func([]*Bar)
barOpType uint
operation struct {
kind barOpType
bar *Bar
result chan bool
}
indexedBarBuffer struct {
index int
buf []byte
}
indexedBar struct {
index int
termWidth int
bar *Bar
}
)
const (
barAdd barOpType = iota
barRemove
)
// default RefreshRate
const rr = 100
// Progress represents the container that renders Progress bars
type Progress struct {
// Context for canceling bars rendering
ctx context.Context
// WaitGroup for internal rendering sync
wg *sync.WaitGroup
out io.Writer
width int
format string
operationCh chan *operation
rrChangeReqCh chan time.Duration
outChangeReqCh chan io.Writer
barCountReqCh chan chan int
brCh chan BeforeRender
done chan struct{}
}
// New creates new Progress instance, which will orchestrate bars rendering
// process. It acceepts context.Context, for cancellation.
// If you don't plan to cancel, it is safe to feed with nil
func New(ctx context.Context) *Progress {
if ctx == nil {
ctx = context.Background()
}
p := &Progress{
width: 70,
operationCh: make(chan *operation),
rrChangeReqCh: make(chan time.Duration),
outChangeReqCh: make(chan io.Writer),
barCountReqCh: make(chan chan int),
brCh: make(chan BeforeRender),
done: make(chan struct{}),
wg: new(sync.WaitGroup),
ctx: ctx,
}
go p.server(cwriter.New(os.Stdout), time.NewTicker(rr*time.Millisecond))
return p
}
// SetWidth sets the width for all underlying bars
func (p *Progress) SetWidth(n int) *Progress {
if n <= 0 {
return p
}
p.width = n
return p
}
// SetOut sets underlying writer of progress. Default is os.Stdout
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) SetOut(w io.Writer) *Progress {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
if w == nil {
return p
}
p.outChangeReqCh <- w
return p
}
// RefreshRate overrides default (30ms) refreshRate value
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) RefreshRate(d time.Duration) *Progress {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
p.rrChangeReqCh <- d
return p
}
func (p *Progress) BeforeRenderFunc(f BeforeRender) *Progress {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
p.brCh <- f
return p
}
// AddBar creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) AddBar(total int64) *Bar {
return p.AddBarWithID(0, total)
}
// AddBarWithID creates a new progress bar and adds to the container
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) AddBarWithID(id, total int64) *Bar {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
result := make(chan bool)
bar := newBar(p.ctx, p.wg, id, total, p.width, p.format)
p.operationCh <- &operation{barAdd, bar, result}
if <-result {
p.wg.Add(1)
}
return bar
}
// RemoveBar removes bar at any time
// pancis, if called on stopped Progress instance, i.e after Stop()
func (p *Progress) RemoveBar(b *Bar) bool {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
result := make(chan bool)
p.operationCh <- &operation{barRemove, b, result}
return <-result
}
// BarCount returns bars count in the container.
// Pancis if called on stopped Progress instance, i.e after Stop()
func (p *Progress) BarCount() int {
if IsClosed(p.done) {
panic(ErrCallAfterStop)
}
respCh := make(chan int)
p.barCountReqCh <- respCh
return <-respCh
}
// Format sets custom format for underlying bar(s).
// The format string, must consist of ASCII characters only.
// The default one is "[=>-]"
func (p *Progress) Format(format string) *Progress {
if utf8.RuneCountInString(format) != 5 {
return p
}
p.format = format
return p
}
// Stop waits for bars to finish rendering and stops the rendering goroutine
func (p *Progress) Stop() {
p.wg.Wait()
if IsClosed(p.done) {
return
}
close(p.operationCh)
}
// server monitors underlying channels and renders any progress bars
func (p *Progress) server(cw *cwriter.Writer, t *time.Ticker) {
defer func() {
t.Stop()
close(p.done)
}()
const numDrawers = 3
bars := make([]*Bar, 0, 4)
var beforeRender BeforeRender
var wg sync.WaitGroup
recoverIfPanic := func() {
if e := recover(); e != nil {
logger.Printf("unexpected panic: %+v\n", e)
}
wg.Done()
}
for {
select {
case w := <-p.outChangeReqCh:
cw.Flush()
cw = cwriter.New(w)
case op, ok := <-p.operationCh:
if !ok {
return
}
switch op.kind {
case barAdd:
bars = append(bars, op.bar)
op.result <- true
case barRemove:
var ok bool
for i, b := range bars {
if b == op.bar {
bars = append(bars[:i], bars[i+1:]...)
ok = true
b.remove()
break
}
}
op.result <- ok
}
case respCh := <-p.barCountReqCh:
respCh <- len(bars)
case beforeRender = <-p.brCh:
case <-t.C:
if beforeRender != nil {
beforeRender(bars)
}
width, _, _ := cwriter.GetTermSize()
ibars := iBarsGen(bars, width)
c := make(chan indexedBarBuffer)
wg.Add(numDrawers)
for i := 0; i < numDrawers; i++ {
go func() {
defer recoverIfPanic()
drawer(ibars, c)
}()
}
go func() {
wg.Wait()
close(c)
}()
m := make(map[int][]byte, len(bars))
for r := range c {
m[r.index] = r.buf
}
for i := 0; i < len(bars); i++ {
cw.Write(m[i])
}
cw.Flush()
for _, b := range bars {
b.flushed()
}
case d := <-p.rrChangeReqCh:
t.Stop()
t = time.NewTicker(d)
case <-p.ctx.Done():
return
}
}
}
func drawer(ibars <-chan indexedBar, c chan<- indexedBarBuffer) {
for b := range ibars {
buf := b.bar.bytes(b.termWidth)
buf = append(buf, '\n')
c <- indexedBarBuffer{b.index, buf}
}
}
func iBarsGen(bars []*Bar, width int) <-chan indexedBar {
ibars := make(chan indexedBar)
go func() {
defer close(ibars)
for i, b := range bars {
ibars <- indexedBar{i, width, b}
}
}()
return ibars
}
// IsClosed check if ch closed
// caution see: http://www.tapirgames.com/blog/golang-channel-closing
func IsClosed(ch <-chan struct{}) bool {
select {
case <-ch:
return true
default:
return false
}
}