From 5ecbdb2863bb72af828437f756bd054a80db31b8 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Fri, 2 Dec 2022 13:34:12 -0500
Subject: [PATCH] e2e: add defer to ensure waitgroup is finished. (#9822)
 (#9824)

* defer wait group completion

* generate correct number of tx

(cherry picked from commit 654e5652e4ce3c521895a97909712c0a5509e8fa)

Co-authored-by: William Banfield <4561443+williambanfield@users.noreply.github.com>
---
 test/e2e/runner/load.go | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/test/e2e/runner/load.go b/test/e2e/runner/load.go
index 1989886b0..b4478bf6d 100644
--- a/test/e2e/runner/load.go
+++ b/test/e2e/runner/load.go
@@ -92,10 +92,12 @@ func loadGenerate(ctx context.Context, txCh chan<- types.Tx, testnet *e2e.Testne
 // is canceled.
 func createTxBatch(ctx context.Context, txCh chan<- types.Tx, testnet *e2e.Testnet, id []byte) {
 	wg := &sync.WaitGroup{}
+	genCh := make(chan struct{})
 	for i := 0; i < workerPoolSize; i++ {
 		wg.Add(1)
 		go func() {
-			for i := 0; i < testnet.LoadTxBatchSize; i++ {
+			defer wg.Done()
+			for range genCh {
 				tx, err := payload.NewBytes(&payload.Payload{
 					Id:          id,
 					Size:        uint64(testnet.LoadTxSizeBytes),
@@ -112,9 +114,16 @@ func createTxBatch(ctx context.Context, txCh chan<- types.Tx, testnet *e2e.Testn
 					return
 				}
 			}
-			wg.Done()
 		}()
 	}
+	for i := 0; i < testnet.LoadTxBatchSize; i++ {
+		select {
+		case genCh <- struct{}{}:
+		case <-ctx.Done():
+			break
+		}
+	}
+	close(genCh)
 	wg.Wait()
 }