diff --git a/view.go b/view.go index 07f844c..2c4662e 100644 --- a/view.go +++ b/view.go @@ -1,6 +1,7 @@ package statview import ( + "fmt" "io" "os" "sort" @@ -19,7 +20,6 @@ type View[K comparable] struct { wg *sync.WaitGroup reportCh chan report[K] - doneCh chan struct{} } func New[K comparable]() *View[K] { @@ -33,22 +33,33 @@ 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{}, } } func (v *View[K]) Show() { - go v.WaitForAllTasks() + v.consumeReports(true) +} + +func (v *View[K]) Hide() { + v.consumeReports(false) +} + +func (v *View[K]) consumeReports(printReports bool) { + go func() { + defer close(v.reportCh) + v.wg.Wait() + }() for status := range v.reportCh { v.lastReports[status.ID] = status v.Add(status.ID) - v.print() - } - v.doneCh <- struct{}{} + if printReports { + v.print() + } + } } func (v *View[K]) SortFunc(f func(i, j K) bool) { @@ -76,14 +87,9 @@ func (v *View[K]) Report(task K, status string) { } func (v *View[K]) Done(task K) { + if _, ok := v.taskMap[task]; !ok { + panic(fmt.Sprintf("unreported task marked as done: %v", v)) + } + v.wg.Done() } - -func (v *View[K]) WaitForAllTasks() { - go func() { - defer close(v.reportCh) - v.wg.Wait() - }() - - <-v.doneCh -}