Merge pull request #18 from winebarrel/wait_progress_close

Wait progress close

progress.go

@@ -20,6 +20,7 @@ type Progress struct {
 	noop        bool
 	prevDPLen   int
 	nDeadAgents int32
+	closed      chan struct{}
 }
 
 func NewProgress(w io.Writer, noop bool) *Progress {
@@ -36,6 +37,7 @@ func (progress *Progress) Start(ctx context.Context, rec *Recorder) {
 		return
 	}
 
+	progress.closed = make(chan struct{})
 	tk := time.NewTicker(InterimReportIntvl)
 
 	go func() {
@@ -50,7 +52,7 @@ func (progress *Progress) Start(ctx context.Context, rec *Recorder) {
 			}
 		}
 
-		progress.clear()
+		close(progress.closed)
 	}()
 }
 
@@ -86,11 +88,12 @@ func (progress *Progress) report(rec *Recorder) {
 	fmt.Fprintf(progress.w, "\r%-*s", width, line)
 }
 
-func (progress *Progress) clear() {
+func (progress *Progress) Close() {
 	if progress.noop {
 		return
 	}
 
+	<-progress.closed
 	width, _, err := term.GetSize(0)
 
 	if err != nil {

recorder.go

@@ -29,13 +29,14 @@ func NewRecorder(id string, options *Options) *Recorder {
 		ID:         id,
 		DataPoints: []DataPoint{},
 		ch:         make(chan []DataPoint, options.Nagents*3),
-		closed:     make(chan struct{}),
 	}
 
 	return rec
 }
 
 func (rec *Recorder) Start() {
+	rec.closed = make(chan struct{})
+
 	push := func(dps []DataPoint) {
 		rec.Lock()
 		defer rec.Unlock()

task.go

@@ -119,7 +119,8 @@ func (task *Task) Run() (*Report, error) {
 	progress.Start(ctx, rec)
 	err = eg.Wait()
 	cancel()
-	rec.Close() // wait for buffer flush
+	progress.Close() // wait for ticker to stop
+	rec.Close()      // wait for buffer flush
 
 	if err != nil && !errors.Is(err, context.DeadlineExceeded) {
 		return nil, err