disttask: dynamic dispatch subtasks

ddl/backfilling_dispatcher.go

@@ -50,11 +50,12 @@ func NewBackfillingDispatcherExt(d DDL) (dispatcher.Extension, error) {
 	}, nil
 }
 
+// OnTick implements dispatcher.Extension interface.
 func (*backfillingDispatcherExt) OnTick(_ context.Context, _ *proto.Task) {
 }
 
-// OnNextStage generate next stage's plan.
-func (h *backfillingDispatcherExt) OnNextStage(ctx context.Context, _ dispatcher.TaskHandle, gTask *proto.Task) ([][]byte, error) {
+// OnNextSubtasksBatch generate batch of next stage's plan.
+func (h *backfillingDispatcherExt) OnNextSubtasksBatch(ctx context.Context, _ dispatcher.TaskHandle, gTask *proto.Task) ([][]byte, error) {
 	var globalTaskMeta BackfillGlobalMeta
 	if err := json.Unmarshal(gTask.Meta, &globalTaskMeta); err != nil {
 		return nil, err
@@ -76,7 +77,6 @@ func (h *backfillingDispatcherExt) OnNextStage(ctx context.Context, _ dispatcher
 		if err != nil {
 			return nil, err
 		}
-		gTask.Step = proto.StepOne
 		return subTaskMetas, nil
 	}
 
@@ -87,7 +87,6 @@ func (h *backfillingDispatcherExt) OnNextStage(ctx context.Context, _ dispatcher
 		if err != nil {
 			return nil, err
 		}
-		gTask.Step = proto.StepOne
 		return subtaskMeta, nil
 	case proto.StepOne:
 		serverNodes, err := dispatcher.GenerateSchedulerNodes(ctx)
@@ -103,15 +102,22 @@ func (h *backfillingDispatcherExt) OnNextStage(ctx context.Context, _ dispatcher
 		for range serverNodes {
 			subTaskMetas = append(subTaskMetas, metaBytes)
 		}
-		gTask.Step = proto.StepTwo
 		return subTaskMetas, nil
-	case proto.StepTwo:
-		return nil, nil
 	default:
 		return nil, nil
 	}
 }
 
+// StageFinished check if current stage finished.
+func (*backfillingDispatcherExt) StageFinished(_ *proto.Task) bool {
+	return true
+}
+
+// Finished check if current task finished.
+func (*backfillingDispatcherExt) Finished(task *proto.Task) bool {
+	return task.Step == proto.StepOne
+}
+
 // OnErrStage generate error handling stage's plan.
 func (*backfillingDispatcherExt) OnErrStage(_ context.Context, _ dispatcher.TaskHandle, task *proto.Task, receiveErr []error) (meta []byte, err error) {
 	// We do not need extra meta info when rolling back

ddl/backfilling_dispatcher_test.go

@@ -38,7 +38,7 @@ func TestBackfillingDispatcher(t *testing.T) {
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
 
-	// test partition table OnNextStage.
+	/// 1. test partition table.
 	tk.MustExec("create table tp1(id int primary key, v int) PARTITION BY RANGE (id) (\n    " +
 		"PARTITION p0 VALUES LESS THAN (10),\n" +
 		"PARTITION p1 VALUES LESS THAN (100),\n" +
@@ -48,23 +48,25 @@ func TestBackfillingDispatcher(t *testing.T) {
 	tbl, err := dom.InfoSchema().TableByName(model.NewCIStr("test"), model.NewCIStr("tp1"))
 	require.NoError(t, err)
 	tblInfo := tbl.Meta()
-	metas, err := dsp.OnNextStage(context.Background(), nil, gTask)
+
+	// 1.1 OnNextSubtasksBatch
+	metas, err := dsp.OnNextSubtasksBatch(context.Background(), nil, gTask)
 	require.NoError(t, err)
-	require.Equal(t, proto.StepOne, gTask.Step)
 	require.Equal(t, len(tblInfo.Partition.Definitions), len(metas))
 	for i, par := range tblInfo.Partition.Definitions {
 		var subTask ddl.BackfillSubTaskMeta
 		require.NoError(t, json.Unmarshal(metas[i], &subTask))
 		require.Equal(t, par.ID, subTask.PhysicalTableID)
 	}
 
-	// test partition table OnNextStage after step1 finished.
+	// 1.2 test partition table OnNextSubtasksBatch after StepInit finished.
 	gTask.State = proto.TaskStateRunning
-	metas, err = dsp.OnNextStage(context.Background(), nil, gTask)
+	gTask.Step++
+	metas, err = dsp.OnNextSubtasksBatch(context.Background(), nil, gTask)
 	require.NoError(t, err)
 	require.Equal(t, 0, len(metas))
 
-	// test partition table OnErrStage.
+	// 1.3 test partition table OnErrStage.
 	errMeta, err := dsp.OnErrStage(context.Background(), nil, gTask, []error{errors.New("mockErr")})
 	require.NoError(t, err)
 	require.Nil(t, errMeta)
@@ -73,10 +75,30 @@ func TestBackfillingDispatcher(t *testing.T) {
 	require.NoError(t, err)
 	require.Nil(t, errMeta)
 
+	/// 2. test non partition table.
+	// 2.1 empty table
 	tk.MustExec("create table t1(id int primary key, v int)")
 	gTask = createAddIndexGlobalTask(t, dom, "test", "t1", ddl.BackfillTaskType)
-	_, err = dsp.OnNextStage(context.Background(), nil, gTask)
+	metas, err = dsp.OnNextSubtasksBatch(context.Background(), nil, gTask)
+	require.NoError(t, err)
+	require.Equal(t, 0, len(metas))
+	// 2.2 non empty table.
+	tk.MustExec("create table t2(id bigint auto_random primary key)")
+	tk.MustExec("insert into t2 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t2 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t2 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t2 values (), (), (), (), (), ()")
+	gTask = createAddIndexGlobalTask(t, dom, "test", "t2", ddl.BackfillTaskType)
+	// 2.2.1 stepInit
+	metas, err = dsp.OnNextSubtasksBatch(context.Background(), nil, gTask)
+	require.NoError(t, err)
+	require.Equal(t, 1, len(metas))
+	// 2.2.2 stepOne
+	gTask.Step++
+	gTask.State = proto.TaskStateRunning
+	metas, err = dsp.OnNextSubtasksBatch(context.Background(), nil, gTask)
 	require.NoError(t, err)
+	require.Equal(t, 1, len(metas))
 }
 
 func createAddIndexGlobalTask(t *testing.T, dom *domain.Domain, dbName, tblName string, taskType string) *proto.Task {

ddl/index.go

@@ -1942,7 +1942,7 @@ func (w *worker) updateJobRowCount(taskKey string, jobID int64) {
 		logutil.BgLogger().Warn("cannot get global task", zap.String("category", "ddl"), zap.String("task_key", taskKey), zap.Error(err))
 		return
 	}
-	rowCount, err := taskMgr.GetSubtaskRowCount(gTask.ID, proto.StepOne)
+	rowCount, err := taskMgr.GetSubtaskRowCount(gTask.ID, proto.StepInit)
 	if err != nil {
 		logutil.BgLogger().Warn("cannot get subtask row count", zap.String("category", "ddl"), zap.String("task_key", taskKey), zap.Error(err))
 		return

ddl/stage_scheduler.go

@@ -98,9 +98,9 @@ func newBackfillDistScheduler(ctx context.Context, id string, taskID int64, task
 
 func (s *backfillDistScheduler) GetSubtaskExecutor(ctx context.Context, task *proto.Task, summary *execute.Summary) (execute.SubtaskExecutor, error) {
 	switch task.Step {
-	case proto.StepOne:
+	case proto.StepInit:
 		return NewBackfillSchedulerHandle(ctx, task.Meta, s.d, false, summary)
-	case proto.StepTwo:
+	case proto.StepOne:
 		return NewBackfillSchedulerHandle(ctx, task.Meta, s.d, true, nil)
 	default:
 		return nil, errors.Errorf("unknown backfill step %d for task %d", task.Step, task.ID)

disttask/framework/BUILD.bazel

@@ -4,14 +4,15 @@ go_test(
     name = "framework_test",
     timeout = "short",
     srcs = [
+        "framework_dynamic_dispatch_test.go",
         "framework_err_handling_test.go",
         "framework_ha_test.go",
         "framework_rollback_test.go",
         "framework_test.go",
     ],
     flaky = True,
     race = "off",
-    shard_count = 24,
+    shard_count = 26,
     deps = [
         "//disttask/framework/dispatcher",
         "//disttask/framework/mock",