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

Signed-off-by: tangenta <[email protected]>

Author: tangenta

pkg/ddl/backfilling_operators.go

@@ -466,6 +466,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 || w.ctx.Err() != nil {
 				w.recycleChunk(srcChk)
@@ -483,17 +487,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/index_cop.go

@@ -149,7 +149,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)
@@ -245,21 +250,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

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

pkg/ddl/ingest/integration_test.go

@@ -376,3 +376,42 @@ 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("set @@global.tidb_enable_dist_task = 0;")
+	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"))
+}