chore(tests/stress): refactor compareBlocksByNumber (#2370)

- Takes in context
- Retry as long as context is not canceled
- Assert errors and common values

Author: qdm12

tests/stress/empty.go

@@ -10,7 +10,5 @@ import (
 var (
 	errFinalizedBlockMismatch = errors.New("node finalised head hashes don't match")
 	errNoFinalizedBlock       = errors.New("did not finalise block for round")
-	errNoBlockAtNumber        = errors.New("no blocks found for given number")
-	errBlocksAtNumberMismatch = errors.New("different blocks found for given number")
 	errChainHeadMismatch      = errors.New("node chain head hashes don't match")
 )

tests/stress/errors.go

@@ -4,9 +4,9 @@
 package stress
 
 import (
+	"context"
 	"errors"
 	"fmt"
-	"sync"
 	"testing"
 	"time"
 
@@ -65,74 +65,60 @@ func compareChainHeadsWithRetry(t *testing.T, nodes []*utils.Node) error {
 
 // compareBlocksByNumber calls getBlockByNumber for each node in the array
 // it returns a map of block hashes to node key names, and an error if the hashes don't all match
-func compareBlocksByNumber(t *testing.T, nodes []*utils.Node, num string) (map[common.Hash][]string, error) {
-	hashes := make(map[common.Hash][]string)
-	var errs []error
-	var mapMu sync.Mutex
-	var wg sync.WaitGroup
-	wg.Add(len(nodes))
+func compareBlocksByNumber(ctx context.Context, t *testing.T, nodes []*utils.Node,
+	num string) (hashToKeys map[common.Hash][]string) {
+	type resultContainer struct {
+		hash    common.Hash
+		nodeKey string
+		err     error
+	}
+	results := make(chan resultContainer)
 
 	for _, node := range nodes {
 		go func(node *utils.Node) {
-			logger.Debugf("calling chain_getBlockHash for node index %d", node.Idx)
-			hash, err := utils.GetBlockHash(t, node, num)
-			mapMu.Lock()
-			defer func() {
-				mapMu.Unlock()
-				wg.Done()
-			}()
-			if err != nil {
-				errs = append(errs, err)
-				return
+			result := resultContainer{
+				nodeKey: node.Key,
 			}
-			logger.Debugf("got hash %s from node with key %s", hash, node.Key)
 
-			hashes[hash] = append(hashes[hash], node.Key)
-		}(node)
-	}
+			for { // retry until context gets canceled
+				result.hash, result.err = utils.GetBlockHash(t, node, num)
 
-	wg.Wait()
-
-	var err error
-	if len(errs) != 0 {
-		err = fmt.Errorf("%v", errs)
-	}
+				if err := ctx.Err(); err != nil {
+					result.err = err
+					break
+				}
 
-	if len(hashes) == 0 {
-		err = errNoBlockAtNumber
-	}
+				if result.err == nil {
+					break
+				}
+			}
 
-	if len(hashes) > 1 {
-		err = errBlocksAtNumberMismatch
+			results <- result
+		}(node)
 	}
 
-	return hashes, err
-}
-
-// compareBlocksByNumberWithRetry calls compareChainHeads, retrying up to maxRetries times if it errors.
-func compareBlocksByNumberWithRetry(t *testing.T, nodes []*utils.Node, num string) (map[common.Hash][]string, error) {
-	var hashes map[common.Hash][]string
 	var err error
+	hashToKeys = make(map[common.Hash][]string, len(nodes))
+	for range nodes {
+		result := <-results
+		if err != nil {
+			continue // one failed, we don't care anymore
+		}
 
-	timeout := time.After(120 * time.Second)
-doneBlockProduction:
-	for {
-		time.Sleep(time.Second)
-		select {
-		case <-timeout:
-			break doneBlockProduction
-		default:
-			hashes, err = compareBlocksByNumber(t, nodes, num)
-			if err == nil {
-				break doneBlockProduction
-			}
+		if result.err != nil {
+			err = result.err
+			continue
 		}
-	}
 
-	if err != nil {
-		err = fmt.Errorf("%w: hashes=%v", err, hashes)
+		hashToKeys[result.hash] = append(hashToKeys[result.hash], result.nodeKey)
 	}
-	return hashes, err
+
+	require.NoError(t, err)
+	require.Lenf(t, hashToKeys, 1,
+		"expected 1 block found for number %s but got %d block(s)",
+		num, len(hashToKeys))
+
+	return hashToKeys
 }
 
 // compareFinalizedHeads calls getFinalizedHeadByRound for each node in the array

tests/stress/helpers.go

@@ -4,6 +4,7 @@
 package stress
 
 import (
+	"context"
 	"fmt"
 	"math/big"
 	"math/rand"
@@ -111,14 +112,18 @@ func TestSync_SingleBlockProducer(t *testing.T) {
 	}()
 
 	numCmps := 10
+	ctx := context.Background()
+
 	for i := 0; i < numCmps; i++ {
 		time.Sleep(3 * time.Second)
 		t.Log("comparing...", i)
-		hashes, err := compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(i))
-		if len(hashes) > 1 || len(hashes) == 0 {
-			require.NoError(t, err, i)
-			continue
-		}
+
+		const comparisonTimeout = 5 * time.Second
+		compareCtx, cancel := context.WithTimeout(ctx, comparisonTimeout)
+
+		hashes := compareBlocksByNumber(compareCtx, t, nodes, strconv.Itoa(i))
+
+		cancel()
 
 		// there will only be one key in the mapping
 		for _, nodesWithHash := range hashes {
@@ -163,13 +168,20 @@ func TestSync_MultipleEpoch(t *testing.T) {
 	// Wait for epoch to pass
 	time.Sleep(time.Duration(uint64(slotDuration.Nanoseconds()) * epochLength))
 
+	ctx := context.Background()
+
 	// Just checking that everythings operating as expected
 	header := utils.GetChainHead(t, nodes[0])
 	currentHeight := header.Number
 	for i := uint(0); i < currentHeight; i++ {
 		t.Log("comparing...", i)
-		_, err = compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(int(i)))
-		require.NoError(t, err, i)
+
+		const compareTimeout = 5 * time.Second
+		compareCtx, cancel := context.WithTimeout(ctx, compareTimeout)
+
+		_ = compareBlocksByNumber(compareCtx, t, nodes, strconv.Itoa(int(i)))
+
+		cancel()
 	}
 }
 
@@ -197,11 +209,18 @@ func TestSync_SingleSyncingNode(t *testing.T) {
 		require.Len(t, errList, 0)
 	}()
 
+	ctx := context.Background()
+
 	numCmps := 100
 	for i := 0; i < numCmps; i++ {
 		t.Log("comparing...", i)
-		_, err = compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(i))
-		require.NoError(t, err, i)
+
+		const compareTimeout = 5 * time.Second
+		compareCtx, cancel := context.WithTimeout(ctx, compareTimeout)
+
+		_ = compareBlocksByNumber(compareCtx, t, nodes, strconv.Itoa(i))
+
+		cancel()
 	}
 }
 
@@ -277,8 +296,16 @@ func TestSync_Bench(t *testing.T) {
 
 	// assert block is correct
 	t.Log("comparing block...", numBlocks)
-	_, err = compareBlocksByNumberWithRetry(t, nodes, fmt.Sprint(numBlocks))
-	require.NoError(t, err, numBlocks)
+
+	ctx := context.Background()
+
+	const compareTimeout = 5 * time.Second
+	compareCtx, cancel := context.WithTimeout(ctx, compareTimeout)
+
+	_ = compareBlocksByNumber(compareCtx, t, nodes, fmt.Sprint(numBlocks))
+
+	cancel()
+
 	time.Sleep(time.Second * 3)
 }
 
@@ -328,11 +355,19 @@ func TestSync_Restart(t *testing.T) {
 		}
 	}()
 
+	ctx := context.Background()
+
 	numCmps := 12
 	for i := 0; i < numCmps; i++ {
 		t.Log("comparing...", i)
-		_, err = compareBlocksByNumberWithRetry(t, nodes, strconv.Itoa(i))
-		require.NoError(t, err, i)
+
+		const compareTimeout = 5 * time.Second
+		compareCtx, cancel := context.WithTimeout(ctx, compareTimeout)
+
+		_ = compareBlocksByNumber(compareCtx, t, nodes, strconv.Itoa(i))
+
+		cancel()
+
 		time.Sleep(time.Second * 5)
 	}
 	close(done)
@@ -476,8 +511,14 @@ func TestSync_SubmitExtrinsic(t *testing.T) {
 
 	require.True(t, included)
 
-	hashes, err := compareBlocksByNumberWithRetry(t, nodes, fmt.Sprint(extInBlock))
-	require.NoError(t, err, hashes)
+	ctx := context.Background()
+
+	const compareTimeout = 5 * time.Second
+	compareCtx, cancel := context.WithTimeout(ctx, compareTimeout)
+
+	_ = compareBlocksByNumber(compareCtx, t, nodes, fmt.Sprint(extInBlock))
+
+	cancel()
 }
 
 func Test_SubmitAndWatchExtrinsic(t *testing.T) {