From 1050a2388c178b7876532d8f3efbbd08e9651398 Mon Sep 17 00:00:00 2001
From: Yi Jin <yi.jin@databricks.com>
Date: Tue, 26 Nov 2024 19:07:20 -0800
Subject: [PATCH] fix data race in testing with some concurrent safe helper
 functions

Signed-off-by: Yi Jin <yi.jin@databricks.com>
---
 pkg/receive/multitsdb.go      | 21 +++++++++++++++++++++
 pkg/receive/multitsdb_test.go |  4 ++--
 pkg/receive/receive_test.go   | 12 ++++++------
 3 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/pkg/receive/multitsdb.go b/pkg/receive/multitsdb.go
index 39c947011e7..ca502ca67ee 100644
--- a/pkg/receive/multitsdb.go
+++ b/pkg/receive/multitsdb.go
@@ -136,6 +136,25 @@ func NewMultiTSDB(
 	return mt
 }
 
+// getTenant returns the tenant with the given tenantID.
+// @testing: This method is only for testing purposes.
+func (t *MultiTSDB) getTenant(tenantID string) *tenant {
+	t.mtx.RLock()
+	defer t.mtx.RUnlock()
+	tenant, _ := t.tenants[tenantID]
+	return tenant
+}
+
+// tsdbLocalClientsCopied returns a copy of tsdbClients.
+// @testing: This method is only for testing purposes.
+func (t *MultiTSDB) tsdbLocalClientsCopied() []store.Client {
+	t.mtx.Lock()
+	defer t.mtx.Unlock()
+	copied := make([]store.Client, len(t.tsdbClients))
+	copy(copied, t.tsdbClients)
+	return copied
+}
+
 func (t *MultiTSDB) updateTSDBClients() {
 	t.tsdbClients = t.tsdbClients[:0]
 	for _, tenant := range t.tenants {
@@ -626,12 +645,14 @@ func (t *MultiTSDB) RemoveLockFilesIfAny() error {
 	return merr.Err()
 }
 
+// TSDBLocalClients should be used as read-only
 func (t *MultiTSDB) TSDBLocalClients() []store.Client {
 	t.mtx.RLock()
 	defer t.mtx.RUnlock()
 	return t.tsdbClients
 }
 
+// TSDBExemplars should be used as read-only
 func (t *MultiTSDB) TSDBExemplars() map[string]*exemplars.TSDB {
 	t.mtx.RLock()
 	defer t.mtx.RUnlock()
diff --git a/pkg/receive/multitsdb_test.go b/pkg/receive/multitsdb_test.go
index a846dfbbdd9..c4da47841ef 100644
--- a/pkg/receive/multitsdb_test.go
+++ b/pkg/receive/multitsdb_test.go
@@ -194,7 +194,7 @@ func TestMultiTSDB(t *testing.T) {
 		testutil.Ok(t, m.Open())
 		testutil.Ok(t, appendSample(m, testTenant, time.Now()))
 
-		tenant := m.tenants[testTenant]
+		tenant := m.getTenant(testTenant)
 		db := tenant.readyStorage().Get()
 
 		testutil.Equals(t, 0, len(db.Blocks()))
@@ -843,7 +843,7 @@ func appendSampleWithLabels(m *MultiTSDB, tenant string, lbls labels.Labels, tim
 
 func queryLabelValues(ctx context.Context, m *MultiTSDB) error {
 	proxy := store.NewProxyStore(nil, nil, func() []store.Client {
-		clients := m.TSDBLocalClients()
+		clients := m.tsdbLocalClientsCopied()
 		if len(clients) > 0 {
 			clients[0] = &slowClient{clients[0]}
 		}
diff --git a/pkg/receive/receive_test.go b/pkg/receive/receive_test.go
index 7e90c3c204a..dd9b33048a6 100644
--- a/pkg/receive/receive_test.go
+++ b/pkg/receive/receive_test.go
@@ -210,7 +210,7 @@ func TestAddingExternalLabelsForTenants(t *testing.T) {
 
 			for _, c := range tc.cfg {
 				for _, tenantId := range c.Tenants {
-					if m.tenants[tenantId] == nil {
+					if m.getTenant(tenantId) == nil {
 						err = appendSample(m, tenantId, time.Now())
 						require.NoError(t, err)
 					}
@@ -294,7 +294,7 @@ func TestLabelSetsOfTenantsWhenAddingTenants(t *testing.T) {
 
 		for _, c := range initialConfig {
 			for _, tenantId := range c.Tenants {
-				if m.tenants[tenantId] == nil {
+				if m.getTenant(tenantId) == nil {
 					err = appendSample(m, tenantId, time.Now())
 					require.NoError(t, err)
 				}
@@ -319,7 +319,7 @@ func TestLabelSetsOfTenantsWhenAddingTenants(t *testing.T) {
 
 		for _, c := range changedConfig {
 			for _, tenantId := range c.Tenants {
-				if m.tenants[tenantId] == nil {
+				if m.getTenant(tenantId) == nil {
 					err = appendSample(m, tenantId, time.Now())
 					require.NoError(t, err)
 				}
@@ -534,7 +534,7 @@ func TestLabelSetsOfTenantsWhenChangingLabels(t *testing.T) {
 
 			for _, c := range initialConfig {
 				for _, tenantId := range c.Tenants {
-					if m.tenants[tenantId] == nil {
+					if m.getTenant(tenantId) == nil {
 						err = appendSample(m, tenantId, time.Now())
 						require.NoError(t, err)
 					}
@@ -704,7 +704,7 @@ func TestAddingLabelsWhenTenantAppearsInMultipleHashrings(t *testing.T) {
 
 			for _, c := range initialConfig {
 				for _, tenantId := range c.Tenants {
-					if m.tenants[tenantId] == nil {
+					if m.getTenant(tenantId) == nil {
 						err = appendSample(m, tenantId, time.Now())
 						require.NoError(t, err)
 					}
@@ -778,7 +778,7 @@ func TestReceiverLabelsNotOverwrittenByExternalLabels(t *testing.T) {
 
 		for _, c := range cfg {
 			for _, tenantId := range c.Tenants {
-				if m.tenants[tenantId] == nil {
+				if m.getTenant(tenantId) == nil {
 					err = appendSample(m, tenantId, time.Now())
 					require.NoError(t, err)
 				}