executor: use statement context for pre-analyze stats flush

Author: 0xPoe

pkg/executor/adapter.go

@@ -419,7 +419,7 @@ func (a *ExecStmt) PointGet(ctx context.Context) (*recordSet, error) {
 	}
 
 	if executor == nil {
-		b := newExecutorBuilder(a.Ctx, a.InfoSchema, a.Ti)
+		b := newExecutorBuilder(ctx, a.Ctx, a.InfoSchema, a.Ti)
 		executor = b.build(a.Plan)
 		if b.err != nil {
 			return nil, b.err
@@ -678,7 +678,7 @@ func (a *ExecStmt) Exec(ctx context.Context) (_ sqlexec.RecordSet, err error) {
 		execStartTime = time.Now()
 	}
 
-	e, err := a.buildExecutor()
+	e, err := a.buildExecutor(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -1371,7 +1371,7 @@ func (a *ExecStmt) handlePessimisticLockError(ctx context.Context, lockErr error
 	a.resetPhaseDurations()
 
 	a.inheritContextFromExecuteStmt()
-	e, err := a.buildExecutor()
+	e, err := a.buildExecutor(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -1397,20 +1397,20 @@ type pessimisticTxn interface {
 }
 
 // buildExecutor build an executor from plan, prepared statement may need additional procedure.
-func (a *ExecStmt) buildExecutor() (exec.Executor, error) {
+func (a *ExecStmt) buildExecutor(ctx context.Context) (exec.Executor, error) {
 	defer func(start time.Time) { a.phaseBuildDurations[0] += time.Since(start) }(time.Now())
-	ctx := a.Ctx
-	stmtCtx := ctx.GetSessionVars().StmtCtx
+	sctx := a.Ctx
+	stmtCtx := sctx.GetSessionVars().StmtCtx
 	if _, ok := a.Plan.(*plannercore.Execute); !ok {
 		if stmtCtx.Priority == mysql.NoPriority && a.LowerPriority {
 			stmtCtx.Priority = kv.PriorityLow
 		}
 	}
-	if _, ok := a.Plan.(*plannercore.Analyze); ok && ctx.GetSessionVars().InRestrictedSQL {
-		ctx.GetSessionVars().StmtCtx.Priority = kv.PriorityLow
+	if _, ok := a.Plan.(*plannercore.Analyze); ok && sctx.GetSessionVars().InRestrictedSQL {
+		sctx.GetSessionVars().StmtCtx.Priority = kv.PriorityLow
 	}
 
-	b := newExecutorBuilder(ctx, a.InfoSchema, a.Ti)
+	b := newExecutorBuilder(ctx, sctx, a.InfoSchema, a.Ti)
 	e := b.build(a.Plan)
 	if b.err != nil {
 		return nil, errors.Trace(b.err)
@@ -1428,7 +1428,7 @@ func (a *ExecStmt) buildExecutor() (exec.Executor, error) {
 			return nil, err
 		}
 		if executorExec.lowerPriority {
-			ctx.GetSessionVars().StmtCtx.Priority = kv.PriorityLow
+			sctx.GetSessionVars().StmtCtx.Priority = kv.PriorityLow
 		}
 		e = executorExec.stmtExec
 	}

pkg/executor/analyze.go

@@ -97,6 +97,9 @@ func flushStatsDeltaForAnalyze(ctx context.Context, sctx sessionctx.Context, pla
 	if len(flushObjects) == 0 {
 		return nil
 	}
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 
 	// HACK: Some tests register in-process TiDB domains but do not start TiDB RPC
 	// endpoints. Broadcasting FLUSH STATS_DELTA CLUSTER to those mock endpoints can

pkg/executor/analyze_test.go

@@ -23,6 +23,7 @@ import (
 	"time"
 
 	"github.com/pingcap/tidb/pkg/domain"
+	"github.com/pingcap/tidb/pkg/executor"
 	"github.com/pingcap/tidb/pkg/infoschema"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/session"
@@ -214,6 +215,25 @@ func TestAnalyzeSaveResultErrorDoesNotHang(t *testing.T) {
 	}
 }
 
+func TestBuildAnalyzePreFlushUsesStatementContext(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t_pre_analyze_flush_ctx(a int, b int, key idx_b(b))")
+	tk.MustExec("insert into t_pre_analyze_flush_ctx values (1, 1), (2, 2)")
+
+	stmtNodes, err := tk.Session().Parse(context.Background(), "analyze table t_pre_analyze_flush_ctx all columns")
+	require.NoError(t, err)
+	require.Len(t, stmtNodes, 1)
+	stmt, err := (&executor.Compiler{Ctx: tk.Session()}).Compile(context.Background(), stmtNodes[0])
+	require.NoError(t, err)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	err = executor.BuildExecutorForTest(ctx, stmt)
+	require.ErrorIs(t, err, context.Canceled)
+}
+
 func TestAnalyzeKillDuringSaveDoesNotHang(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 	tk := testkit.NewTestKit(t, store)

pkg/executor/analyze_utils_test.go

@@ -60,3 +60,10 @@ func TestCanBroadcastToTiDBRPCForTestRejectsInvalidEndpoints(t *testing.T) {
 	// listener. Such targets must not take the broadcast path.
 	require.False(t, canBroadcastToTiDBRPCForTest(context.Background(), []string{"", ""}))
 }
+
+// BuildExecutorForTest builds stmt's executor tree. It is exported only for
+// external package tests that need to assert executor-build behavior.
+func BuildExecutorForTest(ctx context.Context, stmt *ExecStmt) error {
+	_, err := stmt.buildExecutor(ctx)
+	return err
+}

pkg/executor/benchmark_test.go

@@ -67,7 +67,7 @@ func buildHashAggExecutor(ctx sessionctx.Context, src exec.Executor, schema *exp
 	plan.SetSchema(schema)
 	plan.Init(ctx.GetPlanCtx(), nil, 0)
 	plan.SetChildren(nil)
-	b := newExecutorBuilder(ctx, nil, nil)
+	b := newExecutorBuilder(context.Background(), ctx, nil, nil)
 	exec := b.build(plan)
 	hashAgg := exec.(*aggregate.HashAggExec)
 	hashAgg.SetChildren(0, src)
@@ -119,7 +119,7 @@ func buildStreamAggExecutor(ctx sessionctx.Context, srcExec exec.Executor, schem
 		plan = sg
 	}
 
-	b := newExecutorBuilder(ctx, nil, nil)
+	b := newExecutorBuilder(context.Background(), ctx, nil, nil)
 	return b.build(plan)
 }
 
@@ -352,7 +352,7 @@ func buildWindowExecutor(ctx sessionctx.Context, windowFunc string, funcs int, f
 		plan = win
 	}
 
-	b := newExecutorBuilder(ctx, nil, nil)
+	b := newExecutorBuilder(context.Background(), ctx, nil, nil)
 	exec := b.build(plan)
 	return exec
 }
@@ -1253,7 +1253,7 @@ func prepare4IndexInnerHashJoin(tc *IndexJoinTestCase, outerDS *testutil.MockDat
 		keyOff2IdxOff[i] = i
 	}
 
-	readerBuilder, err := newExecutorBuilder(tc.Ctx, nil, nil).
+	readerBuilder, err := newExecutorBuilder(context.Background(), tc.Ctx, nil, nil).
 		newDataReaderBuilder(&mockPhysicalIndexReader{e: innerDS})
 	if err != nil {
 		return nil, err
@@ -1328,7 +1328,7 @@ func prepare4IndexMergeJoin(tc *IndexJoinTestCase, outerDS *testutil.MockDataSou
 		outerCompareFuncs = append(outerCompareFuncs, expression.GetCmpFunction(nil, outerJoinKeys[i], outerJoinKeys[i]))
 	}
 
-	readerBuilder, err := newExecutorBuilder(tc.Ctx, nil, nil).
+	readerBuilder, err := newExecutorBuilder(context.Background(), tc.Ctx, nil, nil).
 		newDataReaderBuilder(&mockPhysicalIndexReader{e: innerDS})
 	if err != nil {
 		return nil, err

pkg/executor/builder.go

@@ -104,6 +104,9 @@ import (
 // executorBuilder builds an Executor from a Plan.
 // The InfoSchema must not change during execution.
 type executorBuilder struct {
+	// stmtCtx is the statement-scoped context.Context for executor build steps that
+	// can block before Executor.Open/Next receives the execution context.
+	stmtCtx context.Context
 	sctx    sessionctx.Context
 	is      infoschema.InfoSchema
 	err     error // err is set when there is error happened during Executor building process.
@@ -146,9 +149,10 @@ type CTEStorages struct {
 	initErr error
 }
 
-func newExecutorBuilder(sctx sessionctx.Context, is infoschema.InfoSchema, ti *TelemetryInfo) *executorBuilder {
+func newExecutorBuilder(stmtCtx context.Context, sctx sessionctx.Context, is infoschema.InfoSchema, ti *TelemetryInfo) *executorBuilder {
 	txnManager := sessiontxn.GetTxnManager(sctx)
 	return &executorBuilder{
+		stmtCtx:          stmtCtx,
 		sctx:             sctx,
 		is:               is,
 		Ti:               ti,
@@ -177,7 +181,7 @@ type MockExecutorBuilder struct {
 // NewMockExecutorBuilderForTest is ONLY used in test.
 func NewMockExecutorBuilderForTest(ctx sessionctx.Context, is infoschema.InfoSchema, ti *TelemetryInfo) *MockExecutorBuilder {
 	return &MockExecutorBuilder{
-		executorBuilder: newExecutorBuilder(ctx, is, ti),
+		executorBuilder: newExecutorBuilder(context.Background(), ctx, is, ti),
 	}
 }
 
@@ -3315,8 +3319,8 @@ func (b *executorBuilder) buildAnalyze(v *plannercore.Analyze) exec.Executor {
 	// buildAnalyzeSamplingPushdown reads base count / modify_count from mysql.stats_meta
 	// while constructing column analyze tasks. Flush pending deltas first so the base
 	// values include pre-analyze changes and later delta dumps cannot double count them.
-	// TODO: Determine whether context.Background is appropriate here; if not, use the proper statement context.
-	if err := flushStatsDeltaForAnalyze(kv.WithInternalSourceType(context.Background(), kv.InternalTxnStats), b.sctx, v); err != nil {
+	intest.Assert(b.stmtCtx != nil, "missing statement context for analyze")
+	if err := flushStatsDeltaForAnalyze(kv.WithInternalSourceType(b.stmtCtx, kv.InternalTxnStats), b.sctx, v); err != nil {
 		b.err = err
 		return nil
 	}

pkg/executor/builder_index_join_cleanup_test.go

@@ -81,7 +81,7 @@ func TestBuildExecutorForIndexJoinHashJoinErrorCleansChildren(t *testing.T) {
 	hashJoinPlan.SetSchema(expression.MergeSchema(lookupSchema, otherSchema))
 	hashJoinPlan.SetChildren(lookupMockPlan, otherMockPlan)
 
-	execBuilder := newExecutorBuilder(ctx, nil, nil)
+	execBuilder := newExecutorBuilder(context.Background(), ctx, nil, nil)
 	execBuilder.forDataReaderBuilder = true
 	execBuilder.dataReaderTS = 1
 	readerBuilder, err := execBuilder.newDataReaderBuilder(hashJoinPlan)
@@ -106,7 +106,7 @@ func TestBuildExecutorForIndexJoinHashJoinErrorCleansChildren(t *testing.T) {
 func TestBuildCTEStorageProducerCleansStoragesOnRecursiveBuildError(t *testing.T) {
 	ctx := mock.NewContext()
 	ctx.GetSessionVars().StmtCtx.CTEStorageMap = map[int]*CTEStorages{}
-	builder := newExecutorBuilder(ctx, nil, nil)
+	builder := newExecutorBuilder(context.Background(), ctx, nil, nil)
 
 	stats := &property.StatsInfo{RowCount: 1}
 	schema := expression.NewSchema(&expression.Column{

pkg/executor/coprocessor.go

@@ -199,7 +199,7 @@ func (h *CoprocessorDAGHandler) buildDAGExecutor(ctx context.Context, req *copro
 	}
 	plan = core.InjectExtraProjection(plan)
 	// Build executor.
-	b := newExecutorBuilder(h.sctx, is, nil)
+	b := newExecutorBuilder(ctx, h.sctx, is, nil)
 	return b.build(plan), nil
 }
 

pkg/executor/executor_required_rows_test.go

@@ -741,7 +741,7 @@ func buildMergeJoinExec(ctx sessionctx.Context, joinType base.JoinType, innerSrc
 		j.CompareFuncs = append(j.CompareFuncs, expression.GetCmpFunction(ctx.GetExprCtx(), j.LeftJoinKeys[i], j.RightJoinKeys[i]))
 	}
 
-	b := newExecutorBuilder(ctx, nil, nil)
+	b := newExecutorBuilder(context.Background(), ctx, nil, nil)
 	return b.build(j)
 }
 

pkg/executor/select.go

@@ -616,7 +616,7 @@ func init() {
 			return nil, err
 		}
 
-		e := newExecutorBuilder(sctx, is, nil)
+		e := newExecutorBuilder(ctx, sctx, is, nil)
 		executor := e.build(p)
 		if e.err != nil {
 			return nil, e.err