diff --git a/view.go b/view.go index f5b645e..ad28135 100644 --- a/view.go +++ b/view.go @@ -19,6 +19,7 @@ type View[K comparable] struct { wg *sync.WaitGroup reportCh chan report[K] + doneCh chan struct{} } func New[K comparable]() *View[K] { @@ -32,6 +33,7 @@ func NewForWriter[K comparable](w io.Writer) *View[K] { taskMap: map[K]struct{}{}, reportCh: make(chan report[K], 100), + doneCh: make(chan struct{}), lastReports: map[K]report[K]{}, wg: &sync.WaitGroup{}, } @@ -48,6 +50,8 @@ func (v *View[K]) Show() { v.Add(status.ID) v.print() } + + v.doneCh <- struct{}{} } func (v *View[K]) SortFunc(f func(i, j K) bool) { @@ -79,5 +83,5 @@ func (v *View[K]) Done(task K) { } func (v *View[K]) WaitForAllTasks() { - v.wg.Wait() + <-v.doneCh }