executor: Fix tidb crash on index merge reader (#40904)

close #40877

executor/index_merge_reader.go

@@ -612,8 +612,13 @@ func (e *IndexMergeReaderExecutor) startIndexMergeTableScanWorker(ctx context.Co
 			defer trace.StartRegion(ctx, "IndexMergeTableScanWorker").End()
 			var task *indexMergeTableTask
 			util.WithRecovery(
-				func() { task = worker.pickAndExecTask(ctx1) },
-				worker.handlePickAndExecTaskPanic(ctx1, task),
+				// Note we use the address of `task` as the argument of both `pickAndExecTask` and `handlePickAndExecTaskPanic`
+				// because `task` is expected to be assigned in `pickAndExecTask`, and this assignment should also be visible
+				// in `handlePickAndExecTaskPanic` since it will get `doneCh` from `task`. Golang always pass argument by value,
+				// so if we don't use the address of `task` as the argument, the assignment to `task` in `pickAndExecTask` is
+				// not visible in `handlePickAndExecTaskPanic`
+				func() { worker.pickAndExecTask(ctx1, &task) },
+				worker.handlePickAndExecTaskPanic(ctx1, &task),
 			)
 			cancel()
 			e.tblWorkerWg.Done()
@@ -1107,38 +1112,41 @@ type indexMergeTableScanWorker struct {
 	memTracker *memory.Tracker
 }
 
-func (w *indexMergeTableScanWorker) pickAndExecTask(ctx context.Context) (task *indexMergeTableTask) {
+func (w *indexMergeTableScanWorker) pickAndExecTask(ctx context.Context, task **indexMergeTableTask) {
 	var ok bool
 	for {
 		waitStart := time.Now()
 		select {
-		case task, ok = <-w.workCh:
+		case *task, ok = <-w.workCh:
 			if !ok {
 				return
 			}
 		case <-w.finished:
 			return
 		}
 		execStart := time.Now()
-		err := w.executeTask(ctx, task)
+		err := w.executeTask(ctx, *task)
 		if w.stats != nil {
 			atomic.AddInt64(&w.stats.WaitTime, int64(execStart.Sub(waitStart)))
 			atomic.AddInt64(&w.stats.FetchRow, int64(time.Since(execStart)))
 			atomic.AddInt64(&w.stats.TableTaskNum, 1)
 		}
-		task.doneCh <- err
+		failpoint.Inject("testIndexMergePickAndExecTaskPanic", nil)
+		(*task).doneCh <- err
 	}
 }
 
-func (w *indexMergeTableScanWorker) handlePickAndExecTaskPanic(ctx context.Context, task *indexMergeTableTask) func(r interface{}) {
+func (w *indexMergeTableScanWorker) handlePickAndExecTaskPanic(ctx context.Context, task **indexMergeTableTask) func(r interface{}) {
 	return func(r interface{}) {
 		if r == nil {
 			return
 		}
 
 		err4Panic := errors.Errorf("panic in IndexMergeReaderExecutor indexMergeTableWorker: %v", r)
 		logutil.Logger(ctx).Error(err4Panic.Error())
-		task.doneCh <- err4Panic
+		if *task != nil {
+			(*task).doneCh <- err4Panic
+		}
 	}
 }
 

executor/index_merge_reader_test.go

@@ -51,6 +51,25 @@ func TestSingleTableRead(t *testing.T) {
 	tk.MustQuery("select /*+ use_index_merge(t1, t1a, t1b) */ sum(a) from t1 where a < 2 or b > 4").Check(testkit.Rows("6"))
 }
 
+func TestIndexMergePickAndExecTaskPanic(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t1, t2")
+	tk.MustExec("create table t1(id int primary key, a int, b int, c int, d int)")
+	tk.MustExec("create index t1a on t1(a)")
+	tk.MustExec("create index t1b on t1(b)")
+	tk.MustExec("insert into t1 values(1,1,1,1,1),(2,2,2,2,2),(3,3,3,3,3),(4,4,4,4,4),(5,5,5,5,5)")
+	tk.MustQuery("select /*+ use_index_merge(t1, primary, t1a) */ * from t1 where id < 2 or a > 4 order by id").Check(testkit.Rows("1 1 1 1 1",
+		"5 5 5 5 5"))
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/executor/testIndexMergePickAndExecTaskPanic", "panic(\"pickAndExecTaskPanic\")"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/executor/testIndexMergePickAndExecTaskPanic"))
+	}()
+	err := tk.QueryToErr("select /*+ use_index_merge(t1, primary, t1a) */ * from t1 where id < 2 or a > 4 order by id")
+	require.Contains(t, err.Error(), "pickAndExecTaskPanic")
+}
+
 func TestJoin(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)