From 61523ca2dbc1dfbfd3d8dfb2f34577b29a81eb39 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 20 May 2026 22:43:50 +0800 Subject: [PATCH 01/21] importinto: move nextgen global-sort prepare into scheduler --- pkg/dxf/importinto/job.go | 11 +++ pkg/dxf/importinto/job_testkit_test.go | 24 +++++ pkg/dxf/importinto/planner.go | 48 ++++++++-- pkg/dxf/importinto/planner_test.go | 32 +++++++ pkg/dxf/importinto/proto.go | 3 + pkg/dxf/importinto/scheduler.go | 120 ++++++++++++++++++++++++- pkg/executor/import_into.go | 33 ++++--- 7 files changed, 248 insertions(+), 23 deletions(-) diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index c983c71bb053a..0488085e42b62 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -64,6 +64,12 @@ func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, * return doSubmitTask(ctx, plan, stmt, nil, nil) } +// ShouldUseScopedPrepareIntegration enables framework prepare integration only +// for nextgen + global-sort path. +func ShouldUseScopedPrepareIntegration(plan *importer.Plan) bool { + 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 { @@ -88,6 +94,11 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc ThreadCnt: plan.ThreadCnt, MaxNodeCnt: plan.MaxNodeCnt, } + if ShouldUseScopedPrepareIntegration(plan) { + logicalPlan.PrepareMode = proto.PrepareModeRequired + planCtx.ThreadCnt = 1 + planCtx.MaxNodeCnt = 1 + } var ( jobID, taskID int64 ) diff --git a/pkg/dxf/importinto/job_testkit_test.go b/pkg/dxf/importinto/job_testkit_test.go index d7d08a836d503..0e8ba31efdd7f 100644 --- a/pkg/dxf/importinto/job_testkit_test.go +++ b/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 scoped prepare integration", 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) { diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index f41dcb231f03d..07af44dc851e5 100644 --- a/pkg/dxf/importinto/planner.go +++ b/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, + PreparedChunkMapExternalPath: 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.PreparedChunkMapExternalPath 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,27 @@ 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() + data, err := store.ReadFile(ctx, externalPath) + if err != nil { + return nil, err + } + var chunkMap map[int32][]importer.Chunk + if err := json.Unmarshal(data, &chunkMap); err != nil { + return nil, errors.Trace(err) + } + return chunkMap, nil +} + func skipMergeSort(kvGroup string, stats []simplesst.MultipleFilesStat, concurrency int) bool { failpoint.Inject("forceMergeSort", func(val failpoint.Value) { in := val.(string) diff --git a/pkg/dxf/importinto/planner_test.go b/pkg/dxf/importinto/planner_test.go index 86d53b7245ca4..21c34a03312ed 100644 --- a/pkg/dxf/importinto/planner_test.go +++ b/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,37 @@ 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 := "100/prepare-chunk-map/mock-attempt/meta.json" + data, err := json.Marshal(preparedChunkMap) + require.NoError(t, err) + require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) + + 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 { diff --git a/pkg/dxf/importinto/proto.go b/pkg/dxf/importinto/proto.go index 6575b327ea10e..6264d7ed4606e 100644 --- a/pkg/dxf/importinto/proto.go +++ b/pkg/dxf/importinto/proto.go @@ -53,6 +53,9 @@ 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 + // PreparedChunkMapExternalPath points to external chunk metadata prepared by + // framework OnPrepare for nextgen global-sort path. + PreparedChunkMapExternalPath string `json:"prepared_chunk_map_external_path,omitempty"` } // ImportStepMeta is the meta of import step. diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 2f30789c9a7e4..443a2a60e2cda 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -19,10 +19,12 @@ import ( "encoding/json" goerrors "errors" "fmt" + "path" "strconv" "sync" "time" + "github.com/google/uuid" "github.com/pingcap/errors" "github.com/pingcap/failpoint" "github.com/pingcap/tidb/br/pkg/utils" @@ -62,10 +64,11 @@ const ( // the target table, it's known to be slow to import in this case. // the value if chosen as most tables have less than 32 indexes, we can adjust // it later if needed. - warningIndexCount = 32 - registerTaskTTL = 10 * time.Minute - refreshTaskTTLInterval = 3 * time.Minute - registerTimeout = 5 * time.Second + warningIndexCount = 32 + registerTaskTTL = 10 * time.Minute + refreshTaskTTLInterval = 3 * time.Minute + registerTimeout = 5 * time.Second + preparedChunkMetaPathPurpose = "prepare-chunk-map" ) var ( @@ -164,6 +167,7 @@ type importScheduler struct { } var _ scheduler.Extension = (*importScheduler)(nil) +var _ scheduler.PrepareExtension = (*importScheduler)(nil) // NewImportScheduler creates a new import scheduler. func NewImportScheduler( @@ -288,6 +292,114 @@ func (sch *importScheduler) checkImportTableEmpty(ctx context.Context, taskMeta }) } +func (sch *importScheduler) checkActiveImportJobs(ctx context.Context, taskMeta *TaskMeta) error { + taskManager, err := sch.getTaskMgrForAccessingImportJob() + if err != nil { + return err + } + return taskManager.WithNewSession(func(se sessionctx.Context) error { + exec := se.GetSQLExecutor() + activeCnt, err := importer.GetActiveJobCnt(ctx, exec, taskMeta.Plan.DBName, taskMeta.Plan.TableInfo.Name.L) + if err != nil { + return err + } + // The current prepare-enabled job is already in pending state, so we only + // fail when there are other active jobs on the same table. + if activeCnt > 1 { + return exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs("there is active job on the target table already") + } + return nil + }) +} + +func (sch *importScheduler) writePreparedChunkMap( + ctx context.Context, + taskID int64, + cloudStorageURI string, + chunkMap map[int32][]importer.Chunk, +) (string, error) { + store, err := importer.GetSortStore(ctx, cloudStorageURI) + if err != nil { + return "", err + } + defer store.Close() + externalPath := path.Join( + strconv.FormatInt(taskID, 10), + preparedChunkMetaPathPurpose, + uuid.NewString(), + "meta.json", + ) + data, err := json.Marshal(chunkMap) + if err != nil { + return "", errors.Trace(err) + } + if err = store.WriteFile(ctx, externalPath, data); err != nil { + return "", err + } + return externalPath, nil +} + +// OnPrepare implements scheduler.PrepareExtension. +func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, task *proto.Task) error { + taskMeta := &TaskMeta{} + if err := json.Unmarshal(task.Meta, taskMeta); err != nil { + return errors.Annotate(err, "unmarshal task meta failed") + } + if !ShouldUseScopedPrepareIntegration(&taskMeta.Plan) { + return nil + } + + logicalPlan := &LogicalPlan{ + Plan: taskMeta.Plan, + Stmt: taskMeta.Stmt, + Logger: sch.GetLogger(), + } + controller, err := buildControllerForPlan(logicalPlan) + if err != nil { + return err + } + defer controller.Close() + + if err = controller.InitDataFiles(ctx); err != nil { + return err + } + if controller.TotalFileSize == 0 { + return exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs( + "No file matched, or the file is empty. Please provide a valid file location.", + ) + } + if err = sch.checkActiveImportJobs(ctx, taskMeta); err != nil { + return err + } + if err = sch.checkImportTableEmpty(ctx, taskMeta); err != nil { + return err + } + if err = controller.CalResourceParams(ctx, sch.TaskStore.GetCodec().GetKeyspace()); err != nil { + return err + } + controller.SetExecuteNodeCnt(controller.MaxNodeCnt) + chunkMap, err := controller.PopulateChunks(ctx) + if err != nil { + return err + } + chunkMapPath, err := sch.writePreparedChunkMap(ctx, task.ID, controller.Plan.CloudStorageURI, chunkMap) + if err != nil { + return err + } + + taskMeta.Plan = *controller.Plan + taskMeta.ChunkMap = nil + taskMeta.PreparedChunkMapExternalPath = chunkMapPath + metaBytes, err := json.Marshal(taskMeta) + if err != nil { + return errors.Trace(err) + } + task.Meta = metaBytes + task.RequiredSlots = controller.ThreadCnt + task.MaxNodeCount = controller.MaxNodeCnt + return nil +} + // OnNextSubtasksBatch generate batch of next stage's plan. func (sch *importScheduler) OnNextSubtasksBatch( ctx context.Context, diff --git a/pkg/executor/import_into.go b/pkg/executor/import_into.go index 42679fdfa9737..72c89c842e357 100644 --- a/pkg/executor/import_into.go +++ b/pkg/executor/import_into.go @@ -106,14 +106,17 @@ func (e *ImportIntoExec) Next(ctx context.Context, req *chunk.Chunk) (err error) return e.importFromSelect(ctx) } - if err2 := e.controller.InitDataFiles(ctx); err2 != nil { - return err2 - } - if kerneltype.IsNextGen() { - ksCodec := e.userSctx.GetStore().GetCodec().GetKeyspace() - if err2 := e.controller.CalResourceParams(ctx, ksCodec); err2 != nil { + useScopedPrepareIntegration := importinto.ShouldUseScopedPrepareIntegration(e.controller.Plan) + if !useScopedPrepareIntegration { + if err2 := e.controller.InitDataFiles(ctx); err2 != nil { return err2 } + if kerneltype.IsNextGen() { + ksCodec := e.userSctx.GetStore().GetCodec().GetKeyspace() + if err2 := e.controller.CalResourceParams(ctx, ksCodec); err2 != nil { + return err2 + } + } } // must use a new session to pre-check, else the stmt in show processlist will be changed. @@ -122,8 +125,10 @@ func (e *ImportIntoExec) Next(ctx context.Context, req *chunk.Chunk) (err error) return err2 } defer CloseSession(newSCtx) - if err2 = e.controller.CheckRequirements(ctx, newSCtx); err2 != nil { - return err2 + if !useScopedPrepareIntegration { + if err2 = e.controller.CheckRequirements(ctx, newSCtx); err2 != nil { + return err2 + } } if err := e.controller.InitTiKVConfigs(ctx, newSCtx); err != nil { @@ -217,12 +222,16 @@ func (e *ImportIntoExec) submitTask(ctx context.Context) (int64, *proto.TaskBase } logutil.Logger(ctx).Info("get job importer", zap.Stringer("param", e.controller.Parameters), zap.Bool("dist-task-enabled", vardef.EnableDistTask.Load())) + useScopedPrepareIntegration := importinto.ShouldUseScopedPrepareIntegration(e.controller.Plan) if importFromServer { - chunkMap, err2 := e.controller.PopulateChunks(ctx) - if err2 != nil { - return 0, nil, err2 + if !useScopedPrepareIntegration { + chunkMap, err2 := e.controller.PopulateChunks(ctx) + if err2 != nil { + return 0, nil, err2 + } + return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap) } - return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap) + return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, nil) } // if tidb_enable_dist_task=true, we import distributively, otherwise we import on current node. if vardef.EnableDistTask.Load() { From dfbc83da6986fe9b1464ac701af0d4c1fd8d5340 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 10:05:23 +0800 Subject: [PATCH 02/21] importer: add preparing phase for prepare-enabled jobs --- pkg/dxf/importinto/scheduler.go | 15 ++++++- pkg/dxf/importinto/scheduler_testkit_test.go | 47 ++++++++++++++++++++ pkg/executor/importer/job.go | 7 +++ pkg/executor/importer/job_test.go | 21 +++++++++ 4 files changed, 88 insertions(+), 2 deletions(-) diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 443a2a60e2cda..df3ba742f14c5 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -345,6 +345,11 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, if err := json.Unmarshal(task.Meta, taskMeta); err != nil { return errors.Annotate(err, "unmarshal task meta failed") } + if task.ExtraParams.PrepareMode == proto.PrepareModeRequired { + if err := sch.startJob(ctx, sch.GetLogger(), taskMeta, importer.JobStepPreparing); err != nil { + return err + } + } if !ShouldUseScopedPrepareIntegration(&taskMeta.Plan) { return nil } @@ -445,8 +450,14 @@ func (sch *importScheduler) OnNextSubtasksBatch( if sch.GlobalSort { jobStep = importer.JobStepGlobalSorting } - if err = sch.startJob(ctx, logger, taskMeta, jobStep); err != nil { - return nil, err + if task.ExtraParams.PrepareMode == proto.PrepareModeRequired { + if err = sch.job2Step(ctx, logger, taskMeta, jobStep); err != nil { + return nil, err + } + } else { + if err = sch.startJob(ctx, logger, taskMeta, jobStep); err != nil { + return nil, err + } } if importer.GetNumOfIndexGenKV(taskMeta.Plan.TableInfo) > warningIndexCount { dxfmetric.ScheduleEventCounter.WithLabelValues(fmt.Sprint(task.ID), dxfmetric.EventTooManyIdx).Inc() diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index 07140d7cfffdd..b0e66370c332e 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -194,6 +194,53 @@ func TestSchedulerExtLocalSort(t *testing.T) { gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) require.NoError(t, err) require.Equal(t, "cancelled", gotJobInfo.Status) + + t.Run("prepare-enabled job transitions from preparing to first business phase", func(t *testing.T) { + jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, + "root", "", &importer.ImportParameters{}, 123) + require.NoError(t, err) + + logicalPlan.JobID = jobID + bs, err := logicalPlan.ToTaskMeta() + require.NoError(t, err) + task.Meta = bs + task.Step = proto.StepInit + task.State = proto.TaskStatePending + task.ExtraParams.PrepareMode = proto.PrepareModeRequired + task.ID, err = manager.CreateTask( + ctx, + importinto.TaskKey(jobID), + proto.ImportInto, + "", + 1, + "", + 1, + proto.ExtraParams{PrepareMode: proto.PrepareModeRequired}, + bs, + ) + require.NoError(t, err) + + prepareExt, ok := ext.(scheduler.PrepareExtension) + require.True(t, ok) + require.NoError(t, prepareExt.OnPrepare(ctx, d, task)) + + gotJobInfo, err := importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) + require.Equal(t, importer.JobStepPreparing, gotJobInfo.Step) + require.False(t, gotJobInfo.StartTime.IsZero()) + startTime := gotJobInfo.StartTime + + task.Step = proto.StepPrepared + subtaskMetas, err := ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, proto.ImportStepImport) + require.NoError(t, err) + require.Len(t, subtaskMetas, 1) + gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) + require.Equal(t, importer.JobStepImporting, gotJobInfo.Step) + require.Equal(t, startTime, gotJobInfo.StartTime) + }) } func TestSchedulerOnDoneCancelResetsTableMode(t *testing.T) { diff --git a/pkg/executor/importer/job.go b/pkg/executor/importer/job.go index e7255367b5526..46dfdceb4ed9f 100644 --- a/pkg/executor/importer/job.go +++ b/pkg/executor/importer/job.go @@ -61,11 +61,18 @@ const ( // when the job is finished, step will be set to none. jobStepNone = "" + // JobStepPreparing is used by prepare-enabled jobs before generating + // first business-step subtasks. + JobStepPreparing = "preparing" // JobStepGlobalSorting is the first step when using global sort, // step goes from none -> global-sorting -> importing -> validating -> none. + // for prepare-enabled global sort, step goes from none -> preparing -> + // global-sorting -> importing -> validating -> none. JobStepGlobalSorting = "global-sorting" // JobStepImporting is the first step when using local sort, // step goes from none -> importing -> validating -> none. + // for prepare-enabled local sort, step goes from none -> preparing -> + // importing -> validating -> none. // when used in global sort, it means importing the sorted data. // when used in local sort, it means encode&sort data and then importing the data. JobStepImporting = "importing" diff --git a/pkg/executor/importer/job_test.go b/pkg/executor/importer/job_test.go index 384c515568791..9f87a8b8b2cc8 100644 --- a/pkg/executor/importer/job_test.go +++ b/pkg/executor/importer/job_test.go @@ -159,6 +159,27 @@ func TestJobHappyPath(t *testing.T) { require.NoError(t, err) require.Equal(t, endTime, gotJobInfo.EndTime) } + + t.Run("prepare phase transition", func(t *testing.T) { + jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, + "root@%", "", &importer.ImportParameters{Format: importer.DataFormatCSV}, 123) + require.NoError(t, err) + + require.NoError(t, importer.StartJob(ctx, conn, jobID, importer.JobStepPreparing)) + info, err := importer.GetJob(ctx, conn, jobID, "root@%", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, info.Status) + require.Equal(t, importer.JobStepPreparing, info.Step) + require.False(t, info.StartTime.IsZero()) + startTime := info.StartTime + + require.NoError(t, importer.Job2Step(ctx, conn, jobID, importer.JobStepImporting)) + info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, info.Status) + require.Equal(t, importer.JobStepImporting, info.Step) + require.Equal(t, startTime, info.StartTime) + }) } func TestGetAndCancelJob(t *testing.T) { From e42b73a77d602847723a1b7886bb85352c8f5789 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 10:15:18 +0800 Subject: [PATCH 03/21] executor: cover preparing phase in show import jobs --- pkg/executor/show_test.go | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/pkg/executor/show_test.go b/pkg/executor/show_test.go index 27c9eca026fed..045cecc159c56 100644 --- a/pkg/executor/show_test.go +++ b/pkg/executor/show_test.go @@ -107,6 +107,15 @@ func TestFillOneImportJobInfo(t *testing.T) { require.Equal(t, "12 conflicts", c.GetRow(5).GetString(fmap["CurStepProcessedSize"])) require.Equal(t, "34 conflicts", c.GetRow(5).GetString(fmap["CurStepTotalSize"])) require.Equal(t, "5 conflicts/s", c.GetRow(5).GetString(fmap["CurStepSpeed"])) + + // preparing phase should be visible while current business step is still init. + jobInfo.Step = importer.JobStepPreparing + ri = &importinto.RuntimeInfo{ + Step: proto.StepInit, + } + executor.FillOneImportJobInfo(c, jobInfo, ri) + require.Equal(t, importer.JobStepPreparing, c.GetRow(6).GetString(fmap["Phase"])) + require.Equal(t, proto.Step2Str(proto.ImportInto, proto.StepInit), c.GetRow(6).GetString(fmap["CurStep"])) } func TestShow(t *testing.T) { From 17e88369a31c2076acd37cb9f9b1b1f37b515222 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 16:13:56 +0800 Subject: [PATCH 04/21] dxf: align import scheduler prepare hook with extension API --- pkg/dxf/importinto/scheduler.go | 3 +-- pkg/dxf/importinto/scheduler_testkit_test.go | 4 +--- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index df3ba742f14c5..af959883bcb1c 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -167,7 +167,6 @@ type importScheduler struct { } var _ scheduler.Extension = (*importScheduler)(nil) -var _ scheduler.PrepareExtension = (*importScheduler)(nil) // NewImportScheduler creates a new import scheduler. func NewImportScheduler( @@ -339,7 +338,7 @@ func (sch *importScheduler) writePreparedChunkMap( return externalPath, nil } -// OnPrepare implements scheduler.PrepareExtension. +// OnPrepare implements scheduler.Extension. func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, task *proto.Task) error { taskMeta := &TaskMeta{} if err := json.Unmarshal(task.Meta, taskMeta); err != nil { diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index b0e66370c332e..200cdca9deb6c 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -220,9 +220,7 @@ func TestSchedulerExtLocalSort(t *testing.T) { ) require.NoError(t, err) - prepareExt, ok := ext.(scheduler.PrepareExtension) - require.True(t, ok) - require.NoError(t, prepareExt.OnPrepare(ctx, d, task)) + require.NoError(t, ext.OnPrepare(ctx, d, task)) gotJobInfo, err := importer.GetJob(ctx, conn, jobID, "root", true) require.NoError(t, err) From cb2566129d2c920992544c9878d0178b0f314531 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 17:44:19 +0800 Subject: [PATCH 05/21] importinto: refine async prepare flow and job metadata updates --- pkg/dxf/importinto/job.go | 10 +++-- pkg/dxf/importinto/job_testkit_test.go | 2 +- pkg/dxf/importinto/planner.go | 9 ++-- pkg/dxf/importinto/scheduler.go | 58 +++++++++++++------------- pkg/executor/import_into.go | 22 +++++----- pkg/executor/importer/job.go | 30 +++++++++++-- pkg/executor/importer/job_test.go | 12 ++++++ pkg/executor/importer/precheck.go | 16 ++++++- pkg/executor/importer/precheck_test.go | 7 ++++ 9 files changed, 108 insertions(+), 58 deletions(-) diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index 0488085e42b62..412bd0c0caa70 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -64,9 +64,11 @@ func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, * return doSubmitTask(ctx, plan, stmt, nil, nil) } -// ShouldUseScopedPrepareIntegration enables framework prepare integration only -// for nextgen + global-sort path. -func ShouldUseScopedPrepareIntegration(plan *importer.Plan) bool { +// ShouldUseAsyncPrepareForImportInto 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 ShouldUseAsyncPrepareForImportInto(plan *importer.Plan) bool { return plan != nil && kerneltype.IsNextGen() && plan.IsGlobalSort() } @@ -94,7 +96,7 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc ThreadCnt: plan.ThreadCnt, MaxNodeCnt: plan.MaxNodeCnt, } - if ShouldUseScopedPrepareIntegration(plan) { + if ShouldUseAsyncPrepareForImportInto(plan) { logicalPlan.PrepareMode = proto.PrepareModeRequired planCtx.ThreadCnt = 1 planCtx.MaxNodeCnt = 1 diff --git a/pkg/dxf/importinto/job_testkit_test.go b/pkg/dxf/importinto/job_testkit_test.go index 0e8ba31efdd7f..6fa2221092c9f 100644 --- a/pkg/dxf/importinto/job_testkit_test.go +++ b/pkg/dxf/importinto/job_testkit_test.go @@ -167,7 +167,7 @@ func TestSubmitTaskNextgen(t *testing.T) { 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 scoped prepare integration", func(t *testing.T) { + 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) { diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index 07af44dc851e5..e4f4a3d579cf9 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -425,13 +425,10 @@ func readPreparedChunkMap( return nil, err } defer store.Close() - data, err := store.ReadFile(ctx, externalPath) - if err != nil { - return nil, err - } + baseMeta := globalsort.BaseExternalMeta{ExternalPath: externalPath} var chunkMap map[int32][]importer.Chunk - if err := json.Unmarshal(data, &chunkMap); err != nil { - return nil, errors.Trace(err) + if err := baseMeta.ReadJSONFromExternalStorage(ctx, store, &chunkMap); err != nil { + return nil, err } return chunkMap, nil } diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index af959883bcb1c..be15b07a96d24 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -39,6 +39,7 @@ import ( "github.com/pingcap/tidb/pkg/dxf/framework/scheduler" "github.com/pingcap/tidb/pkg/dxf/framework/storage" "github.com/pingcap/tidb/pkg/executor/importer" + "github.com/pingcap/tidb/pkg/ingestor/globalsort" "github.com/pingcap/tidb/pkg/kv" "github.com/pingcap/tidb/pkg/lightning/common" "github.com/pingcap/tidb/pkg/lightning/config" @@ -291,26 +292,6 @@ func (sch *importScheduler) checkImportTableEmpty(ctx context.Context, taskMeta }) } -func (sch *importScheduler) checkActiveImportJobs(ctx context.Context, taskMeta *TaskMeta) error { - taskManager, err := sch.getTaskMgrForAccessingImportJob() - if err != nil { - return err - } - return taskManager.WithNewSession(func(se sessionctx.Context) error { - exec := se.GetSQLExecutor() - activeCnt, err := importer.GetActiveJobCnt(ctx, exec, taskMeta.Plan.DBName, taskMeta.Plan.TableInfo.Name.L) - if err != nil { - return err - } - // The current prepare-enabled job is already in pending state, so we only - // fail when there are other active jobs on the same table. - if activeCnt > 1 { - return exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs("there is active job on the target table already") - } - return nil - }) -} - func (sch *importScheduler) writePreparedChunkMap( ctx context.Context, taskID int64, @@ -328,11 +309,8 @@ func (sch *importScheduler) writePreparedChunkMap( uuid.NewString(), "meta.json", ) - data, err := json.Marshal(chunkMap) - if err != nil { - return "", errors.Trace(err) - } - if err = store.WriteFile(ctx, externalPath, data); err != nil { + baseMeta := globalsort.BaseExternalMeta{ExternalPath: externalPath} + if err = baseMeta.WriteJSONToExternalStorage(ctx, store, chunkMap); err != nil { return "", err } return externalPath, nil @@ -349,7 +327,7 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, return err } } - if !ShouldUseScopedPrepareIntegration(&taskMeta.Plan) { + if !ShouldUseAsyncPrepareForImportInto(&taskMeta.Plan) { return nil } @@ -372,9 +350,6 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, "No file matched, or the file is empty. Please provide a valid file location.", ) } - if err = sch.checkActiveImportJobs(ctx, taskMeta); err != nil { - return err - } if err = sch.checkImportTableEmpty(ctx, taskMeta); err != nil { return err } @@ -394,6 +369,9 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, taskMeta.Plan = *controller.Plan taskMeta.ChunkMap = nil taskMeta.PreparedChunkMapExternalPath = chunkMapPath + if err = sch.updatePreparedJobInfo(ctx, sch.GetLogger(), taskMeta); err != nil { + return err + } metaBytes, err := json.Marshal(taskMeta) if err != nil { return errors.Trace(err) @@ -819,6 +797,28 @@ func (sch *importScheduler) job2Step(ctx context.Context, logger *zap.Logger, ta ) } +func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *zap.Logger, taskMeta *TaskMeta) error { + taskManager, err := sch.getTaskMgrForAccessingImportJob() + if err != nil { + return err + } + backoffer := backoff.NewExponential(scheduler.RetrySQLInterval, 2, scheduler.RetrySQLMaxInterval) + return handle.RunWithRetry(ctx, scheduler.RetrySQLTimes, backoffer, logger, + func(ctx context.Context) (bool, error) { + return true, taskManager.WithNewSession(func(se sessionctx.Context) error { + exec := se.GetSQLExecutor() + return importer.UpdateJobPreparedInfo( + ctx, + exec, + taskMeta.JobID, + taskMeta.Plan.TotalFileSize, + taskMeta.Plan.Parameters, + ) + }) + }, + ) +} + func (sch *importScheduler) finishJob(ctx context.Context, logger *zap.Logger, task *proto.Task, taskMeta *TaskMeta) error { // we have already switched import-mode when switch to post-process step. diff --git a/pkg/executor/import_into.go b/pkg/executor/import_into.go index 72c89c842e357..1548974e1dc35 100644 --- a/pkg/executor/import_into.go +++ b/pkg/executor/import_into.go @@ -106,8 +106,8 @@ func (e *ImportIntoExec) Next(ctx context.Context, req *chunk.Chunk) (err error) return e.importFromSelect(ctx) } - useScopedPrepareIntegration := importinto.ShouldUseScopedPrepareIntegration(e.controller.Plan) - if !useScopedPrepareIntegration { + useAsyncPrepare := importinto.ShouldUseAsyncPrepareForImportInto(e.controller.Plan) + if !useAsyncPrepare { if err2 := e.controller.InitDataFiles(ctx); err2 != nil { return err2 } @@ -125,10 +125,12 @@ func (e *ImportIntoExec) Next(ctx context.Context, req *chunk.Chunk) (err error) return err2 } defer CloseSession(newSCtx) - if !useScopedPrepareIntegration { - if err2 = e.controller.CheckRequirements(ctx, newSCtx); err2 != nil { + if useAsyncPrepare { + if err2 = e.controller.CheckRequirementsBeforeInitDataFiles(ctx, newSCtx); err2 != nil { return err2 } + } else if err2 = e.controller.CheckRequirements(ctx, newSCtx); err2 != nil { + return err2 } if err := e.controller.InitTiKVConfigs(ctx, newSCtx); err != nil { @@ -222,16 +224,12 @@ func (e *ImportIntoExec) submitTask(ctx context.Context) (int64, *proto.TaskBase } logutil.Logger(ctx).Info("get job importer", zap.Stringer("param", e.controller.Parameters), zap.Bool("dist-task-enabled", vardef.EnableDistTask.Load())) - useScopedPrepareIntegration := importinto.ShouldUseScopedPrepareIntegration(e.controller.Plan) if importFromServer { - if !useScopedPrepareIntegration { - chunkMap, err2 := e.controller.PopulateChunks(ctx) - if err2 != nil { - return 0, nil, err2 - } - return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap) + chunkMap, err2 := e.controller.PopulateChunks(ctx) + if err2 != nil { + return 0, nil, err2 } - return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, nil) + return importinto.SubmitStandaloneTask(ctx, e.controller.Plan, e.stmt, chunkMap) } // if tidb_enable_dist_task=true, we import distributively, otherwise we import on current node. if vardef.EnableDistTask.Load() { diff --git a/pkg/executor/importer/job.go b/pkg/executor/importer/job.go index 46dfdceb4ed9f..1ea4a57062286 100644 --- a/pkg/executor/importer/job.go +++ b/pkg/executor/importer/job.go @@ -71,8 +71,6 @@ const ( JobStepGlobalSorting = "global-sorting" // JobStepImporting is the first step when using local sort, // step goes from none -> importing -> validating -> none. - // for prepare-enabled local sort, step goes from none -> preparing -> - // importing -> validating -> none. // when used in global sort, it means importing the sorted data. // when used in local sort, it means encode&sort data and then importing the data. JobStepImporting = "importing" @@ -260,13 +258,37 @@ func StartJob(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64, step s func Job2Step(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64, step string) error { ctx = util.WithInternalSourceType(ctx, kv.InternalImportInto) _, err := conn.ExecuteInternal(ctx, `UPDATE mysql.tidb_import_jobs - SET update_time = CURRENT_TIMESTAMP(6), step = %? - WHERE id = %? AND status = %?;`, + SET update_time = CURRENT_TIMESTAMP(6), step = %? + WHERE id = %? AND status = %?;`, step, jobID, JobStatusRunning) return err } +// UpdateJobPreparedInfo updates import job fields that are known only after +// async prepare succeeds. +func UpdateJobPreparedInfo( + ctx context.Context, + conn sqlexec.SQLExecutor, + jobID int64, + sourceFileSize int64, + parameters *ImportParameters, +) error { + if parameters == nil { + parameters = &ImportParameters{} + } + paramsBytes, err := json.Marshal(parameters) + if err != nil { + return err + } + ctx = util.WithInternalSourceType(ctx, kv.InternalImportInto) + _, err = conn.ExecuteInternal(ctx, `UPDATE mysql.tidb_import_jobs + SET update_time = CURRENT_TIMESTAMP(6), source_file_size = %?, parameters = %? + WHERE id = %? AND status = %?;`, + sourceFileSize, paramsBytes, jobID, JobStatusRunning) + return err +} + // FinishJob tries to finish a running job with jobID, change its status to finished, clear its step and update summary. // It will not return error when there's no matched job. func FinishJob(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64, summary *Summary) error { diff --git a/pkg/executor/importer/job_test.go b/pkg/executor/importer/job_test.go index 9f87a8b8b2cc8..928373597b7ae 100644 --- a/pkg/executor/importer/job_test.go +++ b/pkg/executor/importer/job_test.go @@ -179,6 +179,18 @@ func TestJobHappyPath(t *testing.T) { require.Equal(t, importer.JobStatusRunning, info.Status) require.Equal(t, importer.JobStepImporting, info.Step) require.Equal(t, startTime, info.StartTime) + + updatedParams := &importer.ImportParameters{ + Format: importer.DataFormatCSV, + Options: map[string]any{ + "thread": float64(16), + }, + } + require.NoError(t, importer.UpdateJobPreparedInfo(ctx, conn, jobID, 456, updatedParams)) + info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) + require.NoError(t, err) + require.EqualValues(t, 456, info.SourceFileSize) + require.Equal(t, updatedParams.Options, info.Parameters.Options) }) } diff --git a/pkg/executor/importer/precheck.go b/pkg/executor/importer/precheck.go index cc84c1b3618ce..bbe18eff578c3 100644 --- a/pkg/executor/importer/precheck.go +++ b/pkg/executor/importer/precheck.go @@ -47,6 +47,16 @@ var GetEtcdClient = store.NewEtcdCli // // we check them one by one, and return the first error we meet. func (e *LoadDataController) CheckRequirements(ctx context.Context, se sessionctx.Context) error { + return e.checkRequirements(ctx, se, true) +} + +// CheckRequirementsBeforeInitDataFiles checks requirements that don't depend on +// discovered data files, and is used by async-prepare submit path. +func (e *LoadDataController) CheckRequirementsBeforeInitDataFiles(ctx context.Context, se sessionctx.Context) error { + return e.checkRequirements(ctx, se, false) +} + +func (e *LoadDataController) checkRequirements(ctx context.Context, se sessionctx.Context, checkTotalFileSize bool) error { conn := se.GetSQLExecutor() if e.DataSourceType == DataSourceTypeFile { cnt, err := GetActiveJobCnt(ctx, conn, e.Plan.DBName, e.Plan.TableInfo.Name.L) @@ -56,8 +66,10 @@ func (e *LoadDataController) CheckRequirements(ctx context.Context, se sessionct if cnt > 0 { return exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs("there is active job on the target table already") } - if err := e.checkTotalFileSize(); err != nil { - return err + if checkTotalFileSize { + if err := e.checkTotalFileSize(); err != nil { + return err + } } } if err := e.checkTableEmpty(ctx, conn); err != nil { diff --git a/pkg/executor/importer/precheck_test.go b/pkg/executor/importer/precheck_test.go index 44cca15fe6f00..4fb8d919e6085 100644 --- a/pkg/executor/importer/precheck_test.go +++ b/pkg/executor/importer/precheck_test.go @@ -136,9 +136,16 @@ func TestCheckRequirements(t *testing.T) { err = c.CheckRequirements(ctx, tk.Session()) require.ErrorIs(t, err, exeerrors.ErrLoadDataPreCheckFailed) require.ErrorContains(t, err, "there is active job on the target table already") + err = c.CheckRequirementsBeforeInitDataFiles(ctx, tk.Session()) + require.ErrorIs(t, err, exeerrors.ErrLoadDataPreCheckFailed) + require.ErrorContains(t, err, "there is active job on the target table already") // cancel the job require.NoError(t, importer.CancelJob(ctx, conn, jobID)) + // async-prepare submit path skips file-size check before InitDataFiles. + c.DisablePrecheck = true + require.NoError(t, c.CheckRequirementsBeforeInitDataFiles(ctx, tk.Session())) + c.DisablePrecheck = false // source data file size = 0 require.ErrorIs(t, c.CheckRequirements(ctx, tk.Session()), exeerrors.ErrLoadDataPreCheckFailed) From 22bbf3f27a3bbc0b0ab7b5b2e6acab44f795cfd6 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 19:33:33 +0800 Subject: [PATCH 06/21] importinto: simplify prepare hook and generic chunk-map meta --- pkg/dxf/framework/scheduler/scheduler.go | 12 ++++--- pkg/dxf/importinto/job.go | 6 ++-- pkg/dxf/importinto/planner.go | 22 ++++++++++--- pkg/dxf/importinto/planner_test.go | 34 ++++++++++++++++++-- pkg/dxf/importinto/proto.go | 7 ++++ pkg/dxf/importinto/scheduler.go | 33 +++++++++---------- pkg/dxf/importinto/scheduler_testkit_test.go | 4 +++ pkg/executor/import_into.go | 2 +- 8 files changed, 87 insertions(+), 33 deletions(-) diff --git a/pkg/dxf/framework/scheduler/scheduler.go b/pkg/dxf/framework/scheduler/scheduler.go index b4d1668c5330b..568b0d640b600 100644 --- a/pkg/dxf/framework/scheduler/scheduler.go +++ b/pkg/dxf/framework/scheduler/scheduler.go @@ -151,10 +151,12 @@ 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 +// OnPrepare implements Extension with default delegation to Extension. +func (s *BaseScheduler) OnPrepare(ctx context.Context, h storage.TaskHandle, task *proto.Task) error { + if s.Extension == nil || s.Extension == s { + return nil + } + return s.Extension.OnPrepare(ctx, h, task) } // getTaskClone returns a clone of the task. @@ -392,7 +394,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) diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index 412bd0c0caa70..c33e802562d57 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -64,11 +64,11 @@ func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, * return doSubmitTask(ctx, plan, stmt, nil, nil) } -// ShouldUseAsyncPrepareForImportInto returns whether IMPORT INTO should use +// 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 ShouldUseAsyncPrepareForImportInto(plan *importer.Plan) bool { +func ShouldUseAsyncPrepare(plan *importer.Plan) bool { return plan != nil && kerneltype.IsNextGen() && plan.IsGlobalSort() } @@ -96,7 +96,7 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc ThreadCnt: plan.ThreadCnt, MaxNodeCnt: plan.MaxNodeCnt, } - if ShouldUseAsyncPrepareForImportInto(plan) { + if ShouldUseAsyncPrepare(plan) { logicalPlan.PrepareMode = proto.PrepareModeRequired planCtx.ThreadCnt = 1 planCtx.MaxNodeCnt = 1 diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index e4f4a3d579cf9..198a794d07ea9 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -425,12 +425,26 @@ func readPreparedChunkMap( return nil, err } defer store.Close() - baseMeta := globalsort.BaseExternalMeta{ExternalPath: externalPath} - var chunkMap map[int32][]importer.Chunk - if err := baseMeta.ReadJSONFromExternalStorage(ctx, store, &chunkMap); err != nil { + data, err := store.ReadFile(ctx, externalPath) + if err != nil { + return nil, err + } + + var preparedChunkMapMeta PreparedChunkMapMeta + if err := json.Unmarshal(data, &preparedChunkMapMeta); err != nil { + return nil, err + } + if preparedChunkMapMeta.ChunkMap != nil { + return preparedChunkMapMeta.ChunkMap, nil + } + + // Compatible with legacy format where the external file stored chunkMap + // directly as the top-level JSON object. + var legacyChunkMap map[int32][]importer.Chunk + if err := json.Unmarshal(data, &legacyChunkMap); err != nil { return nil, err } - return chunkMap, nil + return legacyChunkMap, nil } func skipMergeSort(kvGroup string, stats []simplesst.MultipleFilesStat, concurrency int) bool { diff --git a/pkg/dxf/importinto/planner_test.go b/pkg/dxf/importinto/planner_test.go index 21c34a03312ed..dcc10838162d8 100644 --- a/pkg/dxf/importinto/planner_test.go +++ b/pkg/dxf/importinto/planner_test.go @@ -150,8 +150,10 @@ func TestToPhysicalPlan(t *testing.T) { preparedChunkMap := map[int32][]importer.Chunk{ 1: {{Path: "gs://test-load/2.csv"}}, } - externalPath := "100/prepare-chunk-map/mock-attempt/meta.json" - data, err := json.Marshal(preparedChunkMap) + externalPath := globalsort.PlanMetaPath(100, "prepare-chunk-map/mock-attempt", 1) + data, err := json.Marshal(PreparedChunkMapMeta{ + ChunkMap: preparedChunkMap, + }) require.NoError(t, err) require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) @@ -171,6 +173,34 @@ func TestToPhysicalPlan(t *testing.T) { require.Len(t, importSpec.Chunks, 1) require.Equal(t, "gs://test-load/2.csv", importSpec.Chunks[0].Path) }) + + t.Run("legacy prepared chunk map format is still readable", 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/legacy.csv"}}, + } + externalPath := globalsort.PlanMetaPath(101, "prepare-chunk-map/mock-attempt", 1) + data, err := json.Marshal(preparedChunkMap) + require.NoError(t, err) + require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) + + specs, err := generateImportSpecs(planner.PlanCtx{Ctx: context.Background()}, &LogicalPlan{ + Plan: importer.Plan{ + CloudStorageURI: cloudStorageURI, + }, + 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/legacy.csv", importSpec.Chunks[0].Path) + }) } func genEncodeStepMetas(t *testing.T, cnt int) [][]byte { diff --git a/pkg/dxf/importinto/proto.go b/pkg/dxf/importinto/proto.go index 6264d7ed4606e..d37d946e21786 100644 --- a/pkg/dxf/importinto/proto.go +++ b/pkg/dxf/importinto/proto.go @@ -58,6 +58,13 @@ type TaskMeta struct { PreparedChunkMapExternalPath string `json:"prepared_chunk_map_external_path,omitempty"` } +// PreparedChunkMapMeta stores chunk-map metadata generated in prepare stage. +// Keep this wrapper generic so future prepare-stage fields can be added without +// changing the on-disk top-level JSON shape again. +type PreparedChunkMapMeta struct { + ChunkMap map[int32][]importer.Chunk `json:"chunk_map,omitempty"` +} + // ImportStepMeta is the meta of import step. // Scheduler will split the task into subtasks(FileInfos -> Chunks) // All the field should be serializable. diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index be15b07a96d24..c4d62318439f2 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -65,11 +65,12 @@ const ( // the target table, it's known to be slow to import in this case. // the value if chosen as most tables have less than 32 indexes, we can adjust // it later if needed. - warningIndexCount = 32 - registerTaskTTL = 10 * time.Minute - refreshTaskTTLInterval = 3 * time.Minute - registerTimeout = 5 * time.Second - preparedChunkMetaPathPurpose = "prepare-chunk-map" + warningIndexCount = 32 + registerTaskTTL = 10 * time.Minute + refreshTaskTTLInterval = 3 * time.Minute + registerTimeout = 5 * time.Second + preparedChunkMapPathStep = "prepare-chunk-map" + preparedChunkMapMetaIndex = 1 ) var ( @@ -303,14 +304,15 @@ func (sch *importScheduler) writePreparedChunkMap( return "", err } defer store.Close() - externalPath := path.Join( - strconv.FormatInt(taskID, 10), - preparedChunkMetaPathPurpose, - uuid.NewString(), - "meta.json", + externalPath := globalsort.PlanMetaPath( + taskID, + path.Join(preparedChunkMapPathStep, uuid.NewString()), + preparedChunkMapMetaIndex, ) baseMeta := globalsort.BaseExternalMeta{ExternalPath: externalPath} - if err = baseMeta.WriteJSONToExternalStorage(ctx, store, chunkMap); err != nil { + if err = baseMeta.WriteJSONToExternalStorage(ctx, store, PreparedChunkMapMeta{ + ChunkMap: chunkMap, + }); err != nil { return "", err } return externalPath, nil @@ -322,13 +324,8 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, if err := json.Unmarshal(task.Meta, taskMeta); err != nil { return errors.Annotate(err, "unmarshal task meta failed") } - if task.ExtraParams.PrepareMode == proto.PrepareModeRequired { - if err := sch.startJob(ctx, sch.GetLogger(), taskMeta, importer.JobStepPreparing); err != nil { - return err - } - } - if !ShouldUseAsyncPrepareForImportInto(&taskMeta.Plan) { - return nil + if err := sch.startJob(ctx, sch.GetLogger(), taskMeta, importer.JobStepPreparing); err != nil { + return err } logicalPlan := &LogicalPlan{ diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index 200cdca9deb6c..92a92ab1d0033 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -196,6 +196,10 @@ func TestSchedulerExtLocalSort(t *testing.T) { require.Equal(t, "cancelled", gotJobInfo.Status) t.Run("prepare-enabled job transitions from preparing to first business phase", func(t *testing.T) { + if !importinto.ShouldUseAsyncPrepare(&logicalPlan.Plan) { + t.Skip("prepare mode only applies when async prepare is enabled") + } + jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, "root", "", &importer.ImportParameters{}, 123) require.NoError(t, err) diff --git a/pkg/executor/import_into.go b/pkg/executor/import_into.go index 1548974e1dc35..2e1753774ef68 100644 --- a/pkg/executor/import_into.go +++ b/pkg/executor/import_into.go @@ -106,7 +106,7 @@ func (e *ImportIntoExec) Next(ctx context.Context, req *chunk.Chunk) (err error) return e.importFromSelect(ctx) } - useAsyncPrepare := importinto.ShouldUseAsyncPrepareForImportInto(e.controller.Plan) + useAsyncPrepare := importinto.ShouldUseAsyncPrepare(e.controller.Plan) if !useAsyncPrepare { if err2 := e.controller.InitDataFiles(ctx); err2 != nil { return err2 From a815e88ed28baf9ed95b23478b2044a54ffa58a0 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 20:02:11 +0800 Subject: [PATCH 07/21] dxf: simplify prepare extension and prepared meta path --- pkg/dxf/framework/scheduler/scheduler.go | 8 ------- pkg/dxf/importinto/planner.go | 11 ++++------ pkg/dxf/importinto/planner_test.go | 14 ++++++------ pkg/dxf/importinto/proto.go | 8 ++++++- pkg/dxf/importinto/scheduler.go | 28 ++++++++++-------------- pkg/ingestor/globalsort/util.go | 5 +++++ pkg/ingestor/globalsort/util_test.go | 2 ++ 7 files changed, 36 insertions(+), 40 deletions(-) diff --git a/pkg/dxf/framework/scheduler/scheduler.go b/pkg/dxf/framework/scheduler/scheduler.go index 568b0d640b600..7ffdc7f506665 100644 --- a/pkg/dxf/framework/scheduler/scheduler.go +++ b/pkg/dxf/framework/scheduler/scheduler.go @@ -151,14 +151,6 @@ func (s *BaseScheduler) GetTask() *proto.Task { return s.task.Load() } -// OnPrepare implements Extension with default delegation to Extension. -func (s *BaseScheduler) OnPrepare(ctx context.Context, h storage.TaskHandle, task *proto.Task) error { - if s.Extension == nil || s.Extension == s { - return nil - } - return s.Extension.OnPrepare(ctx, h, task) -} - // getTaskClone returns a clone of the task. func (s *BaseScheduler) getTaskClone() *proto.Task { clone := *s.GetTask() diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index 198a794d07ea9..7a26ca9d29ba6 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -425,13 +425,10 @@ func readPreparedChunkMap( return nil, err } defer store.Close() - data, err := store.ReadFile(ctx, externalPath) - if err != nil { - return nil, err + preparedChunkMapMeta := PreparedChunkMapMeta{ + BaseExternalMeta: globalsort.BaseExternalMeta{ExternalPath: externalPath}, } - - var preparedChunkMapMeta PreparedChunkMapMeta - if err := json.Unmarshal(data, &preparedChunkMapMeta); err != nil { + if err := preparedChunkMapMeta.ReadJSONFromExternalStorage(ctx, store, &preparedChunkMapMeta); err != nil { return nil, err } if preparedChunkMapMeta.ChunkMap != nil { @@ -441,7 +438,7 @@ func readPreparedChunkMap( // Compatible with legacy format where the external file stored chunkMap // directly as the top-level JSON object. var legacyChunkMap map[int32][]importer.Chunk - if err := json.Unmarshal(data, &legacyChunkMap); err != nil { + if err := preparedChunkMapMeta.ReadJSONFromExternalStorage(ctx, store, &legacyChunkMap); err != nil { return nil, err } return legacyChunkMap, nil diff --git a/pkg/dxf/importinto/planner_test.go b/pkg/dxf/importinto/planner_test.go index dcc10838162d8..5cb94c8bd13e8 100644 --- a/pkg/dxf/importinto/planner_test.go +++ b/pkg/dxf/importinto/planner_test.go @@ -150,12 +150,12 @@ func TestToPhysicalPlan(t *testing.T) { preparedChunkMap := map[int32][]importer.Chunk{ 1: {{Path: "gs://test-load/2.csv"}}, } - externalPath := globalsort.PlanMetaPath(100, "prepare-chunk-map/mock-attempt", 1) - data, err := json.Marshal(PreparedChunkMapMeta{ - ChunkMap: preparedChunkMap, - }) - require.NoError(t, err) - require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) + externalPath := globalsort.PreparedMetaPath(100) + preparedMeta := PreparedChunkMapMeta{ + 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{ @@ -183,7 +183,7 @@ func TestToPhysicalPlan(t *testing.T) { preparedChunkMap := map[int32][]importer.Chunk{ 1: {{Path: "gs://test-load/legacy.csv"}}, } - externalPath := globalsort.PlanMetaPath(101, "prepare-chunk-map/mock-attempt", 1) + externalPath := globalsort.PreparedMetaPath(101) data, err := json.Marshal(preparedChunkMap) require.NoError(t, err) require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) diff --git a/pkg/dxf/importinto/proto.go b/pkg/dxf/importinto/proto.go index d37d946e21786..de6ea7854e45d 100644 --- a/pkg/dxf/importinto/proto.go +++ b/pkg/dxf/importinto/proto.go @@ -62,7 +62,13 @@ type TaskMeta struct { // Keep this wrapper generic so future prepare-stage fields can be added without // changing the on-disk top-level JSON shape again. type PreparedChunkMapMeta struct { - ChunkMap map[int32][]importer.Chunk `json:"chunk_map,omitempty"` + globalsort.BaseExternalMeta + ChunkMap map[int32][]importer.Chunk `json:"chunk_map,omitempty" external:"true"` +} + +// Marshal marshals the prepared chunk map meta to JSON. +func (m *PreparedChunkMapMeta) Marshal() ([]byte, error) { + return m.BaseExternalMeta.Marshal(m) } // ImportStepMeta is the meta of import step. diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index c4d62318439f2..8ddbf0c96d033 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -19,12 +19,10 @@ import ( "encoding/json" goerrors "errors" "fmt" - "path" "strconv" "sync" "time" - "github.com/google/uuid" "github.com/pingcap/errors" "github.com/pingcap/failpoint" "github.com/pingcap/tidb/br/pkg/utils" @@ -65,12 +63,10 @@ const ( // the target table, it's known to be slow to import in this case. // the value if chosen as most tables have less than 32 indexes, we can adjust // it later if needed. - warningIndexCount = 32 - registerTaskTTL = 10 * time.Minute - refreshTaskTTLInterval = 3 * time.Minute - registerTimeout = 5 * time.Second - preparedChunkMapPathStep = "prepare-chunk-map" - preparedChunkMapMetaIndex = 1 + warningIndexCount = 32 + registerTaskTTL = 10 * time.Minute + refreshTaskTTLInterval = 3 * time.Minute + registerTimeout = 5 * time.Second ) var ( @@ -304,18 +300,16 @@ func (sch *importScheduler) writePreparedChunkMap( return "", err } defer store.Close() - externalPath := globalsort.PlanMetaPath( - taskID, - path.Join(preparedChunkMapPathStep, uuid.NewString()), - preparedChunkMapMetaIndex, - ) - baseMeta := globalsort.BaseExternalMeta{ExternalPath: externalPath} - if err = baseMeta.WriteJSONToExternalStorage(ctx, store, PreparedChunkMapMeta{ + preparedMeta := PreparedChunkMapMeta{ + BaseExternalMeta: globalsort.BaseExternalMeta{ + ExternalPath: globalsort.PreparedMetaPath(taskID), + }, ChunkMap: chunkMap, - }); err != nil { + } + if err = preparedMeta.WriteJSONToExternalStorage(ctx, store, preparedMeta); err != nil { return "", err } - return externalPath, nil + return preparedMeta.ExternalPath, nil } // OnPrepare implements scheduler.Extension. diff --git a/pkg/ingestor/globalsort/util.go b/pkg/ingestor/globalsort/util.go index 7c817ae59b849..ae6f913df0d67 100644 --- a/pkg/ingestor/globalsort/util.go +++ b/pkg/ingestor/globalsort/util.go @@ -274,6 +274,11 @@ func PlanMetaPath(taskID int64, step string, idx int) string { return path.Join(strconv.FormatInt(taskID, 10), "plan", step, strconv.Itoa(idx), metaName) } +// PreparedMetaPath returns the path of the prepared meta file. +func PreparedMetaPath(taskID int64) string { + return path.Join(strconv.FormatInt(taskID, 10), "plan", "prepared", metaName) +} + // SubtaskMetaPath returns the path of the subtask meta file. func SubtaskMetaPath(taskID int64, subtaskID int64) string { return path.Join(strconv.FormatInt(taskID, 10), strconv.FormatInt(subtaskID, 10), metaName) diff --git a/pkg/ingestor/globalsort/util_test.go b/pkg/ingestor/globalsort/util_test.go index df920a6a2a8f4..fe58d4a105e2b 100644 --- a/pkg/ingestor/globalsort/util_test.go +++ b/pkg/ingestor/globalsort/util_test.go @@ -375,6 +375,8 @@ func TestReadWriteJSON(t *testing.T) { func TestExternalMetaPath(t *testing.T) { require.Equal(t, "1/plan/merge-sort/1/meta.json", PlanMetaPath(1, "merge-sort", 1)) require.Equal(t, "2/plan/ingest/3/meta.json", PlanMetaPath(2, "ingest", 3)) + require.Equal(t, "1/plan/prepared/meta.json", PreparedMetaPath(1)) + require.Equal(t, "2/plan/prepared/meta.json", PreparedMetaPath(2)) require.Equal(t, "1/1/meta.json", SubtaskMetaPath(1, 1)) require.Equal(t, "2/3/meta.json", SubtaskMetaPath(2, 3)) From b6620fb8e8c31e27f527de41d29820016acec359 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 21 May 2026 20:25:06 +0800 Subject: [PATCH 08/21] importinto: merge detected format into prepared job params --- pkg/dxf/importinto/scheduler.go | 2 +- pkg/executor/importer/import.go | 6 +++++- pkg/executor/importer/job.go | 36 +++++++++++++++++++++++++------ pkg/executor/importer/job_test.go | 22 +++++++++++-------- 4 files changed, 49 insertions(+), 17 deletions(-) diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 8ddbf0c96d033..287644447f496 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -803,7 +803,7 @@ func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *z exec, taskMeta.JobID, taskMeta.Plan.TotalFileSize, - taskMeta.Plan.Parameters, + taskMeta.Plan.Format, ) }) }, diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 044fbf5d197e5..c0f21158898d8 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -1643,7 +1643,11 @@ func (e *LoadDataController) detectAndUpdateFormat(path string) { e.Format = parseFileType(path) e.logger.Info("detect and update import plan format based on file extension", zap.String("file", path), zap.String("detected format", e.Format)) - e.Parameters.Format = e.Format + // Plan.Parameters is not persisted in dist-task task meta. Keep this + // in-memory update best-effort for non-prepare paths. + if e.Parameters != nil { + e.Parameters.Format = e.Format + } } } diff --git a/pkg/executor/importer/job.go b/pkg/executor/importer/job.go index 1ea4a57062286..6495cff8e5f95 100644 --- a/pkg/executor/importer/job.go +++ b/pkg/executor/importer/job.go @@ -272,19 +272,43 @@ func UpdateJobPreparedInfo( conn sqlexec.SQLExecutor, jobID int64, sourceFileSize int64, - parameters *ImportParameters, + format string, ) error { - if parameters == nil { - parameters = &ImportParameters{} + ctx = util.WithInternalSourceType(ctx, kv.InternalImportInto) + rs, err := conn.ExecuteInternal(ctx, `SELECT parameters FROM mysql.tidb_import_jobs + WHERE id = %? AND status = %?;`, + jobID, JobStatusRunning) + if err != nil { + return err + } + defer terror.Call(rs.Close) + rows, err := sqlexec.DrainRecordSet(ctx, rs, 1) + if err != nil { + return err } + // no matched running job, keep historical behavior: no-op. + if len(rows) == 0 { + return nil + } + + parameters := &ImportParameters{} + parametersStr := rows[0].GetString(0) + if len(parametersStr) > 0 { + if err = json.Unmarshal([]byte(parametersStr), parameters); err != nil { + return err + } + } + if format != "" { + parameters.Format = format + } + paramsBytes, err := json.Marshal(parameters) if err != nil { return err } - ctx = util.WithInternalSourceType(ctx, kv.InternalImportInto) _, err = conn.ExecuteInternal(ctx, `UPDATE mysql.tidb_import_jobs - SET update_time = CURRENT_TIMESTAMP(6), source_file_size = %?, parameters = %? - WHERE id = %? AND status = %?;`, + SET update_time = CURRENT_TIMESTAMP(6), source_file_size = %?, parameters = %? + WHERE id = %? AND status = %?;`, sourceFileSize, paramsBytes, jobID, JobStatusRunning) return err } diff --git a/pkg/executor/importer/job_test.go b/pkg/executor/importer/job_test.go index 928373597b7ae..d4624ffc444fb 100644 --- a/pkg/executor/importer/job_test.go +++ b/pkg/executor/importer/job_test.go @@ -161,8 +161,16 @@ func TestJobHappyPath(t *testing.T) { } t.Run("prepare phase transition", func(t *testing.T) { + createdParams := &importer.ImportParameters{ + FileLocation: "s3://bucket/path.csv", + Format: importer.DataFormatAuto, + Options: map[string]any{ + "detached": nil, + "thread": float64(16), + }, + } jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, - "root@%", "", &importer.ImportParameters{Format: importer.DataFormatCSV}, 123) + "root@%", "", createdParams, 123) require.NoError(t, err) require.NoError(t, importer.StartJob(ctx, conn, jobID, importer.JobStepPreparing)) @@ -180,17 +188,13 @@ func TestJobHappyPath(t *testing.T) { require.Equal(t, importer.JobStepImporting, info.Step) require.Equal(t, startTime, info.StartTime) - updatedParams := &importer.ImportParameters{ - Format: importer.DataFormatCSV, - Options: map[string]any{ - "thread": float64(16), - }, - } - require.NoError(t, importer.UpdateJobPreparedInfo(ctx, conn, jobID, 456, updatedParams)) + require.NoError(t, importer.UpdateJobPreparedInfo(ctx, conn, jobID, 456, importer.DataFormatCSV)) info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) require.NoError(t, err) require.EqualValues(t, 456, info.SourceFileSize) - require.Equal(t, updatedParams.Options, info.Parameters.Options) + require.Equal(t, importer.DataFormatCSV, info.Parameters.Format) + require.Equal(t, createdParams.FileLocation, info.Parameters.FileLocation) + require.Equal(t, createdParams.Options, info.Parameters.Options) }) } From 7a42bdfd9b37f2eb5558686ed70ed54f6faac852 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Fri, 22 May 2026 01:02:58 +0800 Subject: [PATCH 09/21] realtikvtest: cover show import job timing around prepare --- tests/realtikvtest/importintotest/job_test.go | 133 ++++++++++++++++++ 1 file changed, 133 insertions(+) diff --git a/tests/realtikvtest/importintotest/job_test.go b/tests/realtikvtest/importintotest/job_test.go index 76ddb5915b1cb..6d6bdc030788f 100644 --- a/tests/realtikvtest/importintotest/job_test.go +++ b/tests/realtikvtest/importintotest/job_test.go @@ -29,6 +29,7 @@ import ( "github.com/pingcap/errors" "github.com/pingcap/failpoint" "github.com/pingcap/tidb/br/pkg/utils" + "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/dxf/framework/storage" "github.com/pingcap/tidb/pkg/dxf/importinto" @@ -368,6 +369,138 @@ func (s *mockGCSSuite) TestShowDetachedJob() { s.compareJobInfoWithoutTime(jobInfo, rows[0]) } +func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { + if !kerneltype.IsNextGen() { + s.T().Skip("async prepare is only enabled in nextgen kernel") + } + + s.prepareAndUseDB("show_job_prepare_timing") + s.tk.MustExec("CREATE TABLE t (i INT PRIMARY KEY);") + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "show-job-prepare-timing-source", + Name: "t.csv", + }, + Content: []byte("1\n2"), + }) + // Ensure the sort bucket exists before scheduler writes prepared metadata. + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "show-job-prepare-timing-sort", + Name: "seed", + }, + Content: []byte("seed"), + }) + + do, err := session.GetDomain(s.store) + s.NoError(err) + tableID := do.MustGetTableID(s.T(), "show_job_prepare_timing", "t") + + // Pause scheduler manager so we can assert show import job before scheduler starts. + testfailpoint.Enable(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeGetSchedulableTasks", "pause") + + importSQL := fmt.Sprintf( + `IMPORT INTO t FROM 'gs://show-job-prepare-timing-source/t.csv?endpoint=%s' + WITH DETACHED, cloud_storage_uri='gs://show-job-prepare-timing-sort/path?endpoint=%s'`, + gcsEndpoint, + gcsEndpoint, + ) + result := s.tk.MustQuery(importSQL).Rows() + s.Len(result, 1) + jobID, err := strconv.Atoi(result[0][0].(string)) + s.NoError(err) + jobID64 := int64(jobID) + + rows := s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + s.Len(rows, 1) + // timing 1: before scheduler starts. + s.Equal(result, rows) + + reachedBeforePrepare := make(chan struct{}) + releaseBeforePrepare := make(chan struct{}) + var beforePrepareOnce sync.Once + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/syncBeforeJobStarted", + func(fpJobID int64) { + if fpJobID != jobID64 { + return + } + beforePrepareOnce.Do(func() { + close(reachedBeforePrepare) + }) + <-releaseBeforePrepare + }, + ) + + reachedPrepared := make(chan struct{}) + releasePrepared := make(chan struct{}) + var preparedOnce sync.Once + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeRefreshTask", + func(task *proto.Task) { + if task.Key != importinto.TaskKey(jobID64) || task.Step != proto.StepPrepared { + return + } + preparedOnce.Do(func() { + close(reachedPrepared) + }) + <-releasePrepared + }, + ) + + // Let scheduler pick the task. + s.NoError(failpoint.Disable("github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeGetSchedulableTasks")) + + select { + case <-reachedBeforePrepare: + case <-time.After(maxWaitTime): + s.T().Fatal("timeout waiting scheduler-start-before-prepare timing point") + } + + rows = s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + s.Len(rows, 1) + // timing 2: after scheduler starts but before prepare starts. + s.Equal(result, rows) + close(releaseBeforePrepare) + + select { + case <-reachedPrepared: + case <-time.After(maxWaitTime): + s.T().Fatal("timeout waiting after-prepare-before-business timing point") + } + + // timing 3: after prepare finishes but before the first business step starts. + rows = s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + s.Len(rows, 1) + jobInfo := &importer.JobInfo{ + ID: jobID64, + TableSchema: "show_job_prepare_timing", + TableName: "t", + TableID: tableID, + CreatedBy: "root@%", + Parameters: importer.ImportParameters{ + FileLocation: fmt.Sprintf(`gs://show-job-prepare-timing-source/t.csv?endpoint=%s`, gcsEndpoint), + Format: importer.DataFormatCSV, + }, + SourceFileSize: 3, + Status: "running", + Step: importer.JobStepPreparing, + ErrorMessage: "", + } + s.compareJobInfoWithoutTime(jobInfo, rows[0]) + close(releasePrepared) + + s.Require().Eventually(func() bool { + curRows := s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + return curRows[0][fmap["Status"]] == importer.JobStatusFinished + }, maxWaitTime, 500*time.Millisecond) + rows = s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + s.Len(rows, 1) + jobInfo.Status = importer.JobStatusFinished + jobInfo.Step = "" + jobInfo.Summary = &importer.Summary{ImportedRows: 2} + s.compareJobInfoWithoutTime(jobInfo, rows[0]) + s.tk.MustQuery("SELECT * FROM t ORDER BY i").Check(testkit.Rows("1", "2")) +} + func (s *mockGCSSuite) getTaskByJobID(ctx context.Context, jobID int64) *proto.Task { taskManager, err := storage.GetTaskManager() s.NoError(err) From 5a55f1e8df31c03e6a9cb8e78a4f77501a84c5af Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Tue, 26 May 2026 18:36:46 +0800 Subject: [PATCH 10/21] change --- pkg/ddl/backfilling_dist_scheduler.go | 5 +++++ pkg/dxf/example/scheduler.go | 4 ++++ pkg/dxf/framework/scheduler/scheduler.go | 3 --- pkg/dxf/importinto/scheduler.go | 2 +- 4 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pkg/ddl/backfilling_dist_scheduler.go b/pkg/ddl/backfilling_dist_scheduler.go index ff01fbf481d0b..9ba5288663e06 100644 --- a/pkg/ddl/backfilling_dist_scheduler.go +++ b/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, diff --git a/pkg/dxf/example/scheduler.go b/pkg/dxf/example/scheduler.go index 75bbbb6ae3c37..5f2d2b85359c7 100644 --- a/pkg/dxf/example/scheduler.go +++ b/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 { diff --git a/pkg/dxf/framework/scheduler/scheduler.go b/pkg/dxf/framework/scheduler/scheduler.go index 7ffdc7f506665..1786ad7d2a419 100644 --- a/pkg/dxf/framework/scheduler/scheduler.go +++ b/pkg/dxf/framework/scheduler/scheduler.go @@ -400,9 +400,6 @@ func (s *BaseScheduler) onPending() error { s.task.Store(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() } diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 287644447f496..0c63b2a297acf 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -607,7 +607,7 @@ func (*importScheduler) IsRetryableErr(err error) bool { // GetNextStep implements scheduler.Extension interface. func (sch *importScheduler) GetNextStep(task *proto.TaskBase) proto.Step { switch task.Step { - case proto.StepInit: + case proto.StepInit, proto.StepPrepared: if sch.GlobalSort { return proto.ImportStepEncodeAndSort } From b243c6d8db9f2f0f7ac80d671659551a69f4afce Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Tue, 26 May 2026 21:38:25 +0800 Subject: [PATCH 11/21] test: cover non-retryable import precheck error --- pkg/dxf/importinto/BUILD.bazel | 1 + pkg/dxf/importinto/scheduler_test.go | 2 ++ 2 files changed, 3 insertions(+) diff --git a/pkg/dxf/importinto/BUILD.bazel b/pkg/dxf/importinto/BUILD.bazel index 142df55be94f2..0959b2565fbe8 100644 --- a/pkg/dxf/importinto/BUILD.bazel +++ b/pkg/dxf/importinto/BUILD.bazel @@ -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", diff --git a/pkg/dxf/importinto/scheduler_test.go b/pkg/dxf/importinto/scheduler_test.go index bb227eacf82c6..716082f91d7b9 100644 --- a/pkg/dxf/importinto/scheduler_test.go +++ b/pkg/dxf/importinto/scheduler_test.go @@ -26,6 +26,7 @@ import ( "github.com/pingcap/tidb/pkg/dxf/framework/scheduler" "github.com/pingcap/tidb/pkg/executor/importer" drivererr "github.com/pingcap/tidb/pkg/store/driver/error" + "github.com/pingcap/tidb/pkg/util/dbterror/exeerrors" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" ) @@ -159,6 +160,7 @@ func (s *importIntoSuite) TestIsRetryable() { ext := &importScheduler{} require.True(s.T(), ext.IsRetryableErr(drivererr.ErrRegionUnavailable)) require.True(s.T(), ext.IsRetryableErr(errors.Annotatef(errGetCrossKSSessionPool, "test"))) + require.False(s.T(), ext.IsRetryableErr(exeerrors.ErrLoadDataPreCheckFailed.FastGenByArgs("target table is not empty"))) } func TestIsImporting2TiKV(t *testing.T) { From 95c27a5d1573c232d335f86eae6df4bfd492f995 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 13:26:12 +0800 Subject: [PATCH 12/21] importinto: stabilize nextgen prepare timing tests --- pkg/dxf/framework/scheduler/scheduler.go | 1 + pkg/dxf/importinto/job.go | 3 + tests/realtikvtest/importintotest/job_test.go | 179 +++++++++++++++++- 3 files changed, 176 insertions(+), 7 deletions(-) diff --git a/pkg/dxf/framework/scheduler/scheduler.go b/pkg/dxf/framework/scheduler/scheduler.go index 1786ad7d2a419..e9ffe1cb57319 100644 --- a/pkg/dxf/framework/scheduler/scheduler.go +++ b/pkg/dxf/framework/scheduler/scheduler.go @@ -398,6 +398,7 @@ 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. } diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index c33e802562d57..46099afc61d02 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -69,6 +69,9 @@ func SubmitTask(ctx context.Context, plan *importer.Plan, stmt string) (int64, * // 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() } diff --git a/tests/realtikvtest/importintotest/job_test.go b/tests/realtikvtest/importintotest/job_test.go index 6d6bdc030788f..3b350e32b1d2c 100644 --- a/tests/realtikvtest/importintotest/job_test.go +++ b/tests/realtikvtest/importintotest/job_test.go @@ -72,6 +72,35 @@ func (s *mockGCSSuite) compareJobInfoWithoutTime(jobInfo *importer.JobInfo, row s.Equal(jobInfo.CreatedBy, row[fmap["CreatedBy"]]) } +func (s *mockGCSSuite) blockSchedulerBeforeGetSchedulableTasks() func() { + reachedBlockPoint := make(chan struct{}) + releaseBlock := make(chan struct{}) + var reachedOnce sync.Once + var releaseOnce sync.Once + + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeGetSchedulableTasks", + func() { + reachedOnce.Do(func() { + close(reachedBlockPoint) + }) + <-releaseBlock + }, + ) + releaseFn := func() { + releaseOnce.Do(func() { + close(releaseBlock) + }) + } + s.T().Cleanup(releaseFn) + + select { + case <-reachedBlockPoint: + case <-time.After(maxWaitTime): + s.T().Fatal("timeout waiting scheduler beforeGetSchedulableTasks block point") + } + return releaseFn +} + func (s *mockGCSSuite) TestShowJob() { s.tk.MustExec("delete from mysql.tidb_import_jobs") s.prepareAndUseDB("test_show_job") @@ -396,8 +425,8 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { s.NoError(err) tableID := do.MustGetTableID(s.T(), "show_job_prepare_timing", "t") - // Pause scheduler manager so we can assert show import job before scheduler starts. - testfailpoint.Enable(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeGetSchedulableTasks", "pause") + // Block scheduler manager so we can assert show import job before scheduler starts. + releaseScheduler := s.blockSchedulerBeforeGetSchedulableTasks() importSQL := fmt.Sprintf( `IMPORT INTO t FROM 'gs://show-job-prepare-timing-source/t.csv?endpoint=%s' @@ -410,6 +439,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { jobID, err := strconv.Atoi(result[0][0].(string)) s.NoError(err) jobID64 := int64(jobID) + createdBy := fmt.Sprintf("%v", result[0][fmap["CreatedBy"]]) rows := s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() s.Len(rows, 1) @@ -434,9 +464,9 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { reachedPrepared := make(chan struct{}) releasePrepared := make(chan struct{}) var preparedOnce sync.Once - testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeRefreshTask", + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/afterTaskPrepared", func(task *proto.Task) { - if task.Key != importinto.TaskKey(jobID64) || task.Step != proto.StepPrepared { + if task.Key != importinto.TaskKey(jobID64) { return } preparedOnce.Do(func() { @@ -447,7 +477,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { ) // Let scheduler pick the task. - s.NoError(failpoint.Disable("github.com/pingcap/tidb/pkg/dxf/framework/scheduler/beforeGetSchedulableTasks")) + releaseScheduler() select { case <-reachedBeforePrepare: @@ -475,7 +505,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { TableSchema: "show_job_prepare_timing", TableName: "t", TableID: tableID, - CreatedBy: "root@%", + CreatedBy: createdBy, Parameters: importer.ImportParameters{ FileLocation: fmt.Sprintf(`gs://show-job-prepare-timing-source/t.csv?endpoint=%s`, gcsEndpoint), Format: importer.DataFormatCSV, @@ -485,7 +515,10 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { Step: importer.JobStepPreparing, ErrorMessage: "", } - s.compareJobInfoWithoutTime(jobInfo, rows[0]) + s.Equal(strconv.FormatInt(jobID64, 10), rows[0][fmap["JobID"]]) + s.Equal(importer.JobStatusRunning, rows[0][fmap["Status"]]) + s.Equal(importer.JobStepPreparing, rows[0][fmap["Phase"]]) + s.Equal(createdBy, rows[0][fmap["CreatedBy"]]) close(releasePrepared) s.Require().Eventually(func() bool { @@ -501,6 +534,138 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { s.tk.MustQuery("SELECT * FROM t ORDER BY i").Check(testkit.Rows("1", "2")) } +func (s *mockGCSSuite) TestPrecheckFailureOnPrepareIsNotRetried() { + if !kerneltype.IsNextGen() { + s.T().Skip("async prepare is only enabled in nextgen kernel") + } + + s.prepareAndUseDB("prepare_precheck_no_retry") + s.tk.MustExec("CREATE TABLE t (i INT PRIMARY KEY);") + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "prepare-precheck-no-retry-source", + Name: "t.csv", + }, + Content: []byte("1\n2"), + }) + // Ensure the sort bucket exists before scheduler writes prepared metadata. + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "prepare-precheck-no-retry-sort", + Name: "seed", + }, + Content: []byte("seed"), + }) + + // Block scheduler so we can install failpoints after detached submit. + releaseScheduler := s.blockSchedulerBeforeGetSchedulableTasks() + importSQL := fmt.Sprintf( + `IMPORT INTO t FROM 'gs://prepare-precheck-no-retry-source/t.csv?endpoint=%s' + WITH DETACHED, cloud_storage_uri='gs://prepare-precheck-no-retry-sort/path?endpoint=%s'`, + gcsEndpoint, + gcsEndpoint, + ) + result := s.tk.MustQuery(importSQL).Rows() + s.Len(result, 1) + jobID, err := strconv.Atoi(result[0][0].(string)) + s.NoError(err) + jobID64 := int64(jobID) + + tk2 := testkit.NewTestKit(s.T(), s.store) + var prepareAttempts atomic.Int32 + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/syncBeforeJobStarted", + func(fpJobID int64) { + if fpJobID != jobID64 { + return + } + if prepareAttempts.Add(1) == 1 { + // Make table non-empty after submit but before OnPrepare precheck. + tk2.MustExec("INSERT INTO prepare_precheck_no_retry.t VALUES (100)") + } + }, + ) + releaseScheduler() + + s.Require().Eventually(func() bool { + rows := s.tk.MustQuery(fmt.Sprintf("SHOW IMPORT JOB %d", jobID)).Rows() + return rows[0][fmap["Status"]] == "failed" && + rows[0][fmap["Phase"]] == importer.JobStepPreparing + }, maxWaitTime, 500*time.Millisecond) + + rows := s.tk.MustQuery(fmt.Sprintf("SHOW IMPORT JOB %d", jobID)).Rows() + s.Len(rows, 1) + s.Regexp("target table is not empty", rows[0][fmap["ResultMessage"]]) + s.tk.MustQuery("SELECT * FROM t").Check(testkit.Rows("100")) + + ctx := util.WithInternalSourceType(context.Background(), "taskManager") + s.Require().Eventually(func() bool { + task := s.getTaskByJobID(ctx, jobID64) + return task.State == proto.TaskStateReverted && prepareAttempts.Load() == 1 + }, maxWaitTime, 500*time.Millisecond) + s.Equal(int32(1), prepareAttempts.Load()) +} + +func (s *mockGCSSuite) TestDetachedJobWithoutPrepareModeStillSucceeds() { + if !kerneltype.IsNextGen() { + s.T().Skip("async prepare is only enabled in nextgen kernel") + } + + dbName := fmt.Sprintf("detached_job_without_prepare_mode_%d", time.Now().UnixNano()) + s.prepareAndUseDB(dbName) + s.T().Cleanup(func() { + s.tk.MustExec("drop database if exists " + dbName) + }) + s.tk.MustExec("CREATE TABLE t (i INT PRIMARY KEY);") + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "detached-job-without-prepare-mode-source", + Name: "t.csv", + }, + Content: []byte("1\n2"), + }) + // Ensure the sort bucket exists before scheduler writes external metadata. + s.server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "detached-job-without-prepare-mode-sort", + Name: "seed", + }, + Content: []byte("seed"), + }) + + testfailpoint.Enable(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/mockDisableAsyncPrepare", "return") + // Block scheduler so we can verify submitted task still has no prepare mode. + releaseScheduler := s.blockSchedulerBeforeGetSchedulableTasks() + s.False(importinto.ShouldUseAsyncPrepare(&importer.Plan{CloudStorageURI: "gs://mock-bucket/mock-path"})) + importSQL := fmt.Sprintf( + `IMPORT INTO t FROM 'gs://detached-job-without-prepare-mode-source/t.csv?endpoint=%s' + WITH DETACHED, cloud_storage_uri='gs://detached-job-without-prepare-mode-sort/path?endpoint=%s'`, + gcsEndpoint, + gcsEndpoint, + ) + result := s.tk.MustQuery(importSQL).Rows() + s.Len(result, 1) + jobID, err := strconv.Atoi(result[0][0].(string)) + s.NoError(err) + jobID64 := int64(jobID) + taskKey := importinto.TaskKey(jobID64) + s.tk.MustQuery(fmt.Sprintf( + "select json_extract(extra_params, '$.prepare_mode') is null from mysql.tidb_global_task where task_key = '%s'", + taskKey, + )).Check(testkit.Rows("1")) + releaseScheduler() + + var finalStatus string + s.Require().Eventually(func() bool { + rows := s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + finalStatus = rows[0][fmap["Status"]].(string) + return finalStatus == importer.JobStatusFinished + }, 3*maxWaitTime, 500*time.Millisecond) + rows := s.tk.MustQuery(fmt.Sprintf("show import job %d", jobID)).Rows() + s.Len(rows, 1) + s.Equal("", rows[0][fmap["Phase"]]) + s.tk.MustQuery("SELECT * FROM t ORDER BY i").Check(testkit.Rows("1", "2")) +} + func (s *mockGCSSuite) getTaskByJobID(ctx context.Context, jobID int64) *proto.Task { taskManager, err := storage.GetTaskManager() s.NoError(err) From 2b9d68b003b72042c6546f236dafddebc062b8ab Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 16:17:17 +0800 Subject: [PATCH 13/21] change --- pkg/dxf/importinto/job.go | 1 + pkg/dxf/importinto/planner.go | 28 ++++++++----------------- pkg/dxf/importinto/planner_test.go | 30 +-------------------------- pkg/dxf/importinto/proto.go | 12 +++++------ pkg/dxf/importinto/scheduler.go | 5 ++--- pkg/executor/importer/import.go | 3 +-- pkg/executor/importer/table_import.go | 5 +++-- 7 files changed, 22 insertions(+), 62 deletions(-) diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index 46099afc61d02..2afb23c0d8f8d 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -101,6 +101,7 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc } if ShouldUseAsyncPrepare(plan) { logicalPlan.PrepareMode = proto.PrepareModeRequired + // below params will be filled later in async prepare, init to 1 temporarily. planCtx.ThreadCnt = 1 planCtx.MaxNodeCnt = 1 } diff --git a/pkg/dxf/importinto/planner.go b/pkg/dxf/importinto/planner.go index 7a26ca9d29ba6..adc4eff13ceec 100644 --- a/pkg/dxf/importinto/planner.go +++ b/pkg/dxf/importinto/planner.go @@ -77,12 +77,12 @@ func (p *LogicalPlan) GetTaskExtraParams() proto.ExtraParams { // 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, - PreparedChunkMapExternalPath: p.PreparedChunkMapExternalPath, + JobID: p.JobID, + Plan: p.Plan, + Stmt: p.Stmt, + EligibleInstances: p.EligibleInstances, + ChunkMap: p.ChunkMap, + PreparedMetaExternalPath: p.PreparedChunkMapExternalPath, } return json.Marshal(taskMeta) } @@ -98,7 +98,7 @@ func (p *LogicalPlan) FromTaskMeta(bs []byte) error { p.Stmt = taskMeta.Stmt p.EligibleInstances = taskMeta.EligibleInstances p.ChunkMap = taskMeta.ChunkMap - p.PreparedChunkMapExternalPath = taskMeta.PreparedChunkMapExternalPath + p.PreparedChunkMapExternalPath = taskMeta.PreparedMetaExternalPath return nil } @@ -425,23 +425,13 @@ func readPreparedChunkMap( return nil, err } defer store.Close() - preparedChunkMapMeta := PreparedChunkMapMeta{ + preparedChunkMapMeta := PreparedMeta{ BaseExternalMeta: globalsort.BaseExternalMeta{ExternalPath: externalPath}, } if err := preparedChunkMapMeta.ReadJSONFromExternalStorage(ctx, store, &preparedChunkMapMeta); err != nil { return nil, err } - if preparedChunkMapMeta.ChunkMap != nil { - return preparedChunkMapMeta.ChunkMap, nil - } - - // Compatible with legacy format where the external file stored chunkMap - // directly as the top-level JSON object. - var legacyChunkMap map[int32][]importer.Chunk - if err := preparedChunkMapMeta.ReadJSONFromExternalStorage(ctx, store, &legacyChunkMap); err != nil { - return nil, err - } - return legacyChunkMap, nil + return preparedChunkMapMeta.ChunkMap, nil } func skipMergeSort(kvGroup string, stats []simplesst.MultipleFilesStat, concurrency int) bool { diff --git a/pkg/dxf/importinto/planner_test.go b/pkg/dxf/importinto/planner_test.go index 5cb94c8bd13e8..2823fc743206d 100644 --- a/pkg/dxf/importinto/planner_test.go +++ b/pkg/dxf/importinto/planner_test.go @@ -151,7 +151,7 @@ func TestToPhysicalPlan(t *testing.T) { 1: {{Path: "gs://test-load/2.csv"}}, } externalPath := globalsort.PreparedMetaPath(100) - preparedMeta := PreparedChunkMapMeta{ + preparedMeta := PreparedMeta{ BaseExternalMeta: globalsort.BaseExternalMeta{ExternalPath: externalPath}, ChunkMap: preparedChunkMap, } @@ -173,34 +173,6 @@ func TestToPhysicalPlan(t *testing.T) { require.Len(t, importSpec.Chunks, 1) require.Equal(t, "gs://test-load/2.csv", importSpec.Chunks[0].Path) }) - - t.Run("legacy prepared chunk map format is still readable", 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/legacy.csv"}}, - } - externalPath := globalsort.PreparedMetaPath(101) - data, err := json.Marshal(preparedChunkMap) - require.NoError(t, err) - require.NoError(t, store.WriteFile(context.Background(), externalPath, data)) - - specs, err := generateImportSpecs(planner.PlanCtx{Ctx: context.Background()}, &LogicalPlan{ - Plan: importer.Plan{ - CloudStorageURI: cloudStorageURI, - }, - 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/legacy.csv", importSpec.Chunks[0].Path) - }) } func genEncodeStepMetas(t *testing.T, cnt int) [][]byte { diff --git a/pkg/dxf/importinto/proto.go b/pkg/dxf/importinto/proto.go index de6ea7854e45d..837f4259a7a09 100644 --- a/pkg/dxf/importinto/proto.go +++ b/pkg/dxf/importinto/proto.go @@ -53,21 +53,19 @@ 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 - // PreparedChunkMapExternalPath points to external chunk metadata prepared by + // PreparedMetaExternalPath points to external chunk metadata prepared by // framework OnPrepare for nextgen global-sort path. - PreparedChunkMapExternalPath string `json:"prepared_chunk_map_external_path,omitempty"` + PreparedMetaExternalPath string `json:"prepared_meta_external_path,omitempty"` } -// PreparedChunkMapMeta stores chunk-map metadata generated in prepare stage. -// Keep this wrapper generic so future prepare-stage fields can be added without -// changing the on-disk top-level JSON shape again. -type PreparedChunkMapMeta struct { +// 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 *PreparedChunkMapMeta) Marshal() ([]byte, error) { +func (m *PreparedMeta) Marshal() ([]byte, error) { return m.BaseExternalMeta.Marshal(m) } diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 0c63b2a297acf..9f97f84e722fb 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -300,7 +300,7 @@ func (sch *importScheduler) writePreparedChunkMap( return "", err } defer store.Close() - preparedMeta := PreparedChunkMapMeta{ + preparedMeta := PreparedMeta{ BaseExternalMeta: globalsort.BaseExternalMeta{ ExternalPath: globalsort.PreparedMetaPath(taskID), }, @@ -358,8 +358,7 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, } taskMeta.Plan = *controller.Plan - taskMeta.ChunkMap = nil - taskMeta.PreparedChunkMapExternalPath = chunkMapPath + taskMeta.PreparedMetaExternalPath = chunkMapPath if err = sch.updatePreparedJobInfo(ctx, sch.GetLogger(), taskMeta); err != nil { return err } diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index c0f21158898d8..60c10a55ded1e 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -1643,8 +1643,7 @@ func (e *LoadDataController) detectAndUpdateFormat(path string) { e.Format = parseFileType(path) e.logger.Info("detect and update import plan format based on file extension", zap.String("file", path), zap.String("detected format", e.Format)) - // Plan.Parameters is not persisted in dist-task task meta. Keep this - // in-memory update best-effort for non-prepare paths. + // Plan.Parameters doesn't exist if we run with async prepare. if e.Parameters != nil { e.Parameters.Format = e.Format } diff --git a/pkg/executor/importer/table_import.go b/pkg/executor/importer/table_import.go index cf1916c833382..7538677a9584e 100644 --- a/pkg/executor/importer/table_import.go +++ b/pkg/executor/importer/table_import.go @@ -461,9 +461,10 @@ func (e *LoadDataController) SetExecuteNodeCnt(cnt int) { // in dist framework, this should be done in the tidb node which is responsible for splitting job into subtasks // then table-importer handles data belongs to the subtask. func (e *LoadDataController) PopulateChunks(ctx context.Context) (chunksMap map[int32][]Chunk, err error) { - task := log.BeginTask(e.logger, "populate chunks") + task := log.BeginTask(e.logger.With(zap.Int("executeNodesCnt", e.ExecuteNodesCnt), + zap.Int64("totalFileSize", e.TotalFileSize)), "populate chunks") defer func() { - task.End(zap.ErrorLevel, err) + task.End(zap.ErrorLevel, err, zap.Int("subtaskCnt", len(chunksMap))) }() tableMeta := &mydump.MDTableMeta{ From 11bab49369262affbd3455e863769cecee60e95d Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 17:44:40 +0800 Subject: [PATCH 14/21] change --- pkg/dxf/importinto/BUILD.bazel | 2 +- pkg/dxf/importinto/scheduler_testkit_test.go | 184 +++++++++++++----- pkg/executor/importer/job_test.go | 14 +- tests/realtikvtest/importintotest/job_test.go | 3 + 4 files changed, 148 insertions(+), 55 deletions(-) diff --git a/pkg/dxf/importinto/BUILD.bazel b/pkg/dxf/importinto/BUILD.bazel index 0959b2565fbe8..7dfb55092111f 100644 --- a/pkg/dxf/importinto/BUILD.bazel +++ b/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", diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index 92a92ab1d0033..41aa48224491d 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -38,6 +38,7 @@ import ( "github.com/pingcap/tidb/pkg/ingestor/simplesst" "github.com/pingcap/tidb/pkg/meta/model" "github.com/pingcap/tidb/pkg/parser/ast" + plannercore "github.com/pingcap/tidb/pkg/planner/core" "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/testkit/testfailpoint" "github.com/stretchr/testify/require" @@ -195,54 +196,143 @@ func TestSchedulerExtLocalSort(t *testing.T) { require.NoError(t, err) require.Equal(t, "cancelled", gotJobInfo.Status) - t.Run("prepare-enabled job transitions from preparing to first business phase", func(t *testing.T) { - if !importinto.ShouldUseAsyncPrepare(&logicalPlan.Plan) { - t.Skip("prepare mode only applies when async prepare is enabled") - } - - jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, - "root", "", &importer.ImportParameters{}, 123) - require.NoError(t, err) - - logicalPlan.JobID = jobID - bs, err := logicalPlan.ToTaskMeta() - require.NoError(t, err) - task.Meta = bs - task.Step = proto.StepInit - task.State = proto.TaskStatePending - task.ExtraParams.PrepareMode = proto.PrepareModeRequired - task.ID, err = manager.CreateTask( - ctx, - importinto.TaskKey(jobID), - proto.ImportInto, - "", - 1, - "", - 1, - proto.ExtraParams{PrepareMode: proto.PrepareModeRequired}, - bs, - ) - require.NoError(t, err) - - require.NoError(t, ext.OnPrepare(ctx, d, task)) - - gotJobInfo, err := importer.GetJob(ctx, conn, jobID, "root", true) - require.NoError(t, err) - require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) - require.Equal(t, importer.JobStepPreparing, gotJobInfo.Step) - require.False(t, gotJobInfo.StartTime.IsZero()) - startTime := gotJobInfo.StartTime - - task.Step = proto.StepPrepared - subtaskMetas, err := ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, proto.ImportStepImport) - require.NoError(t, err) - require.Len(t, subtaskMetas, 1) - gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) - require.NoError(t, err) - require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) - require.Equal(t, importer.JobStepImporting, gotJobInfo.Step) - require.Equal(t, startTime, gotJobInfo.StartTime) +} + +func TestSchedulerPrepareEnabledJobTransitionsFromPreparingToFirstBusinessPhase(t *testing.T) { + if !kerneltype.IsNextGen() { + t.Skip("prepare mode only applies in nextgen kernel") + } + + host := "127.0.0.1" + port := uint16(4447) + opt := fakestorage.Options{ + Scheme: "http", + Host: host, + Port: port, + PublicHost: host, + } + gcsEndpoint := fmt.Sprintf("http://%s:%d/storage/v1/", host, port) + sortStorageURI := fmt.Sprintf("gs://sort-bucket/import?endpoint=%s&access-key=aaaaaa&secret-access-key=bbbbbb", gcsEndpoint) + server, err := fakestorage.NewServerWithOptions(opt) + defer server.Stop() + require.NoError(t, err) + server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "sort-bucket"}) + server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "test-load"}) + server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "test-load", + Name: "1.csv", + }, + Content: []byte("1\n"), }) + + testfailpoint.Enable(t, "github.com/pingcap/tidb/pkg/domain/MockDisableDistTask", "return(true)") + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t (id int)") + tbl, err := domain.GetDomain(tk.Session()).InfoSchema().TableByName( + context.Background(), + ast.NewCIStr("test"), + ast.NewCIStr("t"), + ) + require.NoError(t, err) + tblInfo := tbl.Meta().Clone() + pool := pools.NewResourcePool(func() (pools.Resource, error) { + return tk.Session(), nil + }, 1, 1, time.Second) + defer pool.Close() + ctx := context.WithValue(context.Background(), "etcd", true) + ctx = util.WithInternalSourceType(ctx, "taskManager") + mgr := storage.NewTaskManager(pool) + storage.SetTaskManager(mgr) + sch := scheduler.NewManager(util.WithInternalSourceType(ctx, "scheduler"), store, mgr, "host:port", proto.NodeResourceForTest) + keyspace := store.GetKeyspace() + scope := handle.GetTargetScope() + require.NoError(t, mgr.InitMeta(ctx, ":4000", scope)) + + conn := tk.Session().GetSQLExecutor() + jobID, err := importer.CreateJob(ctx, conn, "test", "t", tblInfo.ID, + "root", "", &importer.ImportParameters{}, 0) + require.NoError(t, err) + defaultCharset := "utf8mb4" + logicalPlan := &importinto.LogicalPlan{ + JobID: jobID, + Plan: importer.Plan{ + Path: fmt.Sprintf("gs://test-load/*.csv?endpoint=%s&access-key=aaaaaa&secret-access-key=bbbbbb", gcsEndpoint), + Format: importer.DataFormatAuto, + DBName: "test", + TableInfo: func() *model.TableInfo { + c := tblInfo.Clone() + c.Name = ast.NewCIStr("t") + c.State = model.StatePublic + return c + }(), + DisableTiKVImportMode: true, + CloudStorageURI: sortStorageURI, + InImportInto: true, + Charset: &defaultCharset, + FieldNullDef: []string{`\N`}, + LineFieldsInfo: plannercore.LineFieldsInfo{ + FieldsTerminatedBy: ",", + FieldsEnclosedBy: `"`, + FieldsEscapedBy: `\`, + LinesStartingBy: ``, + LinesTerminatedBy: ``, + }, + }, + Stmt: `IMPORT INTO test.t FROM 'gs://test-load/*.csv?endpoint=xxx'`, + } + require.True(t, importinto.ShouldUseAsyncPrepare(&logicalPlan.Plan)) + bs, err := logicalPlan.ToTaskMeta() + require.NoError(t, err) + task := &proto.Task{ + TaskBase: proto.TaskBase{ + Type: proto.ImportInto, + Step: proto.StepInit, + State: proto.TaskStatePending, + ExtraParams: proto.ExtraParams{PrepareMode: proto.PrepareModeRequired}, + }, + Meta: bs, + StateUpdateTime: time.Now(), + } + task.ID, err = mgr.CreateTask( + ctx, + importinto.TaskKey(jobID), + proto.ImportInto, + keyspace, + 1, + scope, + 1, + proto.ExtraParams{PrepareMode: proto.PrepareModeRequired}, + bs, + ) + require.NoError(t, err) + d := sch.MockScheduler(task) + ext := importinto.NewImportSchedulerForTest(true, task, scheduler.NewParamForTest(mgr, store)) + + require.NoError(t, ext.OnPrepare(ctx, d, task)) + gotJobInfo, err := importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) + require.Equal(t, importer.JobStepPreparing, gotJobInfo.Step) + require.Equal(t, importer.DataFormatCSV, gotJobInfo.Parameters.Format) + require.EqualValues(t, 2, gotJobInfo.SourceFileSize) + require.False(t, gotJobInfo.StartTime.IsZero()) + startTime := gotJobInfo.StartTime + + task.Step = proto.StepPrepared + nextStep := ext.GetNextStep(&task.TaskBase) + require.Equal(t, proto.ImportStepEncodeAndSort, nextStep) + subtaskMetas, err := ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, nextStep) + require.NoError(t, err) + require.NotEmpty(t, subtaskMetas) + gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, gotJobInfo.Status) + require.Equal(t, importer.JobStepGlobalSorting, gotJobInfo.Step) + require.Equal(t, startTime, gotJobInfo.StartTime) } func TestSchedulerOnDoneCancelResetsTableMode(t *testing.T) { diff --git a/pkg/executor/importer/job_test.go b/pkg/executor/importer/job_test.go index d4624ffc444fb..fc1757c022598 100644 --- a/pkg/executor/importer/job_test.go +++ b/pkg/executor/importer/job_test.go @@ -181,13 +181,6 @@ func TestJobHappyPath(t *testing.T) { require.False(t, info.StartTime.IsZero()) startTime := info.StartTime - require.NoError(t, importer.Job2Step(ctx, conn, jobID, importer.JobStepImporting)) - info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) - require.NoError(t, err) - require.Equal(t, importer.JobStatusRunning, info.Status) - require.Equal(t, importer.JobStepImporting, info.Step) - require.Equal(t, startTime, info.StartTime) - require.NoError(t, importer.UpdateJobPreparedInfo(ctx, conn, jobID, 456, importer.DataFormatCSV)) info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) require.NoError(t, err) @@ -195,6 +188,13 @@ func TestJobHappyPath(t *testing.T) { require.Equal(t, importer.DataFormatCSV, info.Parameters.Format) require.Equal(t, createdParams.FileLocation, info.Parameters.FileLocation) require.Equal(t, createdParams.Options, info.Parameters.Options) + + require.NoError(t, importer.Job2Step(ctx, conn, jobID, importer.JobStepGlobalSorting)) + info, err = importer.GetJob(ctx, conn, jobID, "root@%", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, info.Status) + require.Equal(t, importer.JobStepGlobalSorting, info.Step) + require.Equal(t, startTime, info.StartTime) }) } diff --git a/tests/realtikvtest/importintotest/job_test.go b/tests/realtikvtest/importintotest/job_test.go index 3b350e32b1d2c..5130db83fbd21 100644 --- a/tests/realtikvtest/importintotest/job_test.go +++ b/tests/realtikvtest/importintotest/job_test.go @@ -445,6 +445,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { s.Len(rows, 1) // timing 1: before scheduler starts. s.Equal(result, rows) + s.Equal("0B", result[0][fmap["SourceFileSize"]]) reachedBeforePrepare := make(chan struct{}) releaseBeforePrepare := make(chan struct{}) @@ -519,6 +520,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { s.Equal(importer.JobStatusRunning, rows[0][fmap["Status"]]) s.Equal(importer.JobStepPreparing, rows[0][fmap["Phase"]]) s.Equal(createdBy, rows[0][fmap["CreatedBy"]]) + s.EqualValues("3B", rows[0][fmap["SourceFileSize"]]) close(releasePrepared) s.Require().Eventually(func() bool { @@ -648,6 +650,7 @@ func (s *mockGCSSuite) TestDetachedJobWithoutPrepareModeStillSucceeds() { s.NoError(err) jobID64 := int64(jobID) taskKey := importinto.TaskKey(jobID64) + // prepare_mode is 0 and omit when marshal to JSON, so "is null" return true. s.tk.MustQuery(fmt.Sprintf( "select json_extract(extra_params, '$.prepare_mode') is null from mysql.tidb_global_task where task_key = '%s'", taskKey, From 1faec4bb9d3ab5db6aee9e7e86017f553b29edd5 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 18:37:34 +0800 Subject: [PATCH 15/21] change --- pkg/executor/importer/importer_testkit_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/executor/importer/importer_testkit_test.go b/pkg/executor/importer/importer_testkit_test.go index 0e338a9472409..2ecde8e68a5aa 100644 --- a/pkg/executor/importer/importer_testkit_test.go +++ b/pkg/executor/importer/importer_testkit_test.go @@ -457,7 +457,7 @@ func TestCalResourceParams(t *testing.T) { t.Skip("we only cal resource related params in nextgen") } _, tm, ctx := testutil.InitTableTest(t) - + testutil.MockNodeResource(t, 8) require.NoError(t, tm.InitMeta(ctx, "tidb1", handle.GetTargetScope())) c := &importer.LoadDataController{TotalRealSize: 200 * units.TiB, Plan: &importer.Plan{TableInfo: &model.TableInfo{}}} importer.WithLogger(zap.NewNop())(c) From 1fa8d383eb073b805ed9f3451bb02c58d9398946 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 18:41:59 +0800 Subject: [PATCH 16/21] change --- pkg/executor/importer/import.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 60c10a55ded1e..67004b8eb6e02 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -1622,6 +1622,9 @@ func (e *LoadDataController) CalResourceParams(ctx context.Context, ksCodec []by e.MaxNodeCnt = cal.CalcMaxNodeCountForImportInto() e.DistSQLScanConcurrency = scheduler.CalcDistSQLConcurrency(e.ThreadCnt, e.MaxNodeCnt, targetNodeCPUCnt) e.logger.Info("auto calculate resource related params", + zap.String("db", e.DBName), + zap.String("table", e.TableInfo.Name.O), + zap.Int64("tableID", e.TableInfo.ID), zap.Int("thread", e.ThreadCnt), zap.Int("maxNode", e.MaxNodeCnt), zap.Int("distsqlScanConcurrency", e.DistSQLScanConcurrency), From a35424a406798869f1539e1bada0cb6b6c8c6fd9 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 18:51:47 +0800 Subject: [PATCH 17/21] change --- pkg/dxf/importinto/scheduler.go | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 9f97f84e722fb..aeca9806ed3ae 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -341,10 +341,10 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, "No file matched, or the file is empty. Please provide a valid file location.", ) } - if err = sch.checkImportTableEmpty(ctx, taskMeta); err != nil { + if err = controller.CalResourceParams(ctx, sch.TaskStore.GetCodec().GetKeyspace()); err != nil { return err } - if err = controller.CalResourceParams(ctx, sch.TaskStore.GetCodec().GetKeyspace()); err != nil { + if err = sch.updatePreparedJobInfo(ctx, sch.GetLogger(), taskMeta.JobID, controller.Plan); err != nil { return err } controller.SetExecuteNodeCnt(controller.MaxNodeCnt) @@ -359,9 +359,6 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, taskMeta.Plan = *controller.Plan taskMeta.PreparedMetaExternalPath = chunkMapPath - if err = sch.updatePreparedJobInfo(ctx, sch.GetLogger(), taskMeta); err != nil { - return err - } metaBytes, err := json.Marshal(taskMeta) if err != nil { return errors.Trace(err) @@ -787,7 +784,7 @@ func (sch *importScheduler) job2Step(ctx context.Context, logger *zap.Logger, ta ) } -func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *zap.Logger, taskMeta *TaskMeta) error { +func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *zap.Logger, jobID int64, plan *importer.Plan) error { taskManager, err := sch.getTaskMgrForAccessingImportJob() if err != nil { return err @@ -800,9 +797,9 @@ func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *z return importer.UpdateJobPreparedInfo( ctx, exec, - taskMeta.JobID, - taskMeta.Plan.TotalFileSize, - taskMeta.Plan.Format, + jobID, + plan.TotalFileSize, + plan.Format, ) }) }, From 03d9278e92cad4afb21b98f55484a0ab181a0206 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Wed, 27 May 2026 22:26:09 +0800 Subject: [PATCH 18/21] change show sour-size display --- pkg/executor/importer/job.go | 15 ++++++++++++ pkg/executor/show.go | 6 ++++- pkg/executor/show_test.go | 23 +++++++++++++++++++ tests/realtikvtest/importintotest/job_test.go | 2 +- 4 files changed, 44 insertions(+), 2 deletions(-) diff --git a/pkg/executor/importer/job.go b/pkg/executor/importer/job.go index 6495cff8e5f95..0e2a5d314a189 100644 --- a/pkg/executor/importer/job.go +++ b/pkg/executor/importer/job.go @@ -142,6 +142,21 @@ func (j *JobInfo) IsSuccess() bool { return j.Status == JobStatusFinished } +// IsSourceFileSizeUnknown returns whether source_file_size has not been +// determined yet for async-prepare jobs. +func (j *JobInfo) IsSourceFileSizeUnknown() bool { + // if the job is not using async prepare, source_file_size is determined + // when creating the job, and it will never be 0 since we check total file + // size > 0 in precheck. + if j.SourceFileSize > 0 { + return false + } + if j.Status == jobStatusPending { + return true + } + return j.Status == JobStatusRunning && j.Step == JobStepPreparing +} + // GetJob returns the job with the given id if the user has privilege. // hasSuperPriv: whether the user has super privilege. // If the user has super privilege, the user can show or operate all jobs, diff --git a/pkg/executor/show.go b/pkg/executor/show.go index 790bb2f953892..fd51d1eeb610b 100644 --- a/pkg/executor/show.go +++ b/pkg/executor/show.go @@ -2503,7 +2503,11 @@ func FillOneImportJobInfo(result *chunk.Chunk, info *importer.JobInfo, runInfo * result.AppendInt64(4, info.TableID) result.AppendString(5, info.Step) result.AppendString(6, info.Status) - result.AppendString(7, units.BytesSize(float64(info.SourceFileSize))) + if info.IsSourceFileSizeUnknown() { + result.AppendString(7, "N/A") + } else { + result.AppendString(7, units.BytesSize(float64(info.SourceFileSize))) + } if runInfo != nil { // running import job result.AppendUint64(8, uint64(runInfo.ImportRows)) diff --git a/pkg/executor/show_test.go b/pkg/executor/show_test.go index 045cecc159c56..812358ef33e6d 100644 --- a/pkg/executor/show_test.go +++ b/pkg/executor/show_test.go @@ -54,6 +54,7 @@ func TestFillOneImportJobInfo(t *testing.T) { fmap := plannercore.ImportIntoFieldMap rowCntIdx := fmap["ImportedRows"] + sourceFileSizeIdx := fmap["SourceFileSize"] startIdx := fmap["StartTime"] endIdx := fmap["EndTime"] @@ -116,6 +117,28 @@ func TestFillOneImportJobInfo(t *testing.T) { executor.FillOneImportJobInfo(c, jobInfo, ri) require.Equal(t, importer.JobStepPreparing, c.GetRow(6).GetString(fmap["Phase"])) require.Equal(t, proto.Step2Str(proto.ImportInto, proto.StepInit), c.GetRow(6).GetString(fmap["CurStep"])) + + // source_file_size may be unknown for pending or running(preparing) jobs in async prepare. + jobInfo.Status = "pending" + jobInfo.Step = "" + jobInfo.SourceFileSize = 0 + executor.FillOneImportJobInfo(c, jobInfo, nil) + require.Equal(t, "N/A", c.GetRow(7).GetString(sourceFileSizeIdx)) + + jobInfo.Status = importer.JobStatusRunning + jobInfo.Step = importer.JobStepPreparing + executor.FillOneImportJobInfo(c, jobInfo, nil) + require.Equal(t, "N/A", c.GetRow(8).GetString(sourceFileSizeIdx)) + + // For other states/steps, keep the existing size formatting. + jobInfo.Step = importer.JobStepImporting + executor.FillOneImportJobInfo(c, jobInfo, nil) + require.Equal(t, "0B", c.GetRow(9).GetString(sourceFileSizeIdx)) + + jobInfo.Step = importer.JobStepPreparing + jobInfo.SourceFileSize = 3 + executor.FillOneImportJobInfo(c, jobInfo, nil) + require.Equal(t, "3B", c.GetRow(10).GetString(sourceFileSizeIdx)) } func TestShow(t *testing.T) { diff --git a/tests/realtikvtest/importintotest/job_test.go b/tests/realtikvtest/importintotest/job_test.go index 5130db83fbd21..76a03b13ea9e0 100644 --- a/tests/realtikvtest/importintotest/job_test.go +++ b/tests/realtikvtest/importintotest/job_test.go @@ -445,7 +445,7 @@ func (s *mockGCSSuite) TestShowImportJobTimingAroundPrepare() { s.Len(rows, 1) // timing 1: before scheduler starts. s.Equal(result, rows) - s.Equal("0B", result[0][fmap["SourceFileSize"]]) + s.Equal("N/A", result[0][fmap["SourceFileSize"]]) reachedBeforePrepare := make(chan struct{}) releaseBeforePrepare := make(chan struct{}) From f9b8a603e6a5bd4068625ff3d398a702ab84967a Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 28 May 2026 11:15:19 +0800 Subject: [PATCH 19/21] tests --- pkg/dxf/importinto/scheduler.go | 15 +++++++++-- pkg/executor/importer/import.go | 17 ++++++++----- tests/realtikvtest/importintotest/job_test.go | 25 +++---------------- .../importintotest3/extra_param_test.go | 8 ++++-- 4 files changed, 33 insertions(+), 32 deletions(-) diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index aeca9806ed3ae..6c20e87725d08 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -289,7 +289,7 @@ func (sch *importScheduler) checkImportTableEmpty(ctx context.Context, taskMeta }) } -func (sch *importScheduler) writePreparedChunkMap( +func (*importScheduler) writePreparedChunkMap( ctx context.Context, taskID int64, cloudStorageURI string, @@ -332,7 +332,7 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, return err } defer controller.Close() - + isAutoDetectingFormat := controller.Format == importer.DataFormatAuto if err = controller.InitDataFiles(ctx); err != nil { return err } @@ -347,6 +347,16 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, if err = sch.updatePreparedJobInfo(ctx, sch.GetLogger(), taskMeta.JobID, controller.Plan); err != nil { return err } + // following the old behavior, but seems fine to remove this check, those + // specified options are not used anyway when the format is auto-detected as + // non-CSV. maybe remove them later, as we prepare in async way, if user + // use detached mode, user might get an invalid options after job submitted + // while from common sense, options should be validated before job submission. + if isAutoDetectingFormat && controller.Format != importer.DataFormatCSV { + if err = controller.CheckNonCSVFormatOptions(); err != nil { + return errors.Trace(err) + } + } controller.SetExecuteNodeCnt(controller.MaxNodeCnt) chunkMap, err := controller.PopulateChunks(ctx) if err != nil { @@ -366,6 +376,7 @@ func (sch *importScheduler) OnPrepare(ctx context.Context, _ storage.TaskHandle, task.Meta = metaBytes task.RequiredSlots = controller.ThreadCnt task.MaxNodeCount = controller.MaxNodeCnt + failpoint.InjectCall("afterPrepare", task) return nil } diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 67004b8eb6e02..1596bc8f702d0 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -321,7 +321,8 @@ type Plan struct { // only initialized for IMPORT INTO, used when creating job. Parameters *ImportParameters `json:"-"` // only initialized for IMPORT INTO, used when format is detected automatically - specifiedOptions map[string]*plannercore.LoadDataOpt + specifiedOptions map[string]*plannercore.LoadDataOpt + SpecifiedOptionNames map[string]struct{} `json:",omitempty"` // the user who executes the statement, in the form of user@host // only initialized for IMPORT INTO User string `json:"-"` @@ -764,6 +765,10 @@ func (p *Plan) initOptions(ctx context.Context, seCtx sessionctx.Context, option specifiedOptions[opt.Name] = opt } p.specifiedOptions = specifiedOptions + p.SpecifiedOptionNames = make(map[string]struct{}, len(specifiedOptions)) + for k := range p.specifiedOptions { + p.SpecifiedOptionNames[k] = struct{}{} + } if kerneltype.IsNextGen() && sem.IsEnabled() { if p.DataSourceType == DataSourceTypeQuery { @@ -1574,7 +1579,7 @@ func (e *LoadDataController) InitDataFiles(ctx context.Context) error { } } if e.InImportInto && isAutoDetectingFormat && e.Format != DataFormatCSV { - if err2 = e.checkNonCSVFormatOptions(); err2 != nil { + if err2 = e.CheckNonCSVFormatOptions(); err2 != nil { return err2 } } @@ -1835,11 +1840,11 @@ func (p *Plan) IsGlobalSort() bool { return !p.IsLocalSort() } -// non CSV format should not specify CSV only options, we check it again if the -// format is detected automatically. -func (p *Plan) checkNonCSVFormatOptions() error { +// CheckNonCSVFormatOptions non CSV format should not specify CSV only options, +// we check it again if the format is detected automatically. +func (p *Plan) CheckNonCSVFormatOptions() error { for k := range csvOnlyOptions { - if _, ok := p.specifiedOptions[k]; ok { + if _, ok := p.SpecifiedOptionNames[k]; ok { return exeerrors.ErrLoadDataUnsupportedOption.FastGenByArgs(k, "non-CSV format") } } diff --git a/tests/realtikvtest/importintotest/job_test.go b/tests/realtikvtest/importintotest/job_test.go index 76a03b13ea9e0..88b4b4254d92e 100644 --- a/tests/realtikvtest/importintotest/job_test.go +++ b/tests/realtikvtest/importintotest/job_test.go @@ -548,7 +548,7 @@ func (s *mockGCSSuite) TestPrecheckFailureOnPrepareIsNotRetried() { BucketName: "prepare-precheck-no-retry-source", Name: "t.csv", }, - Content: []byte("1\n2"), + Content: []byte(""), }) // Ensure the sort bucket exists before scheduler writes prepared metadata. s.server.CreateObject(fakestorage.Object{ @@ -559,8 +559,6 @@ func (s *mockGCSSuite) TestPrecheckFailureOnPrepareIsNotRetried() { Content: []byte("seed"), }) - // Block scheduler so we can install failpoints after detached submit. - releaseScheduler := s.blockSchedulerBeforeGetSchedulableTasks() importSQL := fmt.Sprintf( `IMPORT INTO t FROM 'gs://prepare-precheck-no-retry-source/t.csv?endpoint=%s' WITH DETACHED, cloud_storage_uri='gs://prepare-precheck-no-retry-sort/path?endpoint=%s'`, @@ -573,21 +571,6 @@ func (s *mockGCSSuite) TestPrecheckFailureOnPrepareIsNotRetried() { s.NoError(err) jobID64 := int64(jobID) - tk2 := testkit.NewTestKit(s.T(), s.store) - var prepareAttempts atomic.Int32 - testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/syncBeforeJobStarted", - func(fpJobID int64) { - if fpJobID != jobID64 { - return - } - if prepareAttempts.Add(1) == 1 { - // Make table non-empty after submit but before OnPrepare precheck. - tk2.MustExec("INSERT INTO prepare_precheck_no_retry.t VALUES (100)") - } - }, - ) - releaseScheduler() - s.Require().Eventually(func() bool { rows := s.tk.MustQuery(fmt.Sprintf("SHOW IMPORT JOB %d", jobID)).Rows() return rows[0][fmap["Status"]] == "failed" && @@ -596,15 +579,13 @@ func (s *mockGCSSuite) TestPrecheckFailureOnPrepareIsNotRetried() { rows := s.tk.MustQuery(fmt.Sprintf("SHOW IMPORT JOB %d", jobID)).Rows() s.Len(rows, 1) - s.Regexp("target table is not empty", rows[0][fmap["ResultMessage"]]) - s.tk.MustQuery("SELECT * FROM t").Check(testkit.Rows("100")) + s.Regexp("the file is empty", rows[0][fmap["ResultMessage"]]) ctx := util.WithInternalSourceType(context.Background(), "taskManager") s.Require().Eventually(func() bool { task := s.getTaskByJobID(ctx, jobID64) - return task.State == proto.TaskStateReverted && prepareAttempts.Load() == 1 + return task.State == proto.TaskStateReverted }, maxWaitTime, 500*time.Millisecond) - s.Equal(int32(1), prepareAttempts.Load()) } func (s *mockGCSSuite) TestDetachedJobWithoutPrepareModeStillSucceeds() { diff --git a/tests/realtikvtest/importintotest3/extra_param_test.go b/tests/realtikvtest/importintotest3/extra_param_test.go index 654683cd7bb0d..f229a86269e0a 100644 --- a/tests/realtikvtest/importintotest3/extra_param_test.go +++ b/tests/realtikvtest/importintotest3/extra_param_test.go @@ -26,11 +26,15 @@ import ( func (s *mockGCSSuite) TestExtraParamMaxRuntimeSlots() { testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", - func(requiredSlots *int, params *proto.ExtraParams) { - *requiredSlots = 16 + func(_ *int, params *proto.ExtraParams) { params.MaxRuntimeSlots = 12 }, ) + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", + func(task *proto.Task) { + task.RequiredSlots = 16 + }, + ) var callCnt int testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool/NewWorkerPool", func(numWorkers int) { s.EqualValues(12, numWorkers) From b04b92ccbb8c8ff7279fa6e51d91f7babb746200 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 28 May 2026 12:05:51 +0800 Subject: [PATCH 20/21] fix comments --- pkg/dxf/importinto/scheduler_testkit_test.go | 10 ++++------ pkg/executor/importer/import.go | 4 +--- .../importintotest3/extra_param_test.go | 18 ++++++++++++------ 3 files changed, 17 insertions(+), 15 deletions(-) diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index 41aa48224491d..bf615193dbef6 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -195,7 +195,6 @@ func TestSchedulerExtLocalSort(t *testing.T) { gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) require.NoError(t, err) require.Equal(t, "cancelled", gotJobInfo.Status) - } func TestSchedulerPrepareEnabledJobTransitionsFromPreparingToFirstBusinessPhase(t *testing.T) { @@ -204,18 +203,17 @@ func TestSchedulerPrepareEnabledJobTransitionsFromPreparingToFirstBusinessPhase( } host := "127.0.0.1" - port := uint16(4447) opt := fakestorage.Options{ Scheme: "http", Host: host, - Port: port, + Port: 0, PublicHost: host, } - gcsEndpoint := fmt.Sprintf("http://%s:%d/storage/v1/", host, port) - sortStorageURI := fmt.Sprintf("gs://sort-bucket/import?endpoint=%s&access-key=aaaaaa&secret-access-key=bbbbbb", gcsEndpoint) server, err := fakestorage.NewServerWithOptions(opt) - defer server.Stop() require.NoError(t, err) + defer server.Stop() + gcsEndpoint := fmt.Sprintf("%s/storage/v1/", server.URL()) + sortStorageURI := fmt.Sprintf("gs://sort-bucket/import?endpoint=%s&access-key=aaaaaa&secret-access-key=bbbbbb", gcsEndpoint) server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "sort-bucket"}) server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "test-load"}) server.CreateObject(fakestorage.Object{ diff --git a/pkg/executor/importer/import.go b/pkg/executor/importer/import.go index 1596bc8f702d0..967dbc631101c 100644 --- a/pkg/executor/importer/import.go +++ b/pkg/executor/importer/import.go @@ -321,7 +321,6 @@ type Plan struct { // only initialized for IMPORT INTO, used when creating job. Parameters *ImportParameters `json:"-"` // only initialized for IMPORT INTO, used when format is detected automatically - specifiedOptions map[string]*plannercore.LoadDataOpt SpecifiedOptionNames map[string]struct{} `json:",omitempty"` // the user who executes the statement, in the form of user@host // only initialized for IMPORT INTO @@ -764,9 +763,8 @@ func (p *Plan) initOptions(ctx context.Context, seCtx sessionctx.Context, option } specifiedOptions[opt.Name] = opt } - p.specifiedOptions = specifiedOptions p.SpecifiedOptionNames = make(map[string]struct{}, len(specifiedOptions)) - for k := range p.specifiedOptions { + for k := range specifiedOptions { p.SpecifiedOptionNames[k] = struct{}{} } diff --git a/tests/realtikvtest/importintotest3/extra_param_test.go b/tests/realtikvtest/importintotest3/extra_param_test.go index f229a86269e0a..3f1f7df95c848 100644 --- a/tests/realtikvtest/importintotest3/extra_param_test.go +++ b/tests/realtikvtest/importintotest3/extra_param_test.go @@ -18,6 +18,7 @@ import ( "fmt" "github.com/fsouza/fake-gcs-server/fakestorage" + "github.com/pingcap/tidb/pkg/config/kerneltype" "github.com/pingcap/tidb/pkg/dxf/framework/proto" "github.com/pingcap/tidb/pkg/lightning/mydump" "github.com/pingcap/tidb/pkg/testkit" @@ -26,15 +27,20 @@ import ( func (s *mockGCSSuite) TestExtraParamMaxRuntimeSlots() { testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/framework/storage/beforeSubmitTask", - func(_ *int, params *proto.ExtraParams) { + func(requiredSlots *int, params *proto.ExtraParams) { + if kerneltype.IsClassic() { + *requiredSlots = 16 + } params.MaxRuntimeSlots = 12 }, ) - testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", - func(task *proto.Task) { - task.RequiredSlots = 16 - }, - ) + if kerneltype.IsNextGen() { + testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/dxf/importinto/afterPrepare", + func(task *proto.Task) { + task.RequiredSlots = 16 + }, + ) + } var callCnt int testfailpoint.EnableCall(s.T(), "github.com/pingcap/tidb/pkg/resourcemanager/pool/workerpool/NewWorkerPool", func(numWorkers int) { s.EqualValues(12, numWorkers) From eaf4e404354dda6fd97eea5ff8e589a493d3ffaa Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 28 May 2026 15:49:09 +0800 Subject: [PATCH 21/21] refine log --- pkg/dxf/importinto/job.go | 30 +++++++++++++++++++----------- 1 file changed, 19 insertions(+), 11 deletions(-) diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index 2afb23c0d8f8d..78f05cc36d477 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -93,18 +93,19 @@ 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, ThreadCnt: plan.ThreadCnt, MaxNodeCnt: plan.MaxNodeCnt, } - if ShouldUseAsyncPrepare(plan) { - logicalPlan.PrepareMode = proto.PrepareModeRequired - // below params will be filled later in async prepare, init to 1 temporarily. - planCtx.ThreadCnt = 1 - planCtx.MaxNodeCnt = 1 - } var ( jobID, taskID int64 ) @@ -168,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 }