copIterator: return context error to avoid return incorrect result on context cancel/timeout

pkg/ddl/backfilling_operators.go

@@ -750,7 +750,9 @@ func (w *indexIngestBaseWorker) WriteChunk(rs *IndexRecordChunk) (count int, nex
 		failpoint.Return(0, nil, errors.New("mock write local error"))
 	})
 	failpoint.Inject("writeLocalExec", func(_ failpoint.Value) {
-		OperatorCallBackForTest()
+		if rs.Done {
+			OperatorCallBackForTest()
+		}
 	})
 
 	oprStartTime := time.Now()

pkg/disttask/framework/taskexecutor/task_executor.go

@@ -131,14 +131,6 @@ func (e *BaseTaskExecutor) checkBalanceSubtask(ctx context.Context) {
 			e.logger.Error("get subtasks failed", zap.Error(err))
 			continue
 		}
-		if ctx.Err() != nil {
-			// workaround for https://github.com/pingcap/tidb/issues/50089
-			// timeline to trigger this:
-			// 	- this routine runs GetSubtasksByExecIDAndStepAndStates
-			// 	- outer runSubtask finishes and cancel check-context
-			// 	- GetSubtasksByExecIDAndStepAndStates returns with no err and no result
-			return
-		}
 		if len(subtasks) == 0 {
 			e.logger.Info("subtask is scheduled away, cancel running")
 			// cancels runStep, but leave the subtask state unchanged.

pkg/executor/adapter_test.go

@@ -15,11 +15,17 @@
 package executor_test
 
 import (
+	"context"
+	"sync"
 	"testing"
 
 	"github.com/pingcap/tidb/pkg/executor"
+	"github.com/pingcap/tidb/pkg/session"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/store/copr"
+	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/util"
 )
 
 func TestFormatSQL(t *testing.T) {
@@ -32,3 +38,30 @@ func TestFormatSQL(t *testing.T) {
 	val = executor.FormatSQL("aaaaaaaaaaaaaaaaaaaa")
 	require.Equal(t, "aaaaa(len:20)", val.String())
 }
+
+func TestContextCancelWhenReadFromCopIterator(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t(a int)")
+	tk.MustExec("insert into t values(1)")
+
+	testkit.EnableFailPoint(t, "github.com/pingcap/tidb/pkg/store/copr/CtxCancelBeforeReceive", "return(true)")
+	ctx := context.WithValue(context.Background(), "TestContextCancel", "test")
+	ctx, cancelFunc := context.WithCancel(ctx)
+	defer cancelFunc()
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ctx = util.WithInternalSourceType(ctx, "scheduler")
+		rs, err := tk.Session().ExecuteInternal(ctx, "select * from test.t")
+		require.NoError(t, err)
+		_, err2 := session.ResultSetToStringSlice(ctx, tk.Session(), rs)
+		require.ErrorIs(t, err2, context.Canceled)
+	}()
+	<-copr.GlobalSyncChForTest
+	cancelFunc()
+	copr.GlobalSyncChForTest <- struct{}{}
+	wg.Wait()
+}

pkg/store/copr/coprocessor.go

@@ -935,7 +935,16 @@ func (sender *copIteratorTaskSender) run(connID uint64) {
 	}
 }
 
