diff --git a/progress.go b/progress.go index a93838d..759a1a1 100644 --- a/progress.go +++ b/progress.go @@ -33,9 +33,10 @@ type Progress struct { // RefreshInterval in the time duration to wait for refreshing the output RefreshInterval time.Duration - lw *uilive.Writer - stopChan chan struct{} - mtx *sync.RWMutex + lw *uilive.Writer + ticker *time.Ticker + tdone chan bool + mtx *sync.RWMutex } // New returns a new progress bar with defaults @@ -46,9 +47,8 @@ func New() *Progress { Bars: make([]*Bar, 0), RefreshInterval: RefreshInterval, - lw: uilive.New(), - stopChan: make(chan struct{}), - mtx: &sync.RWMutex{}, + lw: uilive.New(), + mtx: &sync.RWMutex{}, } } @@ -85,35 +85,55 @@ func (p *Progress) AddBar(total int) *Bar { // Listen listens for updates and renders the progress bars func (p *Progress) Listen() { + var tickChan = p.ticker.C p.lw.Out = p.Out + for { select { - case <-p.stopChan: - return - default: - time.Sleep(p.RefreshInterval) + case <-tickChan: p.mtx.RLock() - for _, bar := range p.Bars { - fmt.Fprintln(p.lw, bar.String()) + + if p.ticker != nil { + p.print() + p.lw.Flush() } - p.lw.Flush() + p.mtx.RUnlock() + case <-p.tdone: + if p.ticker != nil { + p.ticker.Stop() + p.ticker = nil + return + } } } } +func (p *Progress) print() { + for _, bar := range p.Bars { + fmt.Fprintln(p.lw, bar.String()) + } +} + // Start starts the rendering the progress of progress bars. It listens for updates using `bar.Set(n)` and new bars when added using `AddBar` func (p *Progress) Start() { - if p.stopChan == nil { - p.stopChan = make(chan struct{}) + p.mtx.Lock() + if p.ticker == nil { + p.ticker = time.NewTicker(RefreshInterval) + p.tdone = make(chan bool, 1) } + p.mtx.Unlock() + go p.Listen() } // Stop stops listening func (p *Progress) Stop() { - close(p.stopChan) - p.stopChan = nil + p.mtx.Lock() + close(p.tdone) + p.print() + p.lw.Flush() + p.mtx.Unlock() } // Bypass returns a writer which allows non-buffered data to be written to the underlying output diff --git a/progress_test.go b/progress_test.go new file mode 100644 index 0000000..d6a1b42 --- /dev/null +++ b/progress_test.go @@ -0,0 +1,43 @@ +package uiprogress + +import ( + "bytes" + "fmt" + "strings" + "sync" + "testing" + "time" +) + +func TestStoppingPrintout(t *testing.T) { + progress := New() + progress.RefreshInterval = time.Millisecond * 10 + var buffer = &bytes.Buffer{} + progress.Out = buffer + bar := progress.AddBar(100) + progress.Start() + + var wg sync.WaitGroup + + wg.Add(1) + + go func() { + for i := 0; i <= 80; i = i + 10 { + bar.Set(i) + time.Sleep(time.Millisecond * 5) + } + + wg.Done() + }() + + wg.Wait() + + progress.Stop() + fmt.Fprintf(buffer, "foo") + + var wantSuffix = "[======================================================>-------------]\nfoo" + + if !strings.HasSuffix(buffer.String(), wantSuffix) { + t.Errorf("Content that should be printed after stop not appearing on buffer.") + } +}