disttask, ddl: refine disttask_flow

ddl/ddl.go

@@ -690,15 +690,20 @@
 			return NewBackfillSchedulerHandle(ctx, taskMeta, d, step == proto.StepTwo)
 		})
 
-	dispatcher.RegisterTaskFlowHandle(BackfillTaskType, NewLitBackfillFlowHandle(d))
-	scheduler.RegisterSubtaskExectorConstructor(BackfillTaskType, proto.StepOne,
-		func(proto.MinimalTask, int64) (scheduler.SubtaskExecutor, error) {
-			return &scheduler.EmptyExecutor{}, nil
-		})
-	scheduler.RegisterSubtaskExectorConstructor(BackfillTaskType, proto.StepTwo,
-		func(proto.MinimalTask, int64) (scheduler.SubtaskExecutor, error) {
-			return &scheduler.EmptyExecutor{}, nil
-		})
+	backfillHandle, err := NewLitBackfillFlowHandle(d)
+	if err != nil {
+		logutil.BgLogger().Warn("NewLitBackfillFlowHandle failed", zap.String("category", "ddl"), zap.Error(err))
+	} else {
+		dispatcher.RegisterTaskFlowHandle(BackfillTaskType, backfillHandle)
+		scheduler.RegisterSubtaskExectorConstructor(BackfillTaskType, proto.StepOne,
+			func(proto.MinimalTask, int64) (scheduler.SubtaskExecutor, error) {
+				return &scheduler.EmptyExecutor{}, nil
+			})
+		scheduler.RegisterSubtaskExectorConstructor(BackfillTaskType, proto.StepTwo,
+			func(proto.MinimalTask, int64) (scheduler.SubtaskExecutor, error) {
+				return &scheduler.EmptyExecutor{}, nil
+			})
+	}
 
 	// Register functions for enable/disable ddl when changing system variable `tidb_enable_ddl`.
 	variable.EnableDDL = d.EnableDDL

ddl/disttask_flow.go

@@ -33,142 +33,82 @@
 )
 
 type litBackfillFlowHandle struct {
-	d DDL
+	d *ddl
 }
 
 var _ dispatcher.TaskFlowHandle = (*litBackfillFlowHandle)(nil)
 
 // NewLitBackfillFlowHandle creates a new litBackfillFlowHandle.
