Merge pull request #21 from henvic/fix-race-condition-on-stop

Using time.Ticker for running stoppable listener. Fix edge case where flush happens after the listener is stopped.

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

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.")
+	}
+}