ddl: fix cancelling stuck issue during adding index (pingcap#61541)

close pingcap#61087

Author: tangenta

pkg/ddl/backfilling.go

@@ -338,7 +338,7 @@ func (w *backfillWorker) sendResult(result *backfillResult) {
 }
 
 func (w *backfillWorker) run(d *ddlCtx, bf backfiller, job *model.Job) {
-	logger := ddlLogger.With(zap.Stringer("worker", w), zap.Int64("jobID", job.ID))
+	logger := logutil.BgLogger().With(zap.String("category", "ddl"), zap.Stringer("worker", w), zap.Int64("jobID", job.ID))
 	var (
 		curTaskID int
 		task      *reorgBackfillTask
@@ -634,6 +634,7 @@ func (dc *ddlCtx) writePhysicalTableRecord(
 	bfWorkerType backfillerType,
 	reorgInfo *reorgInfo,
 ) error {
+	logger := logutil.BgLogger().With(zap.String("category", "ddl"))
 	startKey, endKey := reorgInfo.StartKey, reorgInfo.EndKey
 
 	if err := dc.isReorgRunnable(reorgInfo.Job.ID, false); err != nil {
@@ -674,7 +675,7 @@ func (dc *ddlCtx) writePhysicalTableRecord(
 				return egCtx.Err()
 			case result, ok := <-scheduler.resultChan():
 				if !ok {
-					ddlLogger.Info("backfill workers successfully processed",
+					logger.Info("backfill workers successfully processed",
 						zap.Stringer("element", reorgInfo.currElement),
 						zap.Int64("total added count", totalAddedCount),
 						zap.String("start key", hex.EncodeToString(startKey)))
@@ -683,7 +684,7 @@ func (dc *ddlCtx) writePhysicalTableRecord(
 				cnt++
 
 				if result.err != nil {
-					ddlLogger.Warn("backfill worker failed",
+					logger.Warn("backfill worker failed",
 						zap.Int64("job ID", reorgInfo.ID),
 						zap.Int64("total added count", totalAddedCount),
 						zap.String("start key", hex.EncodeToString(startKey)),
@@ -704,15 +705,15 @@ func (dc *ddlCtx) writePhysicalTableRecord(
 				if cnt%(scheduler.currentWorkerSize()*4) == 0 {
 					err2 := reorgInfo.UpdateReorgMeta(keeper.nextKey, sessPool)
 					if err2 != nil {
-						ddlLogger.Warn("update reorg meta failed",
+						logger.Warn("update reorg meta failed",
 							zap.Int64("job ID", reorgInfo.ID),
 							zap.Error(err2))
 					}
 					// We try to adjust the worker size regularly to reduce
 					// the overhead of loading the DDL related global variables.
 					err2 = scheduler.adjustWorkerSize()
 					if err2 != nil {
-						ddlLogger.Warn("cannot adjust backfill worker size",
+						logger.Warn("cannot adjust backfill worker size",
 							zap.Int64("job ID", reorgInfo.ID),
 							zap.Error(err2))
 					}
@@ -735,7 +736,7 @@ func (dc *ddlCtx) writePhysicalTableRecord(
 			if len(kvRanges) == 0 {
 				break
 			}
-			ddlLogger.Info("start backfill workers to reorg record",
+			logutil.BgLogger().Info("start backfill workers to reorg record",
 				zap.Stringer("type", bfWorkerType),
 				zap.Int("workerCnt", scheduler.currentWorkerSize()),
 				zap.Int("regionCnt", len(kvRanges)),

pkg/ddl/backfilling_operators.go

@@ -451,6 +451,10 @@ func (w *tableScanWorker) scanRecords(task TableScanTask, sender func(IndexRecor
 		var done bool
 		for !done {
 			srcChk := w.getChunk()
+			if srcChk == nil {
+				terror.Call(rs.Close)
+				return err
+			}
 			done, err = fetchTableScanResult(w.ctx, w.copCtx.GetBase(), rs, srcChk)
 			if err != nil || util2.IsContextDone(w.ctx) {
 				w.recycleChunk(srcChk)
@@ -468,17 +472,25 @@ func (w *tableScanWorker) scanRecords(task TableScanTask, sender func(IndexRecor
 }
 
 func (w *tableScanWorker) getChunk() *chunk.Chunk {
-	chk := <-w.srcChkPool
-	newCap := copReadBatchSize()
-	if chk.Capacity() != newCap {
-		chk = chunk.NewChunkWithCapacity(w.copCtx.GetBase().FieldTypes, newCap)
+	select {
+	case <-w.ctx.Done():
+		return nil
+	case chk := <-w.srcChkPool:
+		newCap := copReadBatchSize()
+		if chk.Capacity() != newCap {
+			chk = chunk.NewChunkWithCapacity(w.copCtx.GetBase().FieldTypes, newCap)
+		}
+		chk.Reset()
+		return chk
 	}
-	chk.Reset()
-	return chk
 }
 
 func (w *tableScanWorker) recycleChunk(chk *chunk.Chunk) {
-	w.srcChkPool <- chk
+	select {
+	case <-w.ctx.Done():
+		return
+	case w.srcChkPool <- chk:
+	}
 }
 
 // WriteExternalStoreOperator writes index records to external storage.

pkg/ddl/ddl.go

@@ -98,8 +98,6 @@ const (
 	recoverCheckFlagDisableGC
 )
 
-var ddlLogger = logutil.BgLogger().With(zap.String("category", "ddl"))
-
 // OnExist specifies what to do when a new object has a name collision.
 type OnExist uint8
 

pkg/ddl/index_cop.go

@@ -155,7 +155,12 @@ func scanRecords(p *copReqSenderPool, task *reorgBackfillTask, se *sess.Session)
 		var done bool
 		startTime := time.Now()
 		for !done {
+			failpoint.InjectCall("beforeGetChunk")
 			srcChk := p.getChunk()
+			if srcChk == nil {
+				terror.Call(rs.Close)
+				return err
+			}
 			done, err = fetchTableScanResult(p.ctx, p.copCtx.GetBase(), rs, srcChk)
 			if err != nil {
 				p.recycleChunk(srcChk)
@@ -253,21 +258,29 @@ func (c *copReqSenderPool) close(force bool) {
 }
 
 func (c *copReqSenderPool) getChunk() *chunk.Chunk {
-	chk := <-c.srcChkPool
-	newCap := copReadBatchSize()
-	if chk.Capacity() != newCap {
-		chk = chunk.NewChunkWithCapacity(c.copCtx.GetBase().FieldTypes, newCap)
+	select {
+	case <-c.ctx.Done():
+		return nil
+	case chk := <-c.srcChkPool:
+		newCap := copReadBatchSize()
+		if chk.Capacity() != newCap {
+			chk = chunk.NewChunkWithCapacity(c.copCtx.GetBase().FieldTypes, newCap)
+		}
+		chk.Reset()
+		return chk
 	}
-	chk.Reset()
-	return chk
 }
 
 // recycleChunk puts the index record slice and the chunk back to the pool for reuse.
 func (c *copReqSenderPool) recycleChunk(chk *chunk.Chunk) {
 	if chk == nil {
 		return
 	}
-	c.srcChkPool <- chk
+	select {
+	case <-c.ctx.Done():
+		return
+	case c.srcChkPool <- chk:
+	}
 }
 
 func buildTableScan(ctx context.Context, c *copr.CopContextBase, startTS uint64, start, end kv.Key) (distsql.SelectResult, error) {

pkg/ddl/ingest/BUILD.bazel

@@ -67,7 +67,7 @@ go_test(
     embed = [":ingest"],
     flaky = True,
     race = "on",
-    shard_count = 18,
+    shard_count = 19,
     deps = [
         "//pkg/config",
         "//pkg/ddl",

pkg/ddl/ingest/integration_test.go

@@ -402,3 +402,41 @@ func TestMultiSchemaAddIndexMerge(t *testing.T) {
 		tk.MustExec("admin check table t;")
 	}
 }
+
+func TestAddIndexGetChunkCancel(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test;")
+	defer ingesttestutil.InjectMockBackendMgr(t, store)()
+
+	tk.MustExec("create table t (a int primary key, b int);")
+	for i := 0; i < 100; i++ {
+		tk.MustExec(fmt.Sprintf("insert into t values (%d, %d);", i*10000, i*10000))
+	}
+	tk.MustExec("split table t between (0) and (1000000) regions 10;")
+	jobID := int64(0)
+	hook := &callback.TestDDLCallback{Do: dom}
+	hook.OnJobRunBeforeExported = func(job *model.Job) {
+		if jobID != 0 {
+			return
+		}
+		if job.Type == model.ActionAddIndex {
+			jobID = job.ID
+		}
+	}
+	dom.DDL().SetHook(hook)
+
+	cancelled := false
+	err := failpoint.EnableCall("github.com/pingcap/tidb/pkg/ddl/beforeGetChunk", func() {
+		if !cancelled {
+			tk2 := testkit.NewTestKit(t, store)
+			tk2.MustExec(fmt.Sprintf("admin cancel ddl jobs %d", jobID))
+			cancelled = true
+		}
+	})
+	require.NoError(t, err)
+	tk.MustGetErrCode("alter table t add index idx(b);", errno.ErrCancelledDDLJob)
+	require.True(t, cancelled)
+	tk.MustExec("admin check table t;")
+	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/ddl/beforeGetChunk"))
+}