From e27791e3aef9d31a3dd328da180b4c262d0ab200 Mon Sep 17 00:00:00 2001
From: Casie Chen <casie.chen@grafana.com>
Date: Fri, 20 Sep 2024 13:19:05 -0700
Subject: [PATCH] Fix race in benchmark test

---
 ...ing_algorithm_tree_queue_benchmark_test.go | 28 +++++++++++--------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/pkg/scheduler/queue/multi_queuing_algorithm_tree_queue_benchmark_test.go b/pkg/scheduler/queue/multi_queuing_algorithm_tree_queue_benchmark_test.go
index 02f56ca42f..288427e2f3 100644
--- a/pkg/scheduler/queue/multi_queuing_algorithm_tree_queue_benchmark_test.go
+++ b/pkg/scheduler/queue/multi_queuing_algorithm_tree_queue_benchmark_test.go
@@ -375,39 +375,45 @@ func TestMultiDimensionalQueueAlgorithmSlowConsumerEffects(t *testing.T) {
 	for _, weightedQueueDimensionTestCase := range weightedQueueDimensionTestCases {
 		numTenants := len(weightedQueueDimensionTestCase.tenantQueueDimensionsWeights)
 
-		tqa := newTenantQuerierAssignments()
+		tqaNonFlipped := newTenantQuerierAssignments()
+		tqaFlipped := newTenantQuerierAssignments()
+		tqaQuerierWorkerPrioritization := newTenantQuerierAssignments()
 
-		nonFlippedRoundRobinTree, err := NewTree(tqa, &roundRobinState{}, &roundRobinState{})
+		nonFlippedRoundRobinTree, err := NewTree(tqaNonFlipped, &roundRobinState{}, &roundRobinState{})
 		require.NoError(t, err)
 
-		flippedRoundRobinTree, err := NewTree(&roundRobinState{}, tqa, &roundRobinState{})
+		flippedRoundRobinTree, err := NewTree(&roundRobinState{}, tqaFlipped, &roundRobinState{})
 		require.NoError(t, err)
 
-		querierWorkerPrioritizationTree, err := NewTree(NewQuerierWorkerQueuePriorityAlgo(), tqa, &roundRobinState{})
+		querierWorkerPrioritizationTree, err := NewTree(NewQuerierWorkerQueuePriorityAlgo(), tqaQuerierWorkerPrioritization, &roundRobinState{})
 		require.NoError(t, err)
 
-		trees := []struct {
+		treeScenarios := []struct {
 			name string
 			tree Tree
+			tqa  *tenantQuerierAssignments
 		}{
 			// keeping these names the same length keeps logged results aligned
 			{
 				"tenant-querier -> query component round-robin tree",
 				nonFlippedRoundRobinTree,
+				tqaNonFlipped,
 			},
 			{
 				"query component round-robin -> tenant-querier tree",
 				flippedRoundRobinTree,
+				tqaFlipped,
 			},
 			{
 				"worker-queue prioritization -> tenant-querier tree",
 				querierWorkerPrioritizationTree,
+				tqaQuerierWorkerPrioritization,
 			},
 		}
-		for _, tree := range trees {
+		for _, scenario := range treeScenarios {
 			testCaseName := fmt.Sprintf(
 				"tree: %s, %s",
-				tree.name,
+				scenario.name,
 				weightedQueueDimensionTestCase.name,
 			)
 			testCaseObservations := &testScenarioQueueDurationObservations{
@@ -417,7 +423,7 @@ func TestMultiDimensionalQueueAlgorithmSlowConsumerEffects(t *testing.T) {
 			}
 
 			// only the non-flipped tree uses the old tenant -> query component hierarchy
-			prioritizeQueryComponents := tree.tree != nonFlippedRoundRobinTree
+			prioritizeQueryComponents := scenario.tree != nonFlippedRoundRobinTree
 
 			t.Run(testCaseName, func(t *testing.T) {
 				queue, err := NewRequestQueue(
@@ -434,9 +440,9 @@ func TestMultiDimensionalQueueAlgorithmSlowConsumerEffects(t *testing.T) {
 
 				// NewRequestQueue constructor does not allow passing in a tree or tenantQuerierAssignments
 				// so we have to override here to use the same structures as the test case
-				queue.queueBroker.tenantQuerierAssignments = tqa
+				queue.queueBroker.tenantQuerierAssignments = scenario.tqa
 				queue.queueBroker.prioritizeQueryComponents = prioritizeQueryComponents
-				queue.queueBroker.tree = tree.tree
+				queue.queueBroker.tree = scenario.tree
 
 				ctx := context.Background()
 				require.NoError(t, queue.starting(ctx))
@@ -482,7 +488,7 @@ func TestMultiDimensionalQueueAlgorithmSlowConsumerEffects(t *testing.T) {
 				testCaseReports[testCaseName] = report
 
 				// ensure everything was dequeued
-				path, val := tree.tree.Dequeue(&DequeueArgs{querierID: tqa.currentQuerier})
+				path, val := scenario.tree.Dequeue(&DequeueArgs{querierID: scenario.tqa.currentQuerier})
 				assert.Nil(t, val)
 				assert.Equal(t, path, QueuePath{})
 			})