From 603fb384780112bdf9101b5277a21a3642a553cc Mon Sep 17 00:00:00 2001
From: Yi Jin <96499497+jnyi@users.noreply.github.com>
Date: Thu, 4 Apr 2024 23:54:17 -0700
Subject: [PATCH] fix query_test when --race enabled (#7258)

Signed-off-by: Yi Jin <yi.jin@databricks.com>
---
 pkg/query/query_test.go | 35 ++++++++++++++++++++++++++++++-----
 1 file changed, 30 insertions(+), 5 deletions(-)

diff --git a/pkg/query/query_test.go b/pkg/query/query_test.go
index 79d887c42b..b490156e82 100644
--- a/pkg/query/query_test.go
+++ b/pkg/query/query_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"sync"
 	"testing"
 	"time"
 
@@ -25,6 +26,31 @@ func TestMain(m *testing.M) {
 	custom.TolerantVerifyLeakMain(m)
 }
 
+type safeClients struct {
+	sync.RWMutex
+	clients []store.Client
+}
+
+func (sc *safeClients) get() []store.Client {
+	sc.RLock()
+	defer sc.RUnlock()
+	ret := make([]store.Client, len(sc.clients))
+	copy(ret, sc.clients)
+	return ret
+}
+
+func (sc *safeClients) reset() {
+	sc.Lock()
+	defer sc.Unlock()
+	sc.clients = sc.clients[:0]
+}
+
+func (sc *safeClients) append(c store.Client) {
+	sc.Lock()
+	defer sc.Unlock()
+	sc.clients = append(sc.clients, c)
+}
+
 func TestQuerier_Proxy(t *testing.T) {
 	files, err := filepath.Glob("testdata/promql/**/*.test")
 	testutil.Ok(t, err)
@@ -32,24 +58,23 @@ func TestQuerier_Proxy(t *testing.T) {
 
 	logger := log.NewLogfmtLogger(os.Stderr)
 	t.Run("proxy", func(t *testing.T) {
-		var clients []store.Client
+		var sc safeClients
 		q := NewQueryableCreator(
 			logger,
 			nil,
-			store.NewProxyStore(logger, nil, func() []store.Client { return clients },
+			store.NewProxyStore(logger, nil, func() []store.Client { return sc.get() },
 				component.Debug, nil, 5*time.Minute, store.EagerRetrieval),
 			1000000,
 			5*time.Minute,
 		)
 
 		createQueryableFn := func(stores []*testStore) storage.Queryable {
-			clients = clients[:0]
+			sc.reset()
 			for i, st := range stores {
 				m, err := storepb.PromMatchersToMatchers(st.matchers...)
 				testutil.Ok(t, err)
-
 				// TODO(bwplotka): Parse external labels.
-				clients = append(clients, &storetestutil.TestClient{
+				sc.append(&storetestutil.TestClient{
 					Name:        fmt.Sprintf("store number %v", i),
 					StoreClient: storepb.ServerAsClient(selectedStore(store.NewTSDBStore(logger, st.storage.DB, component.Debug, nil), m, st.mint, st.maxt)),
 					MinTime:     st.mint,