test: fix unstable TestFrameworkCancelGTask (#44923)

close #44904
pingcap · Jun 26, 2023 · cae4314 · cae4314
1 parent 8d7d961
commit cae4314
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 22 deletions.
diff --git a/disttask/framework/framework_test.go b/disttask/framework/framework_test.go
@@ -159,29 +159,12 @@ func DispatchTaskAndCheckSuccess(taskKey string, t *testing.T, v *atomic.Int64)
 }
 
 func DispatchAndCancelTask(taskKey string, t *testing.T, v *atomic.Int64) {
-	mgr, err := storage.GetTaskManager()
-	require.NoError(t, err)
-	taskID, err := mgr.AddNewGlobalTask(taskKey, proto.TaskTypeExample, 8, nil)
-	require.NoError(t, err)
-	start := time.Now()
-	mgr.CancelGlobalTask(taskID)
-	var task *proto.Task
-	for {
-		if time.Since(start) > 2*time.Minute {
-			require.FailNow(t, "timeout")
-		}
-
-		time.Sleep(time.Second)
-		task, err = mgr.GetGlobalTaskByID(taskID)
-		require.NoError(t, err)
-		require.NotNil(t, task)
-		if task.State != proto.TaskStatePending && task.State != proto.TaskStateRunning && task.State != proto.TaskStateCancelling && task.State != proto.TaskStateReverting {
-			break
-		}
-	}
-
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/disttask/framework/scheduler/MockExecutorRunCancel", "1*return(1)"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/disttask/framework/scheduler/MockExecutorRunCancel"))
+	}()
+	task := DispatchTask(taskKey, t)
 	require.Equal(t, proto.TaskStateReverted, task.State)
-	require.Equal(t, int64(0), v.Load())
 	v.Store(0)
 }
 

diff --git a/disttask/framework/scheduler/BUILD.bazel b/disttask/framework/scheduler/BUILD.bazel
@@ -13,6 +13,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//disttask/framework/proto",
+        "//disttask/framework/storage",
         "//resourcemanager/pool/spool",
         "//resourcemanager/util",
         "//util/logutil",

diff --git a/disttask/framework/scheduler/scheduler.go b/disttask/framework/scheduler/scheduler.go
@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/disttask/framework/proto"
+	"github.com/pingcap/tidb/disttask/framework/storage"
 	"github.com/pingcap/tidb/util/logutil"
 	"go.uber.org/zap"
 )
@@ -260,6 +261,19 @@ func (s *InternalSchedulerImpl) runMinimalTask(minimalTaskCtx context.Context, m
 			s.onError(errors.New("MockExecutorRunErr"))
 		}
 	})
+	failpoint.Inject("MockExecutorRunCancel", func(val failpoint.Value) {
+		if taskID, ok := val.(int); ok {
+			mgr, err := storage.GetTaskManager()
+			if err != nil {
+				logutil.BgLogger().Error("get task manager failed", zap.Error(err))
+			} else {
+				err = mgr.CancelGlobalTask(int64(taskID))
+				if err != nil {
+					logutil.BgLogger().Error("cancel global task failed", zap.Error(err))
+				}
+			}
+		}
+	})
 	if err = executor.Run(minimalTaskCtx); err != nil {
 		s.onError(err)
 	}