From e3abcbbe6801a59531ac9313eec4c28aaf3a4344 Mon Sep 17 00:00:00 2001
From: EasonBall <592838129@qq.com>
Date: Thu, 11 Jan 2024 13:57:25 +0800
Subject: [PATCH] test: fix race in TestManagerHandleTasks (#50304)

close pingcap/tidb#50297
---
 pkg/disttask/framework/taskexecutor/manager_test.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pkg/disttask/framework/taskexecutor/manager_test.go b/pkg/disttask/framework/taskexecutor/manager_test.go
index 27e71a6e18487..e351aa8cc3c1a 100644
--- a/pkg/disttask/framework/taskexecutor/manager_test.go
+++ b/pkg/disttask/framework/taskexecutor/manager_test.go
@@ -297,6 +297,7 @@ func TestManagerHandleTasks(t *testing.T) {
 
 	// handle pending tasks
 	var task1Ctx context.Context
+	var mu sync.Mutex
 	mockTaskTable.EXPECT().GetTasksInStates(m.ctx, proto.TaskStateRunning, proto.TaskStateReverting, proto.TaskStatePausing).
 		Return([]*proto.Task{task1}, nil)
 	mockTaskTable.EXPECT().HasSubtasksInStates(m.ctx, id, task1.ID, proto.StepOne,
@@ -307,11 +308,15 @@ func TestManagerHandleTasks(t *testing.T) {
 	mockTaskTable.EXPECT().HasSubtasksInStates(m.ctx, id, task1.ID, proto.StepOne,
 		unfinishedSubtaskStates).Return(true, nil)
 	mockInternalExecutor.EXPECT().Run(gomock.Any(), task1).DoAndReturn(func(ctx context.Context, _ *proto.Task) error {
+		mu.Lock()
 		task1Ctx = ctx
+		mu.Unlock()
 		return <-ch
 	})
 	m.handleTasks()
 	require.Eventually(t, func() bool {
+		mu.Lock()
+		defer mu.Unlock()
 		return task1Ctx != nil && ctrl.Satisfied()
 	}, 5*time.Second, 100*time.Millisecond)
 	require.True(t, m.isExecutorStarted(task1.ID))