diff --git a/decorators_test.go b/decorators_test.go index 64b1c99..144ae4a 100644 --- a/decorators_test.go +++ b/decorators_test.go @@ -183,30 +183,26 @@ } for _, columnCase := range testCases { - mpb.SyncWidth(toSyncMatrix(columnCase)) - numBars := len(columnCase) - gott := make([]chan string, numBars) - wg := new(sync.WaitGroup) - wg.Add(numBars) - for i, step := range columnCase { + var wg sync.WaitGroup + mpb.SyncWidth(&wg, toSyncMatrix(columnCase)) + var results []chan string + for _, step := range columnCase { step := step - ch := make(chan string, 1) + ch := make(chan string) go func() { - defer wg.Done() ch <- step.decorator.Decor(step.stat) }() - gott[i] = ch + results = append(results, ch) + } + + for i, ch := range results { + res := <-ch + want := columnCase[i].want + if res != want { + t.Errorf("Want: %q, Got: %q\n", want, res) + } } wg.Wait() - - for i, ch := range gott { - got := <-ch - want := columnCase[i].want - if got != want { - t.Errorf("Want: %q, Got: %q\n", want, got) - } - } - } } diff --git a/export_test.go b/export_test.go index 243429b..7f5cb84 100644 --- a/export_test.go +++ b/export_test.go @@ -2,4 +2,3 @@ // make syncWidth func public in test var SyncWidth = syncWidth -var MaxWidthDistributor = maxWidthDistributor diff --git a/progress.go b/progress.go index 24427c0..f70e2aa 100644 --- a/progress.go +++ b/progress.go @@ -297,12 +297,13 @@ } func (s *pState) render(cw *cwriter.Writer) error { + var wg sync.WaitGroup if s.heapUpdated { s.updateSyncMatrix() s.heapUpdated = false } - syncWidth(s.pMatrix) - syncWidth(s.aMatrix) + syncWidth(&wg, s.pMatrix) + syncWidth(&wg, s.aMatrix) width, height, err := cw.GetTermSize() if err != nil { @@ -314,11 +315,12 @@ go bar.render(width) } - return s.flush(cw, height) -} - -func (s *pState) flush(cw *cwriter.Writer, height int) error { - var wg sync.WaitGroup + err = s.flush(&wg, cw, height) + wg.Wait() + return err +} + +func (s *pState) flush(wg *sync.WaitGroup, cw *cwriter.Writer, height int) error { var popCount int for s.bHeap.Len() > 0 { @@ -373,12 +375,14 @@ } case 1: heap.Push(&s.bHeap, s.pool[0]) + s.pool = s.pool[:0] default: wg.Add(1) go func() { for _, b := range s.pool { heap.Push(&s.bHeap, b) } + s.pool = s.pool[:0] wg.Done() }() } @@ -386,15 +390,13 @@ for i := len(s.rows) - 1; i >= 0; i-- { _, err := cw.ReadFrom(s.rows[i]) if err != nil { - wg.Wait() + s.rows = s.rows[:0] return err } } err := cw.Flush(len(s.rows) - popCount) - wg.Wait() s.rows = s.rows[:0] - s.pool = s.pool[:0] return err } @@ -483,13 +485,14 @@ return bs } -func syncWidth(matrix map[int][]chan int) { +func syncWidth(wg *sync.WaitGroup, matrix map[int][]chan int) { for _, column := range matrix { - go maxWidthDistributor(column) - } -} - -func maxWidthDistributor(column []chan int) { + wg.Add(1) + go maxWidthDistributor(wg, column) + } +} + +func maxWidthDistributor(wg *sync.WaitGroup, column []chan int) { var maxWidth int for _, ch := range column { if w := <-ch; w > maxWidth { @@ -499,4 +502,5 @@ for _, ch := range column { ch <- maxWidth } -} + wg.Done() +} diff --git a/progress_test.go b/progress_test.go index c85465a..7e5db19 100644 --- a/progress_test.go +++ b/progress_test.go @@ -112,70 +112,6 @@ } } -func TestMaxWidthDistributor(t *testing.T) { - makeWrapper := func(f func([]chan int), start, end chan struct{}) func([]chan int) { - return func(column []chan int) { - start <- struct{}{} - f(column) - <-end - } - } - - ready := make(chan struct{}) - start := make(chan struct{}) - end := make(chan struct{}) - // mpb.MaxWidthDistributor shouldn't stuck in the middle while removing or aborting a bar - mpb.MaxWidthDistributor = makeWrapper(mpb.MaxWidthDistributor, start, end) - - total := 100 - numBars := 6 - p := mpb.New(mpb.WithOutput(io.Discard)) - for i := 0; i < numBars; i++ { - bar := p.AddBar(int64(total), - mpb.BarOptional(mpb.BarRemoveOnComplete(), i == 0), - mpb.PrependDecorators(decor.EwmaETA(decor.ET_STYLE_GO, 60, decor.WCSyncSpace)), - ) - go func() { - <-ready - for i := 0; i < total; i++ { - start := time.Now() - if id := bar.ID(); id > 1 && i >= 32 { - if id&1 == 1 { - bar.Abort(true) - } else { - bar.Abort(false) - } - } - time.Sleep(randomDuration(100 * time.Millisecond)) - bar.EwmaIncrInt64(rand.Int63n(5)+1, time.Since(start)) - } - }() - } - - go func() { - <-ready - p.Wait() - close(start) - }() - - res := t.Run("maxWidthDistributor", func(t *testing.T) { - close(ready) - for v := range start { - timer := time.NewTimer(100 * time.Millisecond) - select { - case end <- v: - timer.Stop() - case <-timer.C: - t.FailNow() - } - } - }) - - if !res { - t.Error("maxWidthDistributor stuck in the middle") - } -} - func TestProgressShutdownsWithErrFiller(t *testing.T) { var debug bytes.Buffer shutdown := make(chan struct{})