executor: add test for pr 45284

executor/index_merge_reader.go

@@ -56,7 +56,8 @@ import (
 )
 
 var (
-	_ exec.Executor = &IndexMergeReaderExecutor{}
+	_                           exec.Executor = &IndexMergeReaderExecutor{}
+	IndexMergeCancelFuncForTest func()
 )
 
 const (
@@ -1101,13 +1102,16 @@ func (w *indexMergeProcessWorker) fetchLoopUnionWithOrderByAndPushedLimit(ctx co
 			return
 		case <-finished:
 			return
-		case workCh <- task:
+		case resultCh <- task:
+			failpoint.Inject("testCancelContext", func() {
+				IndexMergeCancelFuncForTest()
+			})
 			select {
 			case <-ctx.Done():
 				return
 			case <-finished:
 				return
-			case resultCh <- task:
+			case workCh <- task:
 				continue
 			}
 		}
@@ -1199,13 +1203,16 @@ func (w *indexMergeProcessWorker) fetchLoopUnion(ctx context.Context, fetchCh <-
 			return
 		case <-finished:
 			return
-		case workCh <- task:
+		case resultCh <- task:
+			failpoint.Inject("testCancelContext", func() {
+				IndexMergeCancelFuncForTest()
+			})
 			select {
 			case <-ctx.Done():
 				return
 			case <-finished:
 				return
-			case resultCh <- task:
+			case workCh <- task:
 			}
 		}
 	}

executor/test/indexmergereadtest/index_merge_reader_test.go

@@ -15,6 +15,7 @@
 package indexmergereadtest
 
 import (
+	"context"
 	"fmt"
 	"math/rand"
 	"regexp"
@@ -25,6 +26,8 @@ import (
 	"time"
 
 	"github.com/pingcap/failpoint"
+	"github.com/pingcap/tidb/executor"
+	"github.com/pingcap/tidb/session"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/testkit/testutil"
 	"github.com/pingcap/tidb/util"
@@ -1164,3 +1167,26 @@ func TestProcessInfoRaceWithIndexScan(t *testing.T) {
 	}
 	wg.Wait()
 }
+
+func TestIndexMergeReaderIssue45279(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+
+	tk.MustExec("drop table if exists reproduce;")
+	tk.MustExec("CREATE TABLE reproduce (c1 int primary key, c2 int, c3 int, key ci2(c2), key ci3(c3));")
+	tk.MustExec("insert into reproduce values (1, 1, 1), (2, 2, 2), (3, 3, 3);")
+	tk.MustQuery("explain select * from reproduce where c1 in (0, 1, 2, 3) or c2 in (0, 1, 2);").Check(testkit.Rows(
+		"IndexMerge_11 33.99 root  type: union",
+		"├─TableRangeScan_8(Build) 4.00 cop[tikv] table:reproduce range:[0,0], [1,1], [2,2], [3,3], keep order:false, stats:pseudo",
+		"├─IndexRangeScan_9(Build) 30.00 cop[tikv] table:reproduce, index:ci2(c2) range:[0,0], [1,1], [2,2], keep order:false, stats:pseudo",
+		"└─TableRowIDScan_10(Probe) 33.99 cop[tikv] table:reproduce keep order:false, stats:pseudo"))
+
+	// This function should return successfully
+	var ctx context.Context
+	ctx, executor.IndexMergeCancelFuncForTest = context.WithCancel(context.Background())
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/executor/testCancelContext", "return()"))
+	rs, _ := tk.ExecWithContext(ctx, "select /*+ use_index_merge(reproduce) */ * from reproduce where (c1 < 10 or c2 < 10) and c3 < 10;")
+	session.ResultSetToStringSlice(ctx, tk.Session(), rs)
+	failpoint.Disable("github.com/pingcap/tidb/br/pkg/checksum/testCancelContext")
+}