From 7e6ebe370ff011d33601476228645dcb03a35a6b Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Tue, 7 Jul 2026 00:58:37 +0800 Subject: [PATCH 1/2] importinto: make dangling import jobs cancellable importinto: simplify dangling job cancellation importinto: clarify dangling cancel fallback executor: document import cancel race change docs(importinto): document job cancellation races executor: cancel dangling import jobs change change importinto: test dangling job cancellation importinto: cover dangling cancel fallback --- pkg/dxf/importinto/BUILD.bazel | 1 + pkg/dxf/importinto/job.go | 1 + pkg/dxf/importinto/job_doc.go | 282 ++++++++++++++++++ pkg/dxf/importinto/job_testkit_test.go | 25 ++ pkg/dxf/importinto/scheduler.go | 49 +++ pkg/dxf/importinto/scheduler_testkit_test.go | 245 ++++++++++----- pkg/executor/BUILD.bazel | 1 + pkg/executor/import_into.go | 44 ++- pkg/executor/import_into_test.go | 57 ++++ pkg/executor/importer/BUILD.bazel | 1 + pkg/executor/importer/job.go | 34 ++- pkg/executor/importer/job_test.go | 73 ++++- tests/realtikvtest/configs/next-gen/pd.toml | 2 +- .../realtikvtest/importintotest3/BUILD.bazel | 2 + .../importintotest3/cross_ks_test.go | 227 ++++++++++++++ 15 files changed, 953 insertions(+), 91 deletions(-) create mode 100644 pkg/dxf/importinto/job_doc.go diff --git a/pkg/dxf/importinto/BUILD.bazel b/pkg/dxf/importinto/BUILD.bazel index 39f544ea50307..8789516ec917c 100644 --- a/pkg/dxf/importinto/BUILD.bazel +++ b/pkg/dxf/importinto/BUILD.bazel @@ -8,6 +8,7 @@ go_library( "conflict_resolution.go", "encode_and_sort_operator.go", "job.go", + "job_doc.go", "metrics.go", "planner.go", "proto.go", diff --git a/pkg/dxf/importinto/job.go b/pkg/dxf/importinto/job.go index 78f05cc36d477..bf8aa269b28b8 100644 --- a/pkg/dxf/importinto/job.go +++ b/pkg/dxf/importinto/job.go @@ -145,6 +145,7 @@ func doSubmitTask(ctx context.Context, plan *importer.Plan, stmt string, instanc // TODO: we need to cleanup the job, if we failed to submit the task to DXF service. dxfTaskMgr := taskManager if runningOnUserKS { + failpoint.InjectCall("afterUserImportJobCreatedBeforeDXFTask", jobID) var err2 error dxfTaskMgr, err2 = storage.GetDXFSvcTaskMgr() if err2 != nil { diff --git a/pkg/dxf/importinto/job_doc.go b/pkg/dxf/importinto/job_doc.go new file mode 100644 index 0000000000000..cbebc8d80957d --- /dev/null +++ b/pkg/dxf/importinto/job_doc.go @@ -0,0 +1,282 @@ +// Copyright 2026 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package importinto + +// IMPORT INTO job and DXF task transaction notes. +// +// This file documents the important transaction boundaries between the +// user-facing import job row and the DXF task row. It is intentionally kept +// close to the importinto code because the behavior depends on several +// helpers in this package plus the DXF framework. +// +// Two rows describe one IMPORT INTO job: +// +// - mysql.tidb_import_jobs stores the user-visible import job status, +// step, summary, and error. In nextgen user keyspace mode this row lives +// in the user keyspace. +// - mysql.tidb_global_task stores the DXF task state, step, task key, and +// keyspace. In nextgen user keyspace mode this row lives in the SYSTEM +// keyspace, because DXF runs as a SYSTEM service. +// +// The normal nextgen submit path is intentionally not atomic across these two +// rows: +// +// 1. Create the import job in the current/user keyspace with +// storage.GetTaskManager().WithNewTxn. +// 2. Commit the import job row as status=pending, step=''. +// 3. Open the DXF service task manager with storage.GetDXFSvcTaskMgr(). +// 4. Create the SYSTEM keyspace DXF task in a second WithNewTxn. +// 5. Commit the DXF task row as state=pending, step=StepInit. +// +// Classic kernel mode and SYSTEM keyspace mode do not have this split: job +// creation and task creation are in the same transaction there. The split only +// exists for nextgen user keyspace imports. +// +// This split is a consequence of how nextgen DXF service and IMPORT INTO +// evolved from the Classic kernel implementation. Classic kernel DXF and +// IMPORT INTO ran in one keyspace, so the import job and DXF task could be +// submitted atomically in one transaction. Nextgen moved DXF service work into +// SYSTEM while keeping the user-facing IMPORT INTO job in the user keyspace. +// That creates a split-brain-style consistency window between the job row and +// the task row. The current design is not the cleanest possible solution, but +// it is the practical one for now because it preserves the existing IMPORT +// INTO job model while allowing DXF to run as a SYSTEM service. +// +// A cleaner long-term design would make the ownership boundary explicit: DXF +// service in SYSTEM owns only DXF tasks, and TiDB/IMPORT INTO owns import jobs +// inside the user keyspace. The two sides would communicate through a clearer +// contract instead of relying on partially coordinated table updates across +// keyspaces. That would avoid this class of split-brain issue more directly, +// but it is a larger design change and is not part of the current +// implementation. +// +// Transaction and statement nodes used below: +// +// J +// The user-keyspace import job transaction commits. The job is externally +// visible as pending. +// +// T +// The SYSTEM-keyspace DXF task transaction commits. The task is externally +// visible as pending/StepInit. +// +// C0 +// CANCEL IMPORT JOB runs checkPrivilegeAndStatus. This reads the import +// job row and allows cancellation only when the job status is pending or +// running and the caller has privilege. +// +// C1 +// CANCEL IMPORT JOB runs CancelTaskByKeySession in a SYSTEM-keyspace +// transaction. It updates a live DXF task to cancelling only when the +// current DXF task state is pending, running, or awaiting-resolution. +// The affected row count is not checked, so "no task", "task not +// committed yet", "already terminal", and "state not matched" all look +// like success to this step. +// +// C2 +// CANCEL IMPORT JOB runs WaitTaskDoneByKey. The first lookup checks live +// and history task tables by task key. If no task is found, it returns +// ErrTaskNotFound. If a task is found, it polls by task ID until the DXF +// task is succeed, reverted, or failed. It does not verify that the task +// was actually cancelled. +// +// C3 +// If C2 returns ErrTaskNotFound, cancelDanglingImportJob runs +// CancelPendingJob in the user keyspace. It directly changes the import +// job from pending to cancelled. Unlike most scheduler writes, this path +// checks affected rows and returns "job state changed during cancel, +// please try again later" if the job is no longer pending. +// +// S0 +// The import scheduler runs checkImportJobNotCancelled. It reads only the +// import job row. A normal visible-task cancel is not visible to this +// check until the scheduler later writes the import job to cancelled. +// +// S1 +// The import scheduler runs StartJob. It updates pending -> running and +// sets the first import job step. The SQL predicate requires status +// pending, but the affected row count is ignored. +// +// S2 +// The import scheduler runs Job2Step. It updates the import job step when +// status is running. The affected row count is ignored. +// +// SF, SE, SC +// The import scheduler runs FinishJob, FailJob, or CancelJob. These are +// terminal import-job writes guarded by the current status. Their affected +// row counts are ignored, so the first terminal writer wins and later +// terminal writers silently no-op. +// +// Submit and cancel interleavings. +// +// The following cases partition the externally different time sequences for +// nextgen user keyspace submit and user cancellation: +// +// 1. C0 < J +// The job row is not visible. CANCEL IMPORT JOB returns job-not-found. +// This normally requires an out-of-band or guessed job ID. +// +// 2. J commits, but T never commits +// The import job remains pending without a DXF task. This can happen if +// the SYSTEM DXF task submission fails after the user job transaction has +// already committed. A later dangling cancel can still cancel the pending +// import job. +// +// 3. J < C0 < C1(0 rows) < C2(no task) < C3, and T never commits +// CancelTaskByKeySession misses the absent DXF task. WaitTaskDoneByKey +// also finds no task. cancelDanglingImportJob changes the import job from +// pending to cancelled. No DXF task exists. +// +// 4. J < C0 < C1(0 rows) < C2(no task) < C3 < T +// The user sees cancel success because the import job row was directly +// marked cancelled. The submitter can still later commit the SYSTEM DXF +// task because the submit path does not recheck the job status before T. +// The scheduler should then see the cancelled import job at S0 and drive +// the DXF task to revert before import work starts. +// +// 5. J < C0 < C1(0 rows) < T < C2(found task) +// This is the known missed-cancel window. C1 ran before the task was +// committed, so no durable cancel marker was recorded in the DXF task +// table. C2 then finds the task, so the dangling-job fallback does not +// run. The cancel command waits for the task's natural terminal state; +// without another cancel, the task can finish or fail normally, and the +// command can still return nil because WaitTaskDoneByKey accepts any done +// state. +// +// 6. J < C0 < C1(0 rows) < C2(no task) < T < S1 < C3 +// The dangling fallback races with a newly committed task and scheduler +// start. If S1 changes the import job to running first, C3 affects zero +// rows and returns "job state changed during cancel, please try again +// later". The task continues unless another cancel succeeds. +// +// 7. J < T < C1, and the DXF task is pending, running, or +// awaiting-resolution +// This is normal visible-task cancellation. C1 changes the DXF task to +// cancelling. The DXF scheduler later changes cancelling -> reverting -> +// reverted, and the import scheduler maps the cancellation error to the +// import job status cancelled. +// +// 8. J < T < C1, but the DXF task is already cancelling, reverting, +// succeed, failed, reverted, or otherwise outside C1's predicate +// C1 affects zero rows and does not report that fact. If C2 finds a live +// or history task, the command waits for or observes terminal completion. +// This can look like cancel success even though this cancel did not +// change the task. +// +// 9. C0 reads a terminal import job, or the caller lacks privilege +// Cancellation stops before touching the DXF task. +// +// Scheduler and DXF framework interleavings after T is visible. +// +// - Normal visible-task cancellation changes only the DXF task first. The +// import job row stays pending or running until the scheduler reaches the +// cancellation OnDone path and runs SC. +// +// - S0 checks only the import job row. Therefore C1 can happen before S0 and +// S0 can still pass, because C1 changed the DXF task state but did not +// change the import job status. +// +// - Between S0 and S1 or S2, a direct import-job cancellation makes S1/S2 +// affect zero rows. A normal visible-task cancellation does not; S1/S2 can +// still advance the import job status or step before the DXF cancellation +// is processed. +// +// - In async prepare, the order is: +// +// S0 -> S1(preparing) -> prepare work -> prepared-info read/update +// -> SwitchTaskStepAfterPrepare +// +// If C1 changes the DXF task to cancelling before the final task-step +// compare-and-set, the task-step update no-ops and the next scheduler +// refresh sees cancelling. The import job can still temporarily show +// running/preparing or prepared metadata. +// +// - During business-step planning, the order is: +// +// S0 -> S1 or S2 -> generate subtask metadata -> SwitchTaskStep +// +// If C1 wins before SwitchTaskStep, the task-step update no-ops and no +// subtasks are inserted in the normal non-batch path. If SwitchTaskStep +// wins first, C1 can still match running and cancellation proceeds. +// +// - The batch step-switch path inserts subtask batches before the task +// state/step update. A concurrent cancel can leave pending subtasks while +// the task is already cancelling. Executors should not run those subtasks +// as normal work while the task is cancelling; once the task becomes +// reverting, executor managers cancel and drain pending/running subtasks. +// +// - If a running subtask fails and manual recovery is enabled, the framework +// tries to move the DXF task to awaiting-resolution. C1 is allowed from +// awaiting-resolution. If C1 wins first, the awaiting-resolution update +// affects zero rows; if awaiting-resolution wins first, C1 changes it to +// cancelling. +// +// - On successful completion, the import scheduler's OnDone hook runs before +// the DXF framework persists succeed. If FinishJob commits before a user +// cancel, the import job can become finished. If the cancel then changes +// the DXF task to cancelling before SucceedTask commits, SucceedTask can +// affect zero rows and the DXF task can later become reverted while the +// import job remains finished. The later CancelJob sees a terminal import +// job and silently no-ops. +// +// - On failure or cancellation completion, the first terminal import-job +// writer wins. FailJob and CancelJob both match pending/running jobs and +// ignore affected rows. A losing terminal writer silently no-ops. +// +// Other import-job-related actors. +// +// - Active-job precheck is a separate read before CreateJob. Two concurrent +// imports can both observe zero active jobs before either creates a job. +// +// - Non-detached IMPORT INTO connection cancellation calls the same +// cancelAndWaitImportJob path and inherits the same races. +// +// - The IMPORT SDK and Lightning do not directly update tidb_import_jobs. +// They enter through SQL front ends such as IMPORT INTO, CANCEL IMPORT JOB, +// SHOW IMPORT JOB, and SHOW IMPORT GROUP. Lightning group cancellation has +// an extra stale-snapshot window: it reads group jobs and then cancels each +// non-completed job one by one. +// +// - SHOW IMPORT JOBS and SHOW IMPORT GROUPS are observers, not writers. They +// read import job rows first and then read DXF runtime/subtask information +// separately for running jobs. Their result can mix snapshots from +// different moments. +// +// - awaiting-resolution is not stored as an import job status. The import +// job row remains running; SHOW IMPORT JOB overlays awaiting-resolution +// from the DXF task state when runtime information is available. +// +// - IMPORT INTO ... FROM SELECT returns through the import-from-select path +// before submitTask and does not use this import-job/DXF-task table path. +// +// Practical risks that need to be fixed later. +// +// 1. Highest: J < C0 < C1(0 rows) < T < C2(found task). The cancel request +// is missed, the dangling fallback is skipped, and the command waits for +// natural task completion. +// +// 2. High: success hook race. FinishJob can commit before a later cancel +// prevents SucceedTask from committing, leaving an import job marked +// finished while the DXF task later reverts. +// +// 3. Medium: C2(no task) < T < S1 < C3. The dangling fallback loses because +// the import job is no longer pending and reports retry-later. +// +// 4. Medium: step/status drift. Normal visible-task cancellation does not +// mark the import job cancelled until scheduler OnDone, so job status and +// step can advance after the user requested cancellation. +// +// 5. Observer-only: SHOW IMPORT, SDK, and Lightning can observe mixed +// snapshots, but they do not introduce new direct import-job writers. diff --git a/pkg/dxf/importinto/job_testkit_test.go b/pkg/dxf/importinto/job_testkit_test.go index 6fa2221092c9f..6979809ac208e 100644 --- a/pkg/dxf/importinto/job_testkit_test.go +++ b/pkg/dxf/importinto/job_testkit_test.go @@ -167,6 +167,31 @@ 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("cancel dangling user keyspace job without DXF task", func(t *testing.T) { + sysKSTK.MustExec("delete from mysql.tidb_import_jobs") + sysKSTK.MustExec("delete from mysql.tidb_global_task") + userKSTK.MustExec("delete from mysql.tidb_import_jobs") + userKSTK.MustExec("delete from mysql.tidb_global_task") + config.UpdateGlobal(func(conf *config.Config) { + conf.KeyspaceName = "ks" + }) + manuallyInitFn(t, userKSStore, sysKSStore) + + conn := userKSTK.Session().GetSQLExecutor() + jobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, + userKSTK.Session().GetSessionVars().User.String(), "", &importer.ImportParameters{}, 123) + require.NoError(t, err) + sysKSTK.MustQuery("select count(1) from mysql.tidb_global_task where task_key = ?", importinto.TaskKey(jobID)). + Check(testkit.Rows("0")) + + userKSTK.MustExec(fmt.Sprintf("cancel import job %d", jobID)) + + userKSTK.MustQuery("select status, error_message from mysql.tidb_import_jobs where id = ?", jobID). + Check(testkit.Rows("cancelled cancelled by user")) + sysKSTK.MustQuery("select count(1) from mysql.tidb_global_task where task_key = ?", importinto.TaskKey(jobID)). + Check(testkit.Rows("0")) + }) + t.Run("submit global-sort task uses async prepare mode", func(t *testing.T) { sysKSTK.MustExec("delete from mysql.tidb_import_jobs") sysKSTK.MustExec("delete from mysql.tidb_global_task") diff --git a/pkg/dxf/importinto/scheduler.go b/pkg/dxf/importinto/scheduler.go index 144e5a9e6cd48..c3820dbc53047 100644 --- a/pkg/dxf/importinto/scheduler.go +++ b/pkg/dxf/importinto/scheduler.go @@ -318,6 +318,9 @@ 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 err := sch.checkImportJobNotCancelled(ctx, sch.GetLogger(), taskMeta); err != nil { + return err + } if err := sch.startJob(ctx, sch.GetLogger(), taskMeta, importer.JobStepPreparing); err != nil { return err } @@ -408,6 +411,9 @@ func (sch *importScheduler) OnNextSubtasksBatch( zap.Int64("table-id", taskMeta.Plan.TableInfo.ID), ) logger.Info("on next subtasks batch") + if err = sch.checkImportJobNotCancelled(ctx, logger, taskMeta); err != nil { + return nil, err + } // Check table emptiness again after the task is started. if kerneltype.IsClassic() && task.Step == proto.StepInit { @@ -796,6 +802,49 @@ func (sch *importScheduler) job2Step(ctx context.Context, logger *zap.Logger, ta ) } +func (sch *importScheduler) checkImportJobNotCancelled( + ctx context.Context, + logger *zap.Logger, + taskMeta *TaskMeta, +) error { + if kerneltype.IsClassic() { + // Classic creates the import job and DXF task in the same transaction, so + // there is no dangling import job to catch here. CANCEL IMPORT JOB moves + // the DXF task to cancelling, and the framework cancellation path handles it. + return nil + } + 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) { + retryable := true + err := taskManager.WithNewSession(func(se sessionctx.Context) error { + job, err := importer.GetJob(ctx, se.GetSQLExecutor(), taskMeta.JobID, "", true) + if err != nil { + // shouldn't happen in normal path unless other clients delete + // it from the system table, just sanity check. + if goerrors.Is(err, exeerrors.ErrLoadDataJobNotFound) { + retryable = false + } + return err + } + // in next-gen, the import job might be taken as a dangling job + // and canceled directly as DXF task might be submitted in another + // transaction. + if job.IsCancelled() { + retryable = false + return errors.Errorf("import job %d cancelled by user", job.ID) + } + return nil + }) + return retryable, err + }, + ) +} + func (sch *importScheduler) updatePreparedJobInfo(ctx context.Context, logger *zap.Logger, jobID int64, plan *importer.Plan) error { taskManager, err := sch.getTaskMgrForAccessingImportJob() if err != nil { diff --git a/pkg/dxf/importinto/scheduler_testkit_test.go b/pkg/dxf/importinto/scheduler_testkit_test.go index 6083f0b42462d..d32700b864db9 100644 --- a/pkg/dxf/importinto/scheduler_testkit_test.go +++ b/pkg/dxf/importinto/scheduler_testkit_test.go @@ -222,6 +222,45 @@ func TestSchedulerExtLocalSort(t *testing.T) { gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) require.NoError(t, err) require.Equal(t, "cancelled", gotJobInfo.Status) + + jobID, err = importer.CreateJob(ctx, conn, "test", "t", 1, + "root", "", &importer.ImportParameters{}, 123) + require.NoError(t, err) + require.NoError(t, importer.CancelJob(ctx, conn, jobID)) + logicalPlan.JobID = jobID + bs, err = logicalPlan.ToTaskMeta() + require.NoError(t, err) + task.Meta = bs + task.Step = proto.StepInit + task.State = proto.TaskStatePending + if kerneltype.IsNextGen() { + // If a nextgen dangling import job was already cancelled before scheduler + // admission, the scheduler should enter the existing cancel/revert path + // without planning work or generating subtasks. + err = ext.OnPrepare(ctx, d, task) + require.Error(t, err) + require.True(t, scheduler.IsCancelledErr(err), err) + gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, "cancelled", gotJobInfo.Status) + require.Equal(t, "", gotJobInfo.Step) + + subtaskMetas, err = ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, ext.GetNextStep(&task.TaskBase)) + require.Error(t, err) + require.True(t, scheduler.IsCancelledErr(err), err) + require.Nil(t, subtaskMetas) + } else { + // Classic has no dangling import job window: CANCEL IMPORT JOB changes the + // DXF task to cancelling, so this import-job status check must not reject + // scheduler planning if a test calls the hook directly. + subtaskMetas, err = ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, ext.GetNextStep(&task.TaskBase)) + require.NoError(t, err) + require.NotNil(t, subtaskMetas) + } + gotJobInfo, err = importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, "cancelled", gotJobInfo.Status) + require.Equal(t, "", gotJobInfo.Step) } func TestSchedulerPrepareEnabledJobTransitionsFromPreparingToFirstBusinessPhase(t *testing.T) { @@ -281,87 +320,139 @@ func TestSchedulerPrepareEnabledJobTransitionsFromPreparingToFirstBusinessPhase( 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: ``, + var taskMgr scheduler.TaskManager = mgr + createPrepareTask := func(t *testing.T) (int64, *proto.Task) { + t.Helper() + 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(), + 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) + return jobID, task } - 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) - var taskMgr scheduler.TaskManager = mgr - ext := importinto.NewImportSchedulerForTest(true, task, scheduler.NewParamForTest(taskMgr, newImportTestRuntime(ctrl, store, pool))) - 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 + t.Run("transitions_to_encode_and_sort", func(t *testing.T) { + jobID, task := createPrepareTask(t) + d := sch.MockScheduler(task) + ext := importinto.NewImportSchedulerForTest(true, task, + scheduler.NewParamForTest(taskMgr, newImportTestRuntime(ctrl, store, pool))) + + require.NoError(t, ext.OnPrepare(ctx, d, task)) + 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.Equal(t, importer.DataFormatCSV, info.Parameters.Format) + require.EqualValues(t, 2, info.SourceFileSize) + require.False(t, info.StartTime.IsZero()) + startTime := info.StartTime + + task.Step = proto.StepPrepared + nextStep := ext.GetNextStep(&task.TaskBase) + require.Equal(t, proto.ImportStepEncodeAndSort, nextStep) + metas, err := ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, nextStep) + require.NoError(t, err) + require.NotEmpty(t, metas) + 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) + }) - 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) + t.Run("cancelled_before_prepare", func(t *testing.T) { + jobID, task := createPrepareTask(t) + require.NoError(t, importer.CancelJob(ctx, conn, jobID)) + d := sch.MockScheduler(task) + ext := importinto.NewImportSchedulerForTest(true, task, + scheduler.NewParamForTest(taskMgr, newImportTestRuntime(ctrl, store, pool))) + + err := ext.OnPrepare(ctx, d, task) + require.Error(t, err) + require.True(t, scheduler.IsCancelledErr(err), err) + info, err := importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, "cancelled", info.Status) + require.Equal(t, "", info.Step) + require.Equal(t, 0, task.RequiredSlots) + require.Equal(t, 0, task.MaxNodeCount) + }) + + t.Run("cancelled_after_prepare", func(t *testing.T) { + jobID, task := createPrepareTask(t) + d := sch.MockScheduler(task) + ext := importinto.NewImportSchedulerForTest(true, task, + scheduler.NewParamForTest(taskMgr, newImportTestRuntime(ctrl, store, pool))) + require.NoError(t, ext.OnPrepare(ctx, d, task)) + 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.NoError(t, importer.CancelJob(ctx, conn, jobID)) + + task.Step = proto.StepPrepared + nextStep := ext.GetNextStep(&task.TaskBase) + require.Equal(t, proto.ImportStepEncodeAndSort, nextStep) + metas, err := ext.OnNextSubtasksBatch(ctx, d, task, []string{":4000"}, nextStep) + require.Error(t, err) + require.True(t, scheduler.IsCancelledErr(err), err) + require.Nil(t, metas) + info, err = importer.GetJob(ctx, conn, jobID, "root", true) + require.NoError(t, err) + require.Equal(t, "cancelled", info.Status) + require.Equal(t, importer.JobStepPreparing, info.Step) + }) } func TestSchedulerOnDoneCancelResetsTableMode(t *testing.T) { diff --git a/pkg/executor/BUILD.bazel b/pkg/executor/BUILD.bazel index 6811ea88d3f88..0c50820197f26 100644 --- a/pkg/executor/BUILD.bazel +++ b/pkg/executor/BUILD.bazel @@ -434,6 +434,7 @@ go_test( "//pkg/domain/affinity", "//pkg/domain/infosync", "//pkg/dxf/framework/proto", + "//pkg/dxf/framework/storage", "//pkg/dxf/importinto", "//pkg/errctx", "//pkg/errno", diff --git a/pkg/executor/import_into.go b/pkg/executor/import_into.go index 2e1753774ef68..deb932928ac28 100644 --- a/pkg/executor/import_into.go +++ b/pkg/executor/import_into.go @@ -16,6 +16,7 @@ package executor import ( "context" + goerrors "errors" "fmt" "github.com/google/uuid" @@ -406,11 +407,50 @@ func cancelAndWaitImportJob(ctx context.Context, jobID int64) error { if err != nil { return err } + taskKey := importinto.TaskKey(jobID) if err := manager.WithNewTxn(ctx, func(se sessionctx.Context) error { ctx = util.WithInternalSourceType(ctx, kv.InternalDistTask) - return manager.CancelTaskByKeySession(ctx, se, importinto.TaskKey(jobID)) + return manager.CancelTaskByKeySession(ctx, se, taskKey) }); err != nil { return err } - return handle.WaitTaskDoneByKey(ctx, importinto.TaskKey(jobID)) + // TODO: Fix the next-gen race where CancelTaskByKeySession can miss a DXF + // task that has not been committed yet. If the task appears before + // WaitTaskDoneByKey's initial lookup, this session waits for the task to + // finish because the cancel request was not recorded; a later CANCEL IMPORT + // JOB can still cancel the now-visible task and unblock the wait. + // see job_doc.go for more detail + if err = handle.WaitTaskDoneByKey(ctx, taskKey); goerrors.Is(err, dxfstorage.ErrTaskNotFound) { + // In next-gen, the import job and DXF task are created in separate + // transactions. The job row can exist before the DXF task row is + // committed, or the task submission can fail after the job is created. If + // no task row is visible here, cancel the job directly. If the task row is + // committed later, the import scheduler checks the job status before + // prepare/planning and will stop before doing import work. + logutil.Logger(ctx).Info("cancel import job directly because dxf task is not found", + zap.Int64("jobID", jobID), + zap.String("taskKey", taskKey)) + failpoint.InjectCall("beforeCancelDanglingImportJob", jobID) + return cancelDanglingImportJob(ctx, jobID) + } + return err +} + +// cancelDanglingImportJob cancels an import job whose DXF task row does not +// exist, so it will not block another import job for the same table. see +// job_doc.go for more detail +func cancelDanglingImportJob(ctx context.Context, jobID int64) error { + manager, err := dxfstorage.GetTaskManager() + if err != nil { + return err + } + return manager.WithNewSession(func(se sessionctx.Context) error { + if err2 := importer.CancelPendingJob(ctx, se.GetSQLExecutor(), jobID); err2 != nil { + return err2 + } + if se.GetSessionVars().StmtCtx.AffectedRows() == 0 { + return errors.New("job state changed during cancel, please try again later") + } + return nil + }) } diff --git a/pkg/executor/import_into_test.go b/pkg/executor/import_into_test.go index 409a9bfa8f7ea..a39ad8a8bff9a 100644 --- a/pkg/executor/import_into_test.go +++ b/pkg/executor/import_into_test.go @@ -23,6 +23,8 @@ import ( "github.com/pingcap/tidb/pkg/config" "github.com/pingcap/tidb/pkg/config/kerneltype" + "github.com/pingcap/tidb/pkg/dxf/framework/storage" + "github.com/pingcap/tidb/pkg/dxf/importinto" "github.com/pingcap/tidb/pkg/executor" "github.com/pingcap/tidb/pkg/executor/importer" "github.com/pingcap/tidb/pkg/kv" @@ -39,6 +41,7 @@ import ( semv1 "github.com/pingcap/tidb/pkg/util/sem" semv2 "github.com/pingcap/tidb/pkg/util/sem/v2" "github.com/stretchr/testify/require" + tikvutil "github.com/tikv/client-go/v2/util" ) var ( @@ -254,6 +257,60 @@ func TestNextGenUnsupportedLocalSortAndOptions(t *testing.T) { } } +func TestCancelImportJobWithoutDXFTask(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("import job and dxf task is submitted together in classic, no such case") + } + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + + ctx := tikvutil.WithInternalSourceType(context.Background(), kv.InternalDistTask) + manager, err := storage.GetTaskManager() + require.NoError(t, err) + require.NoError(t, manager.InitMeta(ctx, ":4000", "")) + + assertNoDXFTask := func(jobID int64) { + taskKey := importinto.TaskKey(jobID) + tk.MustQuery("select count(1) from mysql.tidb_global_task where task_key = ?", taskKey). + Check(testkit.Rows("0")) + tk.MustQuery("select count(1) from mysql.tidb_global_task_history where task_key = ?", taskKey). + Check(testkit.Rows("0")) + } + + jobID, err := importer.CreateJob(ctx, tk.Session().GetSQLExecutor(), "test", "t", 1, + tk.Session().GetSessionVars().User.String(), "", &importer.ImportParameters{ + Format: importer.DataFormatCSV, + }, 0) + require.NoError(t, err) + assertNoDXFTask(jobID) + + tk.MustExec(fmt.Sprintf("cancel import job %d", jobID)) + tk.MustQuery("select status, error_message from mysql.tidb_import_jobs where id = ?", jobID). + Check(testkit.Rows("cancelled cancelled by user")) + assertNoDXFTask(jobID) + + err = tk.ExecToErr(fmt.Sprintf("cancel import job %d", jobID)) + require.ErrorIs(t, err, exeerrors.ErrLoadDataInvalidOperation) + require.ErrorContains(t, err, "The current job status cannot perform the operation. CANCEL") + tk.MustQuery("select status, error_message from mysql.tidb_import_jobs where id = ?", jobID). + Check(testkit.Rows("cancelled cancelled by user")) + assertNoDXFTask(jobID) + + runningJobID, err := importer.CreateJob(ctx, tk.Session().GetSQLExecutor(), "test", "t", 1, + tk.Session().GetSessionVars().User.String(), "", &importer.ImportParameters{ + Format: importer.DataFormatCSV, + }, 0) + require.NoError(t, err) + require.NoError(t, importer.StartJob(ctx, tk.Session().GetSQLExecutor(), runningJobID, importer.JobStepImporting)) + assertNoDXFTask(runningJobID) + + err = tk.ExecToErr(fmt.Sprintf("cancel import job %d", runningJobID)) + require.ErrorContains(t, err, "job state changed during cancel, please try again later") + tk.MustQuery("select status, step from mysql.tidb_import_jobs where id = ?", runningJobID). + Check(testkit.Rows("running importing")) + assertNoDXFTask(runningJobID) +} + func testNextGenUnsupportedLocalSortAndOptions(t *testing.T, store kv.Storage, initFn func(t *testing.T, tk *testkit.TestKit)) { t.Run("import from select", func(t *testing.T) { tk := testkit.NewTestKit(t, store) diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index ee86cb36c5eb6..bf25898913f34 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -121,6 +121,7 @@ go_test( flaky = True, race = "on", shard_count = 42, + shard_count = 43, deps = [ "//br/pkg/mock", "//br/pkg/streamhelper", diff --git a/pkg/executor/importer/job.go b/pkg/executor/importer/job.go index 0e2a5d314a189..bc23e27c9f713 100644 --- a/pkg/executor/importer/job.go +++ b/pkg/executor/importer/job.go @@ -137,6 +137,11 @@ func (j *JobInfo) CanCancel() bool { return j.Status == jobStatusPending || j.Status == JobStatusRunning } +// IsCancelled returns whether the job has been cancelled. +func (j *JobInfo) IsCancelled() bool { + return j.Status == jogStatusCancelled +} + // IsSuccess returns whether the job is successful. func (j *JobInfo) IsSuccess() bool { return j.Status == JobStatusFinished @@ -487,11 +492,32 @@ func GetAllViewableJobs(ctx context.Context, conn sqlexec.SQLExecutor, user stri // CancelJob cancels import into job. Only a running/paused job can be canceled. // check privileges using get before calling this method. func CancelJob(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64) (err error) { + return cancelJobInState(ctx, conn, jobID, jobStatusPending, JobStatusRunning) +} + +// CancelPendingJob cancels a job in pending state. +func CancelPendingJob(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64) (err error) { + return cancelJobInState(ctx, conn, jobID, jobStatusPending) +} + +func cancelJobInState(ctx context.Context, conn sqlexec.SQLExecutor, jobID int64, states ...string) (err error) { + var markerSB strings.Builder + for i := range states { + if i > 0 { + markerSB.WriteString(",") + } + markerSB.WriteString("%?") + } ctx = util.WithInternalSourceType(ctx, kv.InternalImportInto) - sql := `UPDATE mysql.tidb_import_jobs - SET update_time = CURRENT_TIMESTAMP(6), status = %?, error_message = 'cancelled by user' - WHERE id = %? AND status IN (%?, %?);` - args := []any{jogStatusCancelled, jobID, jobStatusPending, JobStatusRunning} + sql := fmt.Sprintf(`UPDATE mysql.tidb_import_jobs + SET update_time = CURRENT_TIMESTAMP(6), status = %%?, error_message = 'cancelled by user' + WHERE id = %%? AND status IN (%s);`, markerSB.String()) + args := make([]any, 0, len(states)+2) + args = append(args, jogStatusCancelled) + args = append(args, jobID) + for _, s := range states { + args = append(args, s) + } _, err = conn.ExecuteInternal(ctx, sql, args...) return err } diff --git a/pkg/executor/importer/job_test.go b/pkg/executor/importer/job_test.go index fc1757c022598..c5ac20e07d055 100644 --- a/pkg/executor/importer/job_test.go +++ b/pkg/executor/importer/job_test.go @@ -238,6 +238,7 @@ func TestGetAndCancelJob(t *testing.T) { // cancel job require.NoError(t, importer.CancelJob(ctx, conn, jobID1)) + require.Equal(t, uint64(1), tk.Session().GetSessionVars().StmtCtx.AffectedRows()) gotJobInfo, err = importer.GetJob(ctx, conn, jobID1, jobInfo.CreatedBy, false) require.NoError(t, err) require.False(t, gotJobInfo.CreateTime.IsZero()) @@ -253,6 +254,7 @@ func TestGetAndCancelJob(t *testing.T) { // call cancel twice is ok, caller should check job status before cancel. require.NoError(t, importer.CancelJob(ctx, conn, jobID1)) + require.Equal(t, uint64(0), tk.Session().GetSessionVars().StmtCtx.AffectedRows()) jobInfo.Status = "pending" jobInfo.ErrorMessage = "" @@ -283,6 +285,7 @@ func TestGetAndCancelJob(t *testing.T) { // cancel job require.NoError(t, importer.CancelJob(ctx, conn, jobID2)) + require.Equal(t, uint64(1), tk.Session().GetSessionVars().StmtCtx.AffectedRows()) gotJobInfo, err = importer.GetJob(ctx, conn, jobID2, jobInfo.CreatedBy, false) require.NoError(t, err) require.False(t, gotJobInfo.CreateTime.IsZero()) @@ -291,6 +294,9 @@ func TestGetAndCancelJob(t *testing.T) { jobInfo.Status = "cancelled" jobInfo.ErrorMessage = "cancelled by user" jobInfoEqual(t, jobInfo, gotJobInfo) + cnt, err = importer.GetActiveJobCnt(ctx, conn, gotJobInfo.TableSchema, gotJobInfo.TableName) + require.NoError(t, err) + require.Equal(t, int64(0), cnt) _, err = importer.GetJob(ctx, conn, 999999999, jobInfo.CreatedBy, false) require.ErrorIs(t, err, exeerrors.ErrLoadDataJobNotFound) @@ -312,6 +318,54 @@ func TestGetAndCancelJob(t *testing.T) { require.Equal(t, jobID2, jobs[1].ID) } +func TestCancelPendingJob(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + ctx := context.Background() + conn := tk.Session().GetSQLExecutor() + parameters := &importer.ImportParameters{ + Format: importer.DataFormatCSV, + } + + pendingJobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, "root@%", "", parameters, 123) + require.NoError(t, err) + cnt, err := importer.GetActiveJobCnt(ctx, conn, "test", "t") + require.NoError(t, err) + require.Equal(t, int64(1), cnt) + + require.NoError(t, importer.CancelPendingJob(ctx, conn, pendingJobID)) + require.Equal(t, uint64(1), tk.Session().GetSessionVars().StmtCtx.AffectedRows()) + pendingJob, err := importer.GetJob(ctx, conn, pendingJobID, "root@%", true) + require.NoError(t, err) + require.Equal(t, "cancelled", pendingJob.Status) + require.True(t, pendingJob.IsCancelled()) + require.Equal(t, "cancelled by user", pendingJob.ErrorMessage) + require.True(t, pendingJob.StartTime.IsZero()) + require.True(t, pendingJob.EndTime.IsZero()) + cnt, err = importer.GetActiveJobCnt(ctx, conn, "test", "t") + require.NoError(t, err) + require.Equal(t, int64(0), cnt) + + runningJobID, err := importer.CreateJob(ctx, conn, "test", "t", 1, "root@%", "", parameters, 123) + require.NoError(t, err) + require.NoError(t, importer.StartJob(ctx, conn, runningJobID, importer.JobStepImporting)) + cnt, err = importer.GetActiveJobCnt(ctx, conn, "test", "t") + require.NoError(t, err) + require.Equal(t, int64(1), cnt) + + require.NoError(t, importer.CancelPendingJob(ctx, conn, runningJobID)) + require.Equal(t, uint64(0), tk.Session().GetSessionVars().StmtCtx.AffectedRows()) + runningJob, err := importer.GetJob(ctx, conn, runningJobID, "root@%", true) + require.NoError(t, err) + require.Equal(t, importer.JobStatusRunning, runningJob.Status) + require.Equal(t, importer.JobStepImporting, runningJob.Step) + require.False(t, runningJob.IsCancelled()) + require.Equal(t, "", runningJob.ErrorMessage) + cnt, err = importer.GetActiveJobCnt(ctx, conn, "test", "t") + require.NoError(t, err) + require.Equal(t, int64(1), cnt) +} + func TestFailJobBeforeStart(t *testing.T) { store := testkit.CreateMockStore(t) tk := testkit.NewTestKit(t, store) @@ -351,17 +405,22 @@ func TestFailJobBeforeStart(t *testing.T) { func TestJobInfo_CanCancel(t *testing.T) { jobInfo := &importer.JobInfo{} for _, c := range []struct { - status string - canCancel bool + status string + canCancel bool + isCancelled bool + isSuccess bool }{ - {"pending", true}, - {"running", true}, - {"finished", false}, - {"failed", false}, - {"canceled", false}, + {status: "pending", canCancel: true}, + {status: "running", canCancel: true}, + {status: "finished", isSuccess: true}, + {status: "failed"}, + {status: "cancelled", isCancelled: true}, + {status: "canceled"}, } { jobInfo.Status = c.status require.Equal(t, c.canCancel, jobInfo.CanCancel(), c.status) + require.Equal(t, c.isCancelled, jobInfo.IsCancelled(), c.status) + require.Equal(t, c.isSuccess, jobInfo.IsSuccess(), c.status) } } diff --git a/tests/realtikvtest/configs/next-gen/pd.toml b/tests/realtikvtest/configs/next-gen/pd.toml index 07569ded421b1..bfe99eff5c1a1 100644 --- a/tests/realtikvtest/configs/next-gen/pd.toml +++ b/tests/realtikvtest/configs/next-gen/pd.toml @@ -1,3 +1,3 @@ # list all the keyspace names here, each keyspace represents a tenant. [keyspace] -pre-alloc = ["keyspace1", "keyspace2", "keyspacecollate"] +pre-alloc = ["keyspace1", "keyspace2", "keyspacecollate", "keyspace_cancel"] diff --git a/tests/realtikvtest/importintotest3/BUILD.bazel b/tests/realtikvtest/importintotest3/BUILD.bazel index c2262524cc068..592024ad22462 100644 --- a/tests/realtikvtest/importintotest3/BUILD.bazel +++ b/tests/realtikvtest/importintotest3/BUILD.bazel @@ -14,6 +14,8 @@ go_test( flaky = True, race = "on", shard_count = 3, + shard_count = 2, + shard_count = 4, deps = [ "//pkg/config", "//pkg/config/deploymode", diff --git a/tests/realtikvtest/importintotest3/cross_ks_test.go b/tests/realtikvtest/importintotest3/cross_ks_test.go index 75a287d3a5ca7..9356d6a9356fd 100644 --- a/tests/realtikvtest/importintotest3/cross_ks_test.go +++ b/tests/realtikvtest/importintotest3/cross_ks_test.go @@ -20,10 +20,12 @@ import ( "fmt" "strconv" "strings" + "sync" "sync/atomic" "testing" "time" + "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/dxf/framework/taskexecutor/execute" @@ -620,3 +622,228 @@ func checkImportTableAndIndexes(tk *testkit.TestKit, tableName string, indexes [ Check(testkit.Rows(expectedCount)) } } + +func TestCancelDanglingImportJobOnUserKeyspace(t *testing.T) { + if kerneltype.IsClassic() { + t.Skip("only runs in nextgen kernel") + } + server, gcsEndpoint := newFakeGCSServer(t) + server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "dangling-cancel-source"}) + server.CreateBucketWithOpts(fakestorage.CreateBucketOpts{Name: "dangling-cancel-sort"}) + server.CreateObject(fakestorage.Object{ + ObjectAttrs: fakestorage.ObjectAttrs{ + BucketName: "dangling-cancel-source", + Name: "data.csv", + }, + Content: []byte("1,a\n2,b\n"), + }) + + const keyspaceName = "keyspace_cancel" + runtimes := realtikvtest.PrepareForCrossKSTest(t, keyspaceName) + userStore := runtimes[keyspaceName].Store + userTK := testkit.NewTestKit(t, userStore) + cancelTK := testkit.NewTestKit(t, userStore) + sysKSTK := testkit.NewTestKit(t, kvstore.GetSystemStorage()) + prepareAndUseDB("cross_ks_cancel", userTK) + cancelTK.MustExec("use cross_ks_cancel") + + sourceURI := fmt.Sprintf("gs://dangling-cancel-source/data.csv?endpoint=%s", gcsEndpoint) + + t.Run("cancel before dxf task submit", func(t *testing.T) { + tableName := "t_cancel_before_task_submit" + createImportTable(userTK, tableName) + + jobIDCh := make(chan int64, 1) + cancelErrCh := make(chan error, 1) + var cancelOnce sync.Once + testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/dxf/importinto/afterUserImportJobCreatedBeforeDXFTask", + func(jobID int64) { + cancelOnce.Do(func() { + jobIDCh <- jobID + cancelErrCh <- cancelTK.ExecToErr(fmt.Sprintf("cancel import job %d", jobID)) + }) + }, + ) + + jobID := submitDetachedJob(t, userTK, tableName, sourceURI, + fmt.Sprintf("gs://dangling-cancel-sort/before-submit?endpoint=%s", gcsEndpoint)) + + var cancelJobID int64 + select { + case cancelJobID = <-jobIDCh: + case <-time.After(30 * time.Second): + t.Fatal("timeout waiting for dangling cancel failpoint") + } + require.EqualValues(t, jobID, cancelJobID) + select { + case err := <-cancelErrCh: + require.NoError(t, err) + case <-time.After(30 * time.Second): + t.Fatal("timeout waiting for cancel import job") + } + + taskKey := importinto.TaskKey(jobID) + userTK.MustQuery("select status, error_message from mysql.tidb_import_jobs where id = ?", jobID). + Check(testkit.Rows("cancelled cancelled by user")) + waitTerminalState(t, sysKSTK, taskKey, proto.TaskStateReverted) + userTK.MustQuery("select count(*) from " + tableName).Check(testkit.Rows("0")) + }) + + t.Run("cancel dangling fallback after dxf task starts next step", func(t *testing.T) { + tableName := "t_cancel_after_task_started" + createImportTable(userTK, tableName) + + cancelAtFallbackCh := make(chan struct{}) + releaseCancelFallbackCh := make(chan struct{}) + var ( + cancelAtFallbackOnce sync.Once + cancelJobID atomic.Int64 + ) + testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/executor/beforeCancelDanglingImportJob", + func(jobID int64) { + cancelJobID.Store(jobID) + cancelAtFallbackOnce.Do(func() { + close(cancelAtFallbackCh) + }) + <-releaseCancelFallbackCh + }, + ) + + sortStartedCh := make(chan struct{}) + releaseSortCh := make(chan struct{}) + var sortStartedOnce sync.Once + testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/dxf/importinto/syncBeforeSortChunk", + func() { + sortStartedOnce.Do(func() { + close(sortStartedCh) + }) + <-releaseSortCh + }, + ) + + cancelErrCh := make(chan error, 1) + var cancelOnce sync.Once + testfailpoint.EnableCall(t, "github.com/pingcap/tidb/pkg/dxf/importinto/afterUserImportJobCreatedBeforeDXFTask", + func(jobID int64) { + cancelOnce.Do(func() { + go func() { + cancelErrCh <- cancelTK.ExecToErr(fmt.Sprintf("cancel import job %d", jobID)) + }() + waitForCancelTestSignal(t, cancelAtFallbackCh, "timeout waiting for cancel to reach dangling fallback") + }) + }, + ) + + jobID := submitDetachedJob(t, userTK, tableName, sourceURI, + fmt.Sprintf("gs://dangling-cancel-sort/after-task-started?endpoint=%s", gcsEndpoint)) + require.EqualValues(t, jobID, cancelJobID.Load()) + + taskKey := importinto.TaskKey(jobID) + waitForCancelTestSignal(t, sortStartedCh, "timeout waiting for import subtask to start") + requireTaskRunningBizStep(t, sysKSTK, taskKey) + + close(releaseCancelFallbackCh) + select { + case err := <-cancelErrCh: + require.ErrorContains(t, err, "job state changed during cancel, please try again later") + case <-time.After(30 * time.Second): + t.Fatal("timeout waiting for cancel import job") + } + userTK.MustQuery("select status from mysql.tidb_import_jobs where id = ?", jobID). + Check(testkit.Rows(importer.JobStatusRunning)) + + close(releaseSortCh) + waitTerminalState(t, sysKSTK, taskKey, proto.TaskStateSucceed) + userTK.MustQuery("select count(*) from " + tableName).Check(testkit.Rows("2")) + }) +} + +func createImportTable(tk *testkit.TestKit, tableName string) { + tk.MustExec("drop table if exists " + tableName) + tk.MustExec(fmt.Sprintf("create table %s (a bigint primary key, b varchar(100));", tableName)) +} + +func submitDetachedJob(t *testing.T, tk *testkit.TestKit, tableName, sourceURI, sortURI string) int64 { + t.Helper() + result := tk.MustQuery(fmt.Sprintf( + "IMPORT INTO %s FROM '%s' WITH DETACHED, cloud_storage_uri='%s'", + tableName, + sourceURI, + sortURI, + )).Rows() + require.Len(t, result, 1) + jobID, err := strconv.ParseInt(result[0][fmap["JobID"]].(string), 10, 64) + require.NoError(t, err) + return jobID +} + +func waitForCancelTestSignal(t *testing.T, ch <-chan struct{}, timeoutMsg string) { + t.Helper() + select { + case <-ch: + case <-time.After(30 * time.Second): + t.Fatal(timeoutMsg) + } +} + +func requireTaskRunningBizStep(t *testing.T, tk *testkit.TestKit, taskKey string) { + t.Helper() + var ( + taskState string + taskStep proto.Step + ) + require.Eventually(t, func() bool { + rows := tk.MustQuery("select state, step from mysql.tidb_global_task where task_key = ?", taskKey).Rows() + if len(rows) != 1 { + return false + } + taskState = rows[0][0].(string) + step, err := strconv.Atoi(rows[0][1].(string)) + if err != nil { + return false + } + taskStep = proto.Step(step) + return taskState == proto.TaskStateRunning.String() && + taskStep != proto.StepInit && + taskStep != proto.StepPrepared + }, 30*time.Second, 500*time.Millisecond) + require.Equal(t, proto.TaskStateRunning.String(), taskState) + require.NotEqual(t, proto.StepInit, taskStep) + require.NotEqual(t, proto.StepPrepared, taskStep) +} + +func waitTerminalState(t *testing.T, tk *testkit.TestKit, taskKey string, expected proto.TaskState) { + t.Helper() + var taskState string + require.Eventually(t, func() bool { + rows := tk.MustQuery(` + select state from mysql.tidb_global_task where task_key = ? + union all + select state from mysql.tidb_global_task_history where task_key = ?`, + taskKey, + taskKey, + ).Rows() + if len(rows) != 1 { + return false + } + taskState = rows[0][0].(string) + return taskState == proto.TaskStateReverted.String() || + taskState == proto.TaskStateSucceed.String() || + taskState == proto.TaskStateFailed.String() + }, 30*time.Second, 500*time.Millisecond) + require.Equal(t, expected.String(), taskState) +} + +func newFakeGCSServer(t *testing.T) (*fakestorage.Server, string) { + t.Helper() + host := "127.0.0.1" + server, err := fakestorage.NewServerWithOptions(fakestorage.Options{ + Scheme: "http", + Host: host, + Port: 0, + PublicHost: host, + }) + require.NoError(t, err) + t.Cleanup(server.Stop) + return server, fmt.Sprintf("%s/storage/v1/", server.URL()) +} From 0e58d402796cf99315d09285e235b839de9c3cb9 Mon Sep 17 00:00:00 2001 From: D3Hunter Date: Thu, 9 Jul 2026 20:33:38 +0800 Subject: [PATCH 2/2] bazel --- pkg/executor/importer/BUILD.bazel | 1 - tests/realtikvtest/importintotest3/BUILD.bazel | 2 -- 2 files changed, 3 deletions(-) diff --git a/pkg/executor/importer/BUILD.bazel b/pkg/executor/importer/BUILD.bazel index bf25898913f34..cd91a27cef493 100644 --- a/pkg/executor/importer/BUILD.bazel +++ b/pkg/executor/importer/BUILD.bazel @@ -120,7 +120,6 @@ go_test( embed = [":importer"], flaky = True, race = "on", - shard_count = 42, shard_count = 43, deps = [ "//br/pkg/mock", diff --git a/tests/realtikvtest/importintotest3/BUILD.bazel b/tests/realtikvtest/importintotest3/BUILD.bazel index 592024ad22462..deca4dd285afd 100644 --- a/tests/realtikvtest/importintotest3/BUILD.bazel +++ b/tests/realtikvtest/importintotest3/BUILD.bazel @@ -13,8 +13,6 @@ go_test( ], flaky = True, race = "on", - shard_count = 3, - shard_count = 2, shard_count = 4, deps = [ "//pkg/config",