Wait for the droppedIterations push in per-vu-iterations

part of #1357
grafana · Jul 20, 2020 · 5cb5ce8 · 5cb5ce8
1 parent 9c8de85
commit 5cb5ce8
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/lib/executor/per_vu_iterations.go b/lib/executor/per_vu_iterations.go
@@ -191,6 +191,8 @@ func (pvi PerVUIterations) Run(parentCtx context.Context, out chan<- stats.Sampl
 	pvi.progress.Modify(pb.WithProgress(progressFn))
 	go trackProgress(parentCtx, maxDurationCtx, regDurationCtx, pvi, progressFn)
 
+	handleVUsWG := &sync.WaitGroup{}
+	defer handleVUsWG.Wait()
 	// Actually schedule the VUs and iterations...
 	activeVUs := &sync.WaitGroup{}
 	defer activeVUs.Wait()
@@ -203,7 +205,9 @@ func (pvi PerVUIterations) Run(parentCtx context.Context, out chan<- stats.Sampl
 			pvi.executionState.ReturnVU(u, true)
 			activeVUs.Done()
 		})
+
 	handleVU := func(initVU lib.InitializedVU) {
+		defer handleVUsWG.Done()
 		ctx, cancel := context.WithCancel(maxDurationCtx)
 		defer cancel()
 
@@ -236,6 +240,7 @@ func (pvi PerVUIterations) Run(parentCtx context.Context, out chan<- stats.Sampl
 			return err
 		}
 		activeVUs.Add(1)
+		handleVUsWG.Add(1)
 		go handleVU(initializedVU)
 	}
 

diff --git a/lib/executor/per_vu_iterations_test.go b/lib/executor/per_vu_iterations_test.go
@@ -63,7 +63,8 @@ func TestPerVUIterationsRun(t *testing.T) {
 		}),
 	)
 	defer cancel()
-	err = executor.Run(ctx, nil)
+	engineOut := make(chan stats.SampleContainer, 1000)
+	err = executor.Run(ctx, engineOut)
 	require.NoError(t, err)
 
 	var totalIters uint64
@@ -105,7 +106,8 @@ func TestPerVUIterationsRunVariableVU(t *testing.T) {
 		}),
 	)
 	defer cancel()
-	err = executor.Run(ctx, nil)
+	engineOut := make(chan stats.SampleContainer, 1000)
+	err = executor.Run(ctx, engineOut)
 	require.NoError(t, err)
 
 	val, ok := result.Load(slowVUID)