-func NewLitBackfillFlowHandle(d DDL) dispatcher.TaskFlowHandle {
-	return &litBackfillFlowHandle{
-		d: d,
+func NewLitBackfillFlowHandle(d DDL) (dispatcher.TaskFlowHandle, error) {
+	ddl, ok := d.(*ddl)
+	if !ok {
+		return nil, errors.New("The getDDL result should be the type of *ddl")
 	}
+	return &litBackfillFlowHandle{
+		d: ddl,
+	}, nil
 }
 
 func (*litBackfillFlowHandle) OnTicker(_ context.Context, _ *proto.Task) {
 }
 
 // ProcessNormalFlow processes the normal flow.
-func (h *litBackfillFlowHandle) ProcessNormalFlow(_ context.Context, _ dispatcher.TaskHandle, gTask *proto.Task) (metas [][]byte, err error) {
+func (h *litBackfillFlowHandle) ProcessNormalFlow(ctx context.Context, _ dispatcher.TaskHandle, gTask *proto.Task) ([][]byte, error) {
 	var globalTaskMeta BackfillGlobalMeta
-	if err = json.Unmarshal(gTask.Meta, &globalTaskMeta); err != nil {
+	if err := json.Unmarshal(gTask.Meta, &globalTaskMeta); err != nil {
 		return nil, err
 	}
-
-	d, ok := h.d.(*ddl)
-	if !ok {
-		return nil, errors.New("The getDDL result should be the type of *ddl")
-	}
-
 	job := &globalTaskMeta.Job
 	var tblInfo *model.TableInfo
-	err = kv.RunInNewTxn(d.ctx, d.store, true, func(ctx context.Context, txn kv.Transaction) error {
-		tblInfo, err = meta.NewMeta(txn).GetTable(job.SchemaID, job.TableID)
-		return err
-	})
+	tblInfo, err := getTblInfo(h.d, job)
 	if err != nil {
 		return nil, err
 	}
-
 	var subTaskMetas [][]byte
-	if tblInfo.Partition == nil {
-		switch gTask.Step {
-		case proto.StepOne:
-			serverNodes, err := dispatcher.GenerateSchedulerNodes(d.ctx)
-			if err != nil {
-				return nil, err
-			}
-			subTaskMetas = make([][]byte, 0, len(serverNodes))
-			dummyMeta := &BackfillSubTaskMeta{}
-			metaBytes, err := json.Marshal(dummyMeta)
-			if err != nil {
-				return nil, err
-			}
-			for range serverNodes {
-				subTaskMetas = append(subTaskMetas, metaBytes)
-			}
-			gTask.Step = proto.StepTwo
-			return subTaskMetas, nil
-		case proto.StepTwo:
+	// generate partition table's plan.
+	if tblInfo.Partition != nil {
+		if gTask.Step != proto.StepInit {
+			// This flow for partition table has only one step
 			return nil, nil
-		default:
 		}
-		tbl, err := getTable(d.store, job.SchemaID, tblInfo)
+		subTaskMetas, err := generatePartitionPlan(tblInfo)
 		if err != nil {
 			return nil, err
 		}
-		ver, err := getValidCurrentVersion(d.store)
+		gTask.Step = proto.StepOne
+		return subTaskMetas, nil
+	}
+
+	// generate non-partition table's plan.
+	switch gTask.Step {
+	case proto.StepInit:
+		subtaskMeta, err := generateNonPartitionPlan(h.d, tblInfo, job)
 		if err != nil {
-			return nil, errors.Trace(err)
-		}
-		startKey, endKey, err := getTableRange(d.jobContext(job.ID), d.ddlCtx, tbl.(table.PhysicalTable), ver.Ver, job.Priority)
-		if startKey == nil && endKey == nil {
-			// Empty table.
-			gTask.Step = proto.StepOne
-			return nil, nil
+			return nil, err
 		}
+		gTask.Step = proto.StepOne
+		return subtaskMeta, nil
+	case proto.StepOne:
+		serverNodes, err := dispatcher.GenerateSchedulerNodes(ctx)
 		if err != nil {
-			return nil, errors.Trace(err)
+			return nil, err
 		}
-		regionCache := d.store.(helper.Storage).GetRegionCache()
-		recordRegionMetas, err := regionCache.LoadRegionsInKeyRange(tikv.NewBackofferWithVars(context.Background(), 20000, nil), startKey, endKey)
+		subTaskMetas = make([][]byte, 0, len(serverNodes))
+		dummyMeta := &BackfillSubTaskMeta{}
+		metaBytes, err := json.Marshal(dummyMeta)
 		if err != nil {
 			return nil, err
 		}
-
-		subTaskMetas = make([][]byte, 0, 100)
-		regionBatch := 20
-		sort.Slice(recordRegionMetas, func(i, j int) bool {
-			return bytes.Compare(recordRegionMetas[i].StartKey(), recordRegionMetas[j].StartKey()) < 0
-		})
-		for i := 0; i < len(recordRegionMetas); i += regionBatch {
-			end := i + regionBatch
-			if end > len(recordRegionMetas) {
-				end = len(recordRegionMetas)
-			}
-			batch := recordRegionMetas[i:end]
-			subTaskMeta := &BackfillSubTaskMeta{StartKey: batch[0].StartKey(), EndKey: batch[len(batch)-1].EndKey()}
-			if i == 0 {
-				subTaskMeta.StartKey = startKey
-			}
-			if end == len(recordRegionMetas) {
-				subTaskMeta.EndKey = endKey
-			}
-			metaBytes, err := json.Marshal(subTaskMeta)
-			if err != nil {
-				return nil, err
-			}
-			subTaskMetas = append(subTaskMetas, metaBytes)
-		}
-	} else {
-		if gTask.State != proto.TaskStatePending {
-			// This flow for partition table has only one step, finish task when it is not pending
-			return nil, nil
-		}
-
-		defs := tblInfo.Partition.Definitions
-		physicalIDs := make([]int64, len(defs))
-		for i := range defs {
-			physicalIDs[i] = defs[i].ID
-		}
-
-		subTaskMetas = make([][]byte, 0, len(physicalIDs))
-		for _, physicalID := range physicalIDs {
-			subTaskMeta := &BackfillSubTaskMeta{
-				PhysicalTableID: physicalID,
-			}
-
-			metaBytes, err := json.Marshal(subTaskMeta)
-			if err != nil {
-				return nil, err
-			}
-
+		for range serverNodes {
 			subTaskMetas = append(subTaskMetas, metaBytes)
 		}
+		gTask.Step = proto.StepTwo
+		return subTaskMetas, nil
+	case proto.StepTwo:
+		return nil, nil
+	default:
+		return nil, nil
 	}
-
-	gTask.Step = proto.StepOne
-	return subTaskMetas, nil
 }
 
 func (*litBackfillFlowHandle) ProcessErrFlow(_ context.Context, _ dispatcher.TaskHandle, task *proto.Task, receiveErr []error) (meta []byte, err error) {
@@ -187,3 +127,88 @@
 func (*litBackfillFlowHandle) IsRetryableErr(error) bool {
 	return true
 }
+
+func getTblInfo(d *ddl, job *model.Job) (tblInfo *model.TableInfo, err error) {
+	err = kv.RunInNewTxn(d.ctx, d.store, true, func(ctx context.Context, txn kv.Transaction) error {
+		tblInfo, err = meta.NewMeta(txn).GetTable(job.SchemaID, job.TableID)
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return tblInfo, nil
+}
+
+func generatePartitionPlan(tblInfo *model.TableInfo) (metas [][]byte, err error) {
+	defs := tblInfo.Partition.Definitions
+	physicalIDs := make([]int64, len(defs))
+	for i := range defs {
+		physicalIDs[i] = defs[i].ID
+	}
+
+	subTaskMetas := make([][]byte, 0, len(physicalIDs))
+	for _, physicalID := range physicalIDs {
+		subTaskMeta := &BackfillSubTaskMeta{
+			PhysicalTableID: physicalID,
+		}
+
+		metaBytes, err := json.Marshal(subTaskMeta)
+		if err != nil {
+			return nil, err
+		}
+
+		subTaskMetas = append(subTaskMetas, metaBytes)
+	}
+	return subTaskMetas, nil
+}
+
+func generateNonPartitionPlan(d *ddl, tblInfo *model.TableInfo, job *model.Job) (metas [][]byte, err error) {
+	tbl, err := getTable(d.store, job.SchemaID, tblInfo)
+	if err != nil {
+		return nil, err
+	}
+	ver, err := getValidCurrentVersion(d.store)
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	startKey, endKey, err := getTableRange(d.jobContext(job.ID), d.ddlCtx, tbl.(table.PhysicalTable), ver.Ver, job.Priority)
+	if startKey == nil && endKey == nil {
+		// Empty table.
+		return nil, nil
+	}
+	if err != nil {
+		return nil, errors.Trace(err)
+	}
+	regionCache := d.store.(helper.Storage).GetRegionCache()
+	recordRegionMetas, err := regionCache.LoadRegionsInKeyRange(tikv.NewBackofferWithVars(context.Background(), 20000, nil), startKey, endKey)
+	if err != nil {
+		return nil, err
+	}
+
+	subTaskMetas := make([][]byte, 0, 100)
+	regionBatch := 20
+	sort.Slice(recordRegionMetas, func(i, j int) bool {
+		return bytes.Compare(recordRegionMetas[i].StartKey(), recordRegionMetas[j].StartKey()) < 0
+	})
+	for i := 0; i < len(recordRegionMetas); i += regionBatch {
+		end := i + regionBatch
+		if end > len(recordRegionMetas) {
+			end = len(recordRegionMetas)
+		}
+		batch := recordRegionMetas[i:end]
+		subTaskMeta := &BackfillSubTaskMeta{StartKey: batch[0].StartKey(), EndKey: batch[len(batch)-1].EndKey()}
+		if i == 0 {
+			subTaskMeta.StartKey = startKey
+		}
+		if end == len(recordRegionMetas) {
+			subTaskMeta.EndKey = endKey
+		}
+		metaBytes, err := json.Marshal(subTaskMeta)
+		if err != nil {
+			return nil, err
+		}
+		subTaskMetas = append(subTaskMetas, metaBytes)
+	}
+	return subTaskMetas, nil
+}

ddl/disttask_flow_test.go

@@ -33,8 +33,8 @@ import (
 
 func TestBackfillFlowHandle(t *testing.T) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
-	handler := ddl.NewLitBackfillFlowHandle(dom.DDL())
-
+	handler, err := ddl.NewLitBackfillFlowHandle(dom.DDL())
+	require.NoError(t, err)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
 

tests/realtikvtest/addindextest/add_index_test.go

@@ -150,6 +150,15 @@ func TestAddIndexDistBasic(t *testing.T) {
 	tk.MustExec("split table t between (3) and (8646911284551352360) regions 50;")
 	tk.MustExec("alter table t add index idx(a);")
 	tk.MustExec("admin check index t idx;")
+
+	tk.MustExec("create table t1(a bigint auto_random primary key);")
+	tk.MustExec("insert into t1 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t1 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t1 values (), (), (), (), (), ()")
+	tk.MustExec("insert into t1 values (), (), (), (), (), ()")
+	tk.MustExec("split table t1 between (3) and (8646911284551352360) regions 50;")
+	tk.MustExec("alter table t1 add index idx(a);")
+	tk.MustExec("admin check index t1 idx;")
 	tk.MustExec(`set global tidb_enable_dist_task=0;`)
 }