Refactor progress bar life cycle

2 files changed, 36 insertions, 21 deletions

diff --git a/llm/perplexity.go b/llm/perplexity.go
index ab31889..5d9056e 100644
--- a/llm/perplexity.go
+++ b/llm/perplexity.go
@@ -1,7 +1,6 @@
 package llm
 
 import (
-	"context"
 	"fmt"
 	"sync"
 	"time"
@@ -47,16 +46,14 @@ func (e *Evaluator[R]) Run(data dataset.Reader, window, stride int) error {
 		}()
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-
-	defer cancel()
-
 	pb := tui.NewProgressBar("Perplexity", 20, 0, time.Now())
 
-	go pb.Watch(ctx, 1*time.Second, func() int {
+	pb.Start(1*time.Second, func() int {
 		return int(e.completed.Load())
 	})
 
+	defer pb.Close()
+
 	n := 0
 
 	for d := range data.Texts("text") {
diff --git a/tui/progressbar.go b/tui/progressbar.go
index 7cd0b44..0cd7a29 100644
--- a/tui/progressbar.go
+++ b/tui/progressbar.go
@@ -14,6 +14,8 @@ type ProgressBar struct {
 	total     atomic.Int64
 	completed atomic.Int64
 	start     time.Time
+	done      chan struct{}
+	cancel    context.CancelFunc
 }
 
 func NewProgressBar(title string, width, total int, start time.Time) *ProgressBar {
@@ -21,6 +23,7 @@ func NewProgressBar(title string, width, total int, start time.Time) *ProgressBa
 		title: title,
 		width: width,
 		start: start,
+		done:  make(chan struct{}),
 	}
 
 	pb.total.Store(int64(total))
@@ -98,32 +101,47 @@ func (pb *ProgressBar) String() string {
 	)
 }
 
-func (pb *ProgressBar) Watch(
-	ctx context.Context,
+func (pb *ProgressBar) Start(
 	interval time.Duration,
 	callback func() int,
 ) {
-	ticker := time.NewTicker(interval)
+	ctx, cancel := context.WithCancel(context.Background())
 
-	defer ticker.Stop()
+	pb.cancel = cancel
 
-	for {
-		select {
-		case <-ctx.Done():
-			n := callback()
+	go func() {
+		ticker := time.NewTicker(interval)
 
-			pb.SetCompleted(n)
+		defer ticker.Stop()
 
-			pb.Finish()
+		defer close(pb.done)
 
-			return
+		for {
+			select {
+			case <-ctx.Done():
+				n := callback()
 
-		case <-ticker.C:
-			n := callback()
+				pb.SetCompleted(n)
 
-			pb.SetCompleted(n)
+				pb.Finish()
 
-			pb.Print()
+				return
+
+			case <-ticker.C:
+				n := callback()
+
+				pb.SetCompleted(n)
+
+				pb.Print()
+			}
 		}
+	}()
+}
+
+func (pb *ProgressBar) Close() {
+	if pb.cancel != nil {
+		pb.cancel()
+
+		<-pb.done
 	}
 }