From 0866673e37e03bae70ca320099acac1460f92b4e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Emilio=20Alvarez=20Pi=C3=B1eiro?=
 <95703246+emilioalvap@users.noreply.github.com>
Date: Wed, 14 Jun 2023 22:01:51 +0200
Subject: [PATCH] [Heartbeat] Fix race condition in scheduler tests (#35728)

---
 heartbeat/scheduler/scheduler_test.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/heartbeat/scheduler/scheduler_test.go b/heartbeat/scheduler/scheduler_test.go
index a50e5cdcce8..867a50c2962 100644
--- a/heartbeat/scheduler/scheduler_test.go
+++ b/heartbeat/scheduler/scheduler_test.go
@@ -231,7 +231,7 @@ func TestSchedTaskLimits(t *testing.T) {
 			numJobs: 50,
 			limit:   math.MaxInt64,
 			expect: func(events []int) {
-				require.GreaterOrEqual(t, len(events), 50)
+				require.GreaterOrEqual(t, len(events), 250)
 			},
 		},
 		{
@@ -239,7 +239,7 @@ func TestSchedTaskLimits(t *testing.T) {
 			numJobs: 100,
 			limit:   0,
 			expect: func(events []int) {
-				require.GreaterOrEqual(t, len(events), 100)
+				require.GreaterOrEqual(t, len(events), 500)
 			},
 		},
 	}
@@ -254,11 +254,15 @@ func TestSchedTaskLimits(t *testing.T) {
 			}
 			s := Create(math.MaxInt64, monitoring.NewRegistry(), tarawaTime(), jobConfigByType, false)
 			var taskArr []int
+			mtx := sync.Mutex{}
 			wg := sync.WaitGroup{}
 			wg.Add(tt.numJobs)
 			for i := 0; i < tt.numJobs; i++ {
 				num := i
 				tf := makeTasks(4, func() {
+					mtx.Lock()
+					defer mtx.Unlock() // taskArr is shared across goroutines
+
 					taskArr = append(taskArr, num)
 				})
 				go func(tff TaskFunc) {