copIterator: return context error to avoid return incorrect result on context cancel/timeout (#53489)

executor/adapter_test.go

@@ -15,13 +15,19 @@
 package executor_test
 
 import (
+	"context"
+	"sync"
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/executor"
+	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/store/copr"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/stretchr/testify/require"
+	"github.com/tikv/client-go/v2/util"
 )
 
 func TestQueryTime(t *testing.T) {
@@ -52,3 +58,33 @@ 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)")
+
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/store/copr/CtxCancelBeforeReceive", "return(true)"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/store/copr/CtxCancelBeforeReceive"))
+	}()
+	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()
+}

store/copr/coprocessor.go

@@ -746,7 +746,16 @@ func (sender *copIteratorTaskSender) run() {
 	}
 }
 
+// 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 {
@@ -868,7 +877,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() {
@@ -886,8 +895,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