fix data race in testing with some concurrent safe helper functions

Signed-off-by: Yi Jin <[email protected]>
thanos-io · Nov 28, 2024 · ca6c149 · ca6c149
1 parent 7193180
commit ca6c149
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 8 deletions.
diff --git a/pkg/receive/multitsdb.go b/pkg/receive/multitsdb.go
@@ -136,6 +136,24 @@ 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()
+	return t.tenants[tenantID]
+}
+
+// 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 +644,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
@@ -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
@@ -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)
 				}