+// GlobalSyncChForTest is a global channel for test.
+var GlobalSyncChForTest = make(chan struct{})
+
 func (it *copIterator) recvFromRespCh(ctx context.Context, respCh <-chan *copResponse) (resp *copResponse, ok bool, exit bool) {
+	failpoint.Inject("CtxCancelBeforeReceive", func(_ failpoint.Value) {
+		if ctx.Value("TestContextCancel") == "test" {
+			GlobalSyncChForTest <- struct{}{}
+			<-GlobalSyncChForTest
+		}
+	})
 	ticker := time.NewTicker(3 * time.Second)
 	defer ticker.Stop()
 	for {
@@ -1072,7 +1081,7 @@ func (it *copIterator) Next(ctx context.Context) (kv.ResultSubset, error) {
 		resp, ok, closed = it.recvFromRespCh(ctx, it.respChan)
 		if !ok || closed {
 			it.actionOnExceed.close()
-			return nil, nil
+			return nil, errors.Trace(ctx.Err())
 		}
 		if resp == finCopResp {
 			it.actionOnExceed.destroyTokenIfNeeded(func() {
@@ -1090,8 +1099,8 @@ func (it *copIterator) Next(ctx context.Context) (kv.ResultSubset, error) {
 			task := it.tasks[it.curr]
 			resp, ok, closed = it.recvFromRespCh(ctx, task.respChan)
 			if closed {
-				// Close() is already called, so Next() is invalid.
-				return nil, nil
+				// Close() is called or context cancelled/timeout, so Next() is invalid.
+				return nil, errors.Trace(ctx.Err())
 			}
 			if ok {
 				break

tests/realtikvtest/addindextest3/operator_test.go

@@ -17,6 +17,7 @@ package addindextest
 import (
 	"context"
 	"fmt"
+	"strings"
 	"sync/atomic"
 	"testing"
 
@@ -234,7 +235,7 @@ func TestBackfillOperatorPipelineException(t *testing.T) {
 		{
 			failPointPath:  "github.com/pingcap/tidb/pkg/ddl/scanRecordExec",
 			closeErrMsg:    "context canceled",
-			operatorErrMsg: "",
+			operatorErrMsg: "context canceled",
 		},
 		{
 			failPointPath:  "github.com/pingcap/tidb/pkg/ddl/mockWriteLocalError",
@@ -254,44 +255,62 @@ func TestBackfillOperatorPipelineException(t *testing.T) {
 	}
 
 	for _, tc := range testCase {
-		require.NoError(t, failpoint.Enable(tc.failPointPath, `return`))
-		ctx, cancel := context.WithCancel(context.Background())
-		ddl.OperatorCallBackForTest = func() {
+		t.Run(tc.failPointPath, func(t *testing.T) {
+			require.NoError(t, failpoint.Enable(tc.failPointPath, `return`))
+			defer func() {
+				require.NoError(t, failpoint.Disable(tc.failPointPath))
+			}()
+			ctx, cancel := context.WithCancel(context.Background())
+			if strings.Contains(tc.failPointPath, "writeLocalExec") {
+				var counter atomic.Int32
+				ddl.OperatorCallBackForTest = func() {
+					// we need to want all tableScanWorkers finish scanning, else
+					// fetchTableScanResult will might return context error, and cause
+					// the case fail.
+					// 10 is the table scan task count.
+					counter.Add(1)
+					if counter.Load() == 10 {
+						cancel()
+					}
+				}
+			} else {
+				ddl.OperatorCallBackForTest = func() {
+					cancel()
+				}
+			}
+			opCtx := ddl.NewOperatorCtx(ctx, 1, 1)
+			pipeline, err := ddl.NewAddIndexIngestPipeline(
+				opCtx, store,
+				sessPool,
+				mockBackendCtx,
+				[]ingest.Engine{mockEngine},
+				tk.Session(),
+				1, // job id
+				tbl.(table.PhysicalTable),
+				[]*model.IndexInfo{idxInfo},
+				startKey,
+				endKey,
+				&atomic.Int64{},
+				nil,
+				ddl.NewDDLReorgMeta(tk.Session()),
+				0,
+				2,
+			)
+			require.NoError(t, err)
+			err = pipeline.Execute()
+			require.NoError(t, err)
+			err = pipeline.Close()
+			comment := fmt.Sprintf("case: %s", tc.failPointPath)
+			require.ErrorContains(t, err, tc.closeErrMsg, comment)
+			opCtx.Cancel()
+			if tc.operatorErrMsg == "" {
+				require.NoError(t, opCtx.OperatorErr())
+			} else {
+				require.Error(t, opCtx.OperatorErr())
+				require.Equal(t, tc.operatorErrMsg, opCtx.OperatorErr().Error())
+			}
 			cancel()
-		}
-		opCtx := ddl.NewOperatorCtx(ctx, 1, 1)
-		pipeline, err := ddl.NewAddIndexIngestPipeline(
-			opCtx, store,
-			sessPool,
-			mockBackendCtx,
-			[]ingest.Engine{mockEngine},
-			tk.Session(),
-			1, // job id
-			tbl.(table.PhysicalTable),
-			[]*model.IndexInfo{idxInfo},
-			startKey,
-			endKey,
-			&atomic.Int64{},
-			nil,
-			ddl.NewDDLReorgMeta(tk.Session()),
-			0,
-			2,
-		)
-		require.NoError(t, err)
-		err = pipeline.Execute()
-		require.NoError(t, err)
-		err = pipeline.Close()
-		comment := fmt.Sprintf("case: %s", tc.failPointPath)
-		require.ErrorContains(t, err, tc.closeErrMsg, comment)
-		opCtx.Cancel()
-		if tc.operatorErrMsg == "" {
-			require.NoError(t, opCtx.OperatorErr())
-		} else {
-			require.Error(t, opCtx.OperatorErr())
-			require.Equal(t, tc.operatorErrMsg, opCtx.OperatorErr().Error())
-		}
-		require.NoError(t, failpoint.Disable(tc.failPointPath))
-		cancel()
+		})
 	}
 }