From 3cb1984024916b8b80ad1bb72518c202dc9966d5 Mon Sep 17 00:00:00 2001
From: milarin <milaring@proton.me>
Date: Mon, 6 Mar 2023 14:09:42 +0100
Subject: [PATCH] fixed implementation of WaitForAllTasks

---
 view.go | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

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
 }