importinto: integrate async prepare flow for nextgen

pkg/ddl/backfilling_dist_scheduler.go

@@ -106,6 +106,11 @@ func (sch *LitBackfillScheduler) Close() {
 func (*LitBackfillScheduler) OnTick(_ context.Context, _ *proto.Task) {
 }
 
+// OnPrepare implements scheduler.Extension interface.
+func (*LitBackfillScheduler) OnPrepare(context.Context, diststorage.TaskHandle, *proto.Task) error {
+	return nil
+}
+
 // OnNextSubtasksBatch generate batch of next step's plan.
 func (sch *LitBackfillScheduler) OnNextSubtasksBatch(
 	ctx context.Context,

pkg/dxf/example/scheduler.go

@@ -103,6 +103,10 @@ func (*schedulerImpl) GetNextStep(task *proto.TaskBase) proto.Step {
 	}
 }
 
+func (*schedulerImpl) OnPrepare(context.Context, storage.TaskHandle, *proto.Task) error {
+	return nil
+}
+
 type postCleanupImpl struct{}
 
 func (*postCleanupImpl) CleanUp(_ context.Context, task *proto.Task) error {

pkg/dxf/framework/scheduler/scheduler.go

@@ -151,12 +151,6 @@ func (s *BaseScheduler) GetTask() *proto.Task {
 	return s.task.Load()
 }
 
-// OnPrepare implements Extension with a no-op default.
-// will be removed in later PR.
-func (*BaseScheduler) OnPrepare(context.Context, storage.TaskHandle, *proto.Task) error {
-	return nil
-}
-
 // getTaskClone returns a clone of the task.
 func (s *BaseScheduler) getTaskClone() *proto.Task {
 	clone := *s.GetTask()
@@ -392,7 +386,7 @@ func (s *BaseScheduler) onPending() error {
 	s.logger.Debug("on pending state", zap.Stringer("state", task.State),
 		zap.String("step", proto.Step2Str(task.Type, task.Step)))
 	if task.Step == proto.StepInit && task.ExtraParams.PrepareMode == proto.PrepareModeRequired {
-		if err := s.Extension.OnPrepare(s.ctx, s, task); err != nil {
+		if err := s.OnPrepare(s.ctx, s, task); err != nil {
 			return s.handlePrepareOrPlanErr(err)
 		}
 		switched, err := s.taskMgr.SwitchTaskStepAfterPrepare(s.ctx, task)
@@ -404,11 +398,9 @@ func (s *BaseScheduler) onPending() error {
 		}
 		task.Step = proto.StepPrepared
 		s.task.Store(task)
+		failpoint.InjectCall("afterTaskPrepared", task)
 		// fall through to switch to next step to avoid wait another tick to
 		// schedule subtasks after prepare.
-		// NOTE: no real business task enables prepare mode yet, so StepPrepared
-		// is only exercised by framework tests for now. Business GetNextStep
-		// StepPrepared integration will be done in a follow-up PR.
 	}
 	return s.switch2NextStep()
 }

pkg/dxf/importinto/BUILD.bazel

@@ -110,7 +110,7 @@ go_test(
     ],
     embed = [":importinto"],
     flaky = True,
-    shard_count = 27,
+    shard_count = 28,
     deps = [
         "//pkg/config",
         "//pkg/config/kerneltype",
@@ -156,6 +156,7 @@ go_test(
         "//pkg/testkit/testfailpoint",
         "//pkg/types",
         "//pkg/util",
+        "//pkg/util/dbterror/exeerrors",
         "//pkg/util/etcd",
         "//pkg/util/mock",
         "@com_github_docker_go_units//:go-units",

pkg/dxf/importinto/job.go

@@ -64,6 +64,17 @@ func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, *
 	return doSubmitTask(ctx, plan, stmt, nil, nil)
 }
 
+// ShouldUseAsyncPrepare returns whether IMPORT INTO should use
+// DXF prepare-mode asynchronous prepare.
+// Nextgen only supports global sort for IMPORT INTO in production, but local
+// sort can still be exercised in tests, so keep the IsGlobalSort check.
+func ShouldUseAsyncPrepare(plan *importer.Plan) bool {
+	failpoint.Inject("mockDisableAsyncPrepare", func() {
+		failpoint.Return(false)
+	})
+	return plan != nil && kerneltype.IsNextGen() && plan.IsGlobalSort()
+}
+
 func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instance *serverinfo.ServerInfo, chunkMap map[int32][]importer.Chunk) (int64, *proto.TaskBase, error) {
 	var instances []*serverinfo.ServerInfo
 	if instance != nil {
@@ -82,6 +93,13 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc
 		EligibleInstances: instances,
 		ChunkMap:          chunkMap,
 	}
+	asyncPrepare := ShouldUseAsyncPrepare(plan)
+	if asyncPrepare {
+		logicalPlan.PrepareMode = proto.PrepareModeRequired
+		// below params will be filled later in async prepare, init to 1 temporarily.
+		plan.ThreadCnt = 1
+		plan.MaxNodeCnt = 1
+	}
 	planCtx := planner.PlanCtx{
 		Ctx:        ctx,
 		TaskType:   proto.ImportInto,
@@ -151,14 +169,21 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc
 	if err != nil {
 		return 0, nil, err
 	}
-
-	logutil.BgLogger().Info("job submitted to task queue",
+	logFields := []zap.Field{
 		zap.Int64("job-id", jobID),
 		zap.String("task-key", task.Key),
 		zap.Int64("task-id", task.ID),
-		zap.String("data-size", units.BytesSize(float64(plan.TotalFileSize))),
-		zap.Int("thread-cnt", plan.ThreadCnt),
-		zap.Bool("global-sort", plan.IsGlobalSort()))
+		zap.Bool("global-sort", plan.IsGlobalSort()),
+		zap.Bool("async-prepare", asyncPrepare),
+	}
+	if !asyncPrepare {
+		logFields = append(logFields,
+			zap.String("data-size", units.BytesSize(float64(plan.TotalFileSize))),
+			zap.Int("thread-cnt", plan.ThreadCnt),
+			zap.Int("max-node-cnt", plan.MaxNodeCnt),
+		)
+	}
+	logutil.BgLogger().Info("job submitted to task queue", logFields...)
 
 	return jobID, task, nil
 }

pkg/dxf/importinto/job_testkit_test.go

@@ -166,6 +166,30 @@ func TestSubmitTaskNextgen(t *testing.T) {
 		sysKSTK.MustQuery("select count(1) from mysql.tidb_import_jobs where id = ?", jobID).Check(testkit.Rows("0"))
 		userKSTK.MustQuery("select count(1) from mysql.tidb_global_task where id = ?", task.ID).Check(testkit.Rows("0"))
 	})
+
+	t.Run("submit global-sort task uses async prepare mode", func(t *testing.T) {
+		sysKSTK.MustExec("delete from mysql.tidb_import_jobs")
+		sysKSTK.MustExec("delete from mysql.tidb_global_task")
+		config.UpdateGlobal(func(conf *config.Config) {
+			conf.KeyspaceName = keyspace.System
+		})
+		manuallyInitFn(t, sysKSStore, sysKSStore)
+
+		jobID, task, err := importinto.SubmitTask(ctx, &importer.Plan{
+			TableInfo:       &model.TableInfo{},
+			Parameters:      &importer.ImportParameters{},
+			ThreadCnt:       16,
+			MaxNodeCnt:      8,
+			CloudStorageURI: "local:///tmp/prepare-mode-sort",
+		}, "import into t from 's3://bucket/test.csv'")
+		require.NoError(t, err)
+		sysKSTK.MustQuery("select count(1) from mysql.tidb_import_jobs where id = ?", jobID).Check(testkit.Rows("1"))
+		sysKSTK.MustQuery(
+			"select concurrency, max_node_count, json_extract(extra_params, '$.prepare_mode') "+
+				"from mysql.tidb_global_task where id = ?",
+			task.ID,
+		).Check(testkit.Rows("1 1 1"))
+	})
 }
 
 func TestGetTaskImportedRows(t *testing.T) {

pkg/dxf/importinto/planner.go

@@ -56,7 +56,11 @@ type LogicalPlan struct {
 	Stmt              string
 	EligibleInstances []*serverinfo.ServerInfo
 	ChunkMap          map[int32][]importer.Chunk
-	Logger            *zap.Logger
+	PrepareMode       proto.PrepareMode
+	// PreparedChunkMapExternalPath points to externally persisted chunk
+	// metadata produced during framework prepare stage.
+	PreparedChunkMapExternalPath string
+	Logger                       *zap.Logger
 
 	// summary for next step
 	summary importer.StepSummary
@@ -66,17 +70,19 @@ type LogicalPlan struct {
 func (p *LogicalPlan) GetTaskExtraParams() proto.ExtraParams {
 	return proto.ExtraParams{
 		ManualRecovery: p.Plan.ManualRecovery,
+		PrepareMode:    p.PrepareMode,
 	}
 }
 
 // ToTaskMeta converts the logical plan to task meta.
 func (p *LogicalPlan) ToTaskMeta() ([]byte, error) {
 	taskMeta := TaskMeta{
-		JobID:             p.JobID,
-		Plan:              p.Plan,
-		Stmt:              p.Stmt,
-		EligibleInstances: p.EligibleInstances,
-		ChunkMap:          p.ChunkMap,
+		JobID:                    p.JobID,
+		Plan:                     p.Plan,
+		Stmt:                     p.Stmt,
+		EligibleInstances:        p.EligibleInstances,
+		ChunkMap:                 p.ChunkMap,
+		PreparedMetaExternalPath: p.PreparedChunkMapExternalPath,
 	}
 	return json.Marshal(taskMeta)
 }
@@ -92,6 +98,7 @@ func (p *LogicalPlan) FromTaskMeta(bs []byte) error {
 	p.Stmt = taskMeta.Stmt
 	p.EligibleInstances = taskMeta.EligibleInstances
 	p.ChunkMap = taskMeta.ChunkMap
+	p.PreparedChunkMapExternalPath = taskMeta.PreparedMetaExternalPath
 	return nil
 }
 
@@ -362,7 +369,13 @@ func buildControllerForPlan(p *LogicalPlan) (*importer.LoadDataController, error
 
 func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.PipelineSpec, error) {
 	var chunkMap map[int32][]importer.Chunk
-	if len(p.ChunkMap) > 0 {
+	if p.PreparedChunkMapExternalPath != "" {
+		var err error
+		chunkMap, err = readPreparedChunkMap(pCtx.Ctx, &p.Plan, p.PreparedChunkMapExternalPath)
+		if err != nil {
+			return nil, err
+		}
+	} else if len(p.ChunkMap) > 0 {
 		chunkMap = p.ChunkMap
 	} else {
 		controller, err2 := buildControllerForPlan(p)
@@ -402,6 +415,25 @@ func generateImportSpecs(pCtx planner.PlanCtx, p *LogicalPlan) ([]planner.Pipeli
 	return importSpecs, nil
 }
 
+func readPreparedChunkMap(
+	ctx context.Context,
+	plan *importer.Plan,
+	externalPath string,
+) (map[int32][]importer.Chunk, error) {
+	store, err := importer.GetSortStore(ctx, plan.CloudStorageURI)
+	if err != nil {
+		return nil, err
+	}
+	defer store.Close()
+	preparedChunkMapMeta := PreparedMeta{
+		BaseExternalMeta: globalsort.BaseExternalMeta{ExternalPath: externalPath},
+	}
+	if err := preparedChunkMapMeta.ReadJSONFromExternalStorage(ctx, store, &preparedChunkMapMeta); err != nil {
+		return nil, err
+	}
+	return preparedChunkMapMeta.ChunkMap, nil
+}
+
 func skipMergeSort(kvGroup string, stats []simplesst.MultipleFilesStat, concurrency int) bool {
 	failpoint.Inject("forceMergeSort", func(val failpoint.Value) {
 		in := val.(string)

pkg/dxf/importinto/planner_test.go

@@ -18,6 +18,7 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"path/filepath"
 	"testing"
 
 	"github.com/pingcap/errors"
@@ -139,6 +140,39 @@ func TestToPhysicalPlan(t *testing.T) {
 	_, err = logicalPlan.ToPhysicalPlan(planCtx)
 	// error when build controller plan
 	require.ErrorContains(t, err, "provide a valid URI")
+
+	t.Run("prepared chunk map external path is preferred", func(t *testing.T) {
+		cloudStorageURI := "local://" + filepath.ToSlash(t.TempDir())
+		store, err := importer.GetSortStore(context.Background(), cloudStorageURI)
+		require.NoError(t, err)
+		defer store.Close()
+
+		preparedChunkMap := map[int32][]importer.Chunk{
+			1: {{Path: "gs://test-load/2.csv"}},
+		}
+		externalPath := globalsort.PreparedMetaPath(100)
+		preparedMeta := PreparedMeta{
+			BaseExternalMeta: globalsort.BaseExternalMeta{ExternalPath: externalPath},
+			ChunkMap:         preparedChunkMap,
+		}
+		require.NoError(t, preparedMeta.WriteJSONToExternalStorage(context.Background(), store, preparedMeta))
+
+		specs, err := generateImportSpecs(planner.PlanCtx{Ctx: context.Background()}, &LogicalPlan{
+			Plan: importer.Plan{
+				CloudStorageURI: cloudStorageURI,
+			},
+			ChunkMap: map[int32][]importer.Chunk{
+				2: {{Path: "gs://test-load/ignored.csv"}},
+			},
+			PreparedChunkMapExternalPath: externalPath,
+		})
+		require.NoError(t, err)
+		require.Len(t, specs, 1)
+		importSpec := specs[0].(*ImportSpec)
+		require.Equal(t, int32(1), importSpec.ID)
+		require.Len(t, importSpec.Chunks, 1)
+		require.Equal(t, "gs://test-load/2.csv", importSpec.Chunks[0].Path)
+	})
 }
 
 func genEncodeStepMetas(t *testing.T, cnt int) [][]byte {

pkg/dxf/importinto/proto.go

@@ -53,6 +53,20 @@ type TaskMeta struct {
 	// we use a map from engine ID to chunks since we need support split_file for CSV,
 	// so need to split them into engines before passing to scheduler.
 	ChunkMap map[int32][]importer.Chunk
+	// PreparedMetaExternalPath points to external chunk metadata prepared by
+	// framework OnPrepare for nextgen global-sort path.
+	PreparedMetaExternalPath string `json:"prepared_meta_external_path,omitempty"`
+}
+
+// PreparedMeta stores metadata generated in prepare stage.
+type PreparedMeta struct {
+	globalsort.BaseExternalMeta
+	ChunkMap map[int32][]importer.Chunk `json:"chunk_map,omitempty" external:"true"`
+}
+
+// Marshal marshals the prepared chunk map meta to JSON.
+func (m *PreparedMeta) Marshal() ([]byte, error) {
+	return m.BaseExternalMeta.Marshal(m)
 }
 
 // ImportStepMeta is the meta of import step.