From 1359faf968e5f1d37541b2872f67bdec60012093 Mon Sep 17 00:00:00 2001
From: Paul Rogers <129207811+paul1r@users.noreply.github.com>
Date: Tue, 5 Mar 2024 17:07:17 -0500
Subject: [PATCH] test: Fix for a few data races (#12129)

---
 pkg/querier/ingester_querier_test.go             | 15 +++++++++------
 pkg/querier/queryrange/split_by_interval_test.go |  6 +++++-
 2 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/pkg/querier/ingester_querier_test.go b/pkg/querier/ingester_querier_test.go
index a5cfd9a54dd8d..2eee6dc0ae072 100644
--- a/pkg/querier/ingester_querier_test.go
+++ b/pkg/querier/ingester_querier_test.go
@@ -3,6 +3,7 @@ package querier
 import (
 	"context"
 	"errors"
+	"go.uber.org/atomic"
 	"sync"
 	"testing"
 	"time"
@@ -65,6 +66,8 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 		},
 	}
 
+	var cnt atomic.Int32
+
 	for testName, testData := range tests {
 		for _, retErr := range []bool{true, false} {
 			testName, testData, retErr := testName, testData, retErr
@@ -75,9 +78,9 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 			}
 
 			t.Run(testName, func(t *testing.T) {
-				cnt := 0
 				wg := sync.WaitGroup{}
 				wait := make(chan struct{})
+				cnt.Store(0)
 
 				runFn := func(args mock.Arguments) {
 					wg.Done()
@@ -88,7 +91,7 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 						// ctx should be cancelled after the first two replicas return
 						require.ErrorIs(t, ctx.Err(), context.Canceled)
 					case <-wait:
-						cnt++
+						cnt.Add(1)
 					case <-time.After(time.Second):
 						t.Error("timed out waiting for ctx cancellation")
 					}
@@ -121,7 +124,7 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 
 				err = testData.testFn(ingesterQuerier)
 				ingesterClient.AssertNumberOfCalls(t, testData.method, 3)
-				require.Equal(t, 2, cnt)
+				require.Equal(t, int32(2), cnt.Load())
 				if retErr {
 					require.ErrorContains(t, err, testData.method+" failed")
 				} else {
@@ -176,7 +179,7 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 			}
 
 			t.Run(testName, func(t *testing.T) {
-				cnt := 0
+				cnt.Store(0)
 				wg := sync.WaitGroup{}
 				wait := make(chan struct{})
 
@@ -189,7 +192,7 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 						// should not be cancelled by the tracker
 						require.NoError(t, ctx.Err())
 					case <-wait:
-						cnt++
+						cnt.Add(1)
 					case <-time.After(time.Second):
 					}
 				}
@@ -221,7 +224,7 @@ func TestIngesterQuerier_earlyExitOnQuorum(t *testing.T) {
 
 				err = testData.testFn(ingesterQuerier)
 				ingesterClient.AssertNumberOfCalls(t, testData.method, 3)
-				require.Equal(t, 2, cnt)
+				require.Equal(t, int32(2), cnt.Load())
 				if retErr {
 					require.ErrorContains(t, err, testData.method+" failed")
 				} else {
diff --git a/pkg/querier/queryrange/split_by_interval_test.go b/pkg/querier/queryrange/split_by_interval_test.go
index c559bad17be1d..6c6d66e3bb662 100644
--- a/pkg/querier/queryrange/split_by_interval_test.go
+++ b/pkg/querier/queryrange/split_by_interval_test.go
@@ -1852,7 +1852,11 @@ func Test_ExitEarly(t *testing.T) {
 
 	res, err := split.Do(ctx, req)
 
-	require.Equal(t, int(req.Limit), callCt)
+	mtx.Lock()
+	count := callCt
+	mtx.Unlock()
+
+	require.Equal(t, int(req.Limit), count)
 	require.NoError(t, err)
 	require.Equal(t, expected, res)
 }