diff --git a/view.go b/view.go index f3de34e..893cd59 100644 --- a/view.go +++ b/view.go @@ -20,6 +20,7 @@ type View[K comparable] struct { wg *sync.WaitGroup reportCh chan report[K] + doneCh chan struct{} } func New[K comparable]() *View[K] { @@ -33,6 +34,7 @@ func NewForWriter[K comparable](w io.Writer) *View[K] { taskMap: map[K]struct{}{}, reportCh: make(chan report[K], 100), + doneCh: make(chan struct{}, 1), lastReports: map[K]report[K]{}, wg: &sync.WaitGroup{}, } @@ -63,6 +65,8 @@ func (v *View[K]) consumeReports(printReports bool) { v.print() } } + + v.doneCh <- struct{}{} } func (v *View[K]) SortFunc(f func(i, j K) bool) { @@ -96,3 +100,7 @@ func (v *View[K]) Done(task K) { v.wg.Done() } + +func (v *View[K]) WaitForAllTasks() { + <-v.doneCh +}