planner: revise optimize tracer code structure (#31739)

ref #29661

executor/trace.go

@@ -41,6 +41,7 @@ import (
 	"github.com/pingcap/tidb/util/chunk"
 	"github.com/pingcap/tidb/util/logutil"
 	"github.com/pingcap/tidb/util/sqlexec"
+	"github.com/pingcap/tidb/util/tracing"
 	"go.uber.org/zap"
 	"sourcegraph.com/sourcegraph/appdash"
 	traceImpl "sourcegraph.com/sourcegraph/appdash/opentracing"
@@ -161,7 +162,9 @@ func (e *TraceExec) nextOptimizerPlanTrace(ctx context.Context, se sessionctx.Co
 	jsonEncoder := json.NewEncoder(&writer)
 	// If we do not set this to false, ">", "<", "&"... will be escaped to "\u003c","\u003e", "\u0026"...
 	jsonEncoder.SetEscapeHTML(false)
-	err = jsonEncoder.Encode(se.GetSessionVars().StmtCtx.LogicalOptimizeTrace)
+	logical := se.GetSessionVars().StmtCtx.LogicalOptimizeTrace
+	physical := se.GetSessionVars().StmtCtx.PhysicalOptimizeTrace
+	err = jsonEncoder.Encode(&tracing.OptimizeTracer{Logical: logical, Physical: physical})
 	if err != nil {
 		return errors.AddStack(err)
 	}

planner/core/find_best_task.go

@@ -207,11 +207,12 @@ func (p *baseLogicalPlan) rebuildChildTasks(childTasks *[]task, pp PhysicalPlan,
 func (p *baseLogicalPlan) enumeratePhysicalPlans4Task(physicalPlans []PhysicalPlan, prop *property.PhysicalProperty, addEnforcer bool, planCounter *PlanCounterTp, opt *physicalOptimizeOp) (task, int64, error) {
 	var bestTask task = invalidTask
 	var curCntPlan, cntPlan int64
+	childProps := make(map[task]*property.PhysicalProperty)
 	childTasks := make([]task, 0, len(p.children))
 	childCnts := make([]int64, len(p.children))
 	cntPlan = 0
 	for _, pp := range physicalPlans {
-		candidateInfo := opt.appendCandidate(p, pp, prop.String())
+
 		// Find best child tasks firstly.
 		childTasks = childTasks[:0]
 		// The curCntPlan records the number of possible plans for pp
@@ -230,9 +231,7 @@ func (p *baseLogicalPlan) enumeratePhysicalPlans4Task(physicalPlans []PhysicalPl
 				break
 			}
 			childTasks = append(childTasks, childTask)
-			if opt != nil && childTask != nil {
-				opt.appendChildToCandidate(candidateInfo, childTask.plan())
-			}
+			childProps[childTask] = childProp
 		}
 
 		// This check makes sure that there is no invalid child task.
@@ -280,9 +279,14 @@ func (p *baseLogicalPlan) enumeratePhysicalPlans4Task(physicalPlans []PhysicalPl
 			bestTask = curTask
 			break
 		}
-		if candidateInfo != nil {
-			candidateInfo.SetCost(curTask.cost())
-			candidateInfo.Info = curTask.plan().ExplainInfo()
+		c := opt.appendCandidate(p, curTask.plan(), prop)
+		if c != nil {
+			for _, childTask := range childTasks {
+				if childTask == nil {
+					continue
+				}
+				opt.appendChildToCandidate(c, childTask.plan(), childProps[childTask])
+			}
 		}
 		// Get the most efficient one.
 		if curTask.cost() < bestTask.cost() || (bestTask.invalid() && !curTask.invalid()) {
@@ -322,34 +326,25 @@ func (op *physicalOptimizeOp) buildPhysicalOptimizeTraceInfo(p LogicalPlan, prop
 	return traceInfo
 }
 
-func (op *physicalOptimizeOp) appendChildToCandidate(candidateInfo *tracing.PhysicalPlanTrace, plan PhysicalPlan) {
+func (op *physicalOptimizeOp) appendChildToCandidate(candidateInfo *tracing.PlanTrace, plan PhysicalPlan, prop *property.PhysicalProperty) {
 	if op == nil || op.tracer == nil || candidateInfo == nil {
 		return
 	}
-	childPhysicalPlanTrace := &tracing.PhysicalPlanTrace{TP: plan.TP(), ID: plan.ID(), Info: plan.ExplainInfo(), Cost: plan.Cost()}
+	childPhysicalPlanTrace := &tracing.PlanTrace{TP: plan.TP(), ID: plan.ID(), ExplainInfo: plan.ExplainInfo(), Cost: plan.Cost(), ProperType: prop.String()}
 	candidateInfo.Children = append(candidateInfo.Children, childPhysicalPlanTrace)
 }
 
-func (op *physicalOptimizeOp) setBest(lp LogicalPlan, pp PhysicalPlan, prop string) {
-	if op == nil || op.tracer == nil || pp == nil {
-		return
-	}
-	traceInfo := op.tracer.State[fmt.Sprintf("%v_%v", lp.TP(), lp.ID())][prop]
-	if traceInfo == nil {
-		return
-	}
-	traceInfo.BestTask = &tracing.PhysicalPlanTrace{ID: pp.ID(), TP: pp.TP(), Cost: pp.Cost(), Info: pp.ExplainInfo()}
-}
-
-func (op *physicalOptimizeOp) appendCandidate(logicalPlan *baseLogicalPlan, physicalPlan PhysicalPlan, prop string) *tracing.PhysicalPlanTrace {
+func (op *physicalOptimizeOp) appendCandidate(logicalPlan LogicalPlan, physicalPlan PhysicalPlan, prop *property.PhysicalProperty) *tracing.PlanTrace {
 	if op == nil || op.tracer == nil {
 		return nil
 	}
-	PhysicalPlanTrace := &tracing.PhysicalPlanTrace{TP: physicalPlan.TP(), ID: physicalPlan.ID()}
+	key := string(prop.HashCode())
+	PhysicalPlanTrace := &tracing.PlanTrace{TP: physicalPlan.TP(), ID: physicalPlan.ID(),
+		ExplainInfo: physicalPlan.ExplainInfo(), Cost: physicalPlan.Cost(), ProperType: prop.String()}
 	name := tracing.CodecPlanName(logicalPlan.TP(), logicalPlan.ID())
-	traceInfo := op.tracer.State[name][prop]
+	traceInfo := op.tracer.State[name][key]
 	if traceInfo == nil {
-		return nil
+		traceInfo = op.buildPhysicalOptimizeTraceInfo(logicalPlan, key)
 	}
 	traceInfo.Candidates = append(traceInfo.Candidates, PhysicalPlanTrace)
 	return PhysicalPlanTrace
@@ -426,9 +421,7 @@ func (p *baseLogicalPlan) findBestTask(prop *property.PhysicalProperty, planCoun
 
 	var cnt int64
 	var curTask task
-	if opt != nil {
-		opt.buildPhysicalOptimizeTraceInfo(p, newProp.String())
-	}
+	opt.buildPhysicalOptimizeTraceInfo(p, newProp.String())
 	if bestTask, cnt, err = p.enumeratePhysicalPlans4Task(plansFitsProp, newProp, false, planCounter, opt); err != nil {
 		return nil, 0, err
 	}
@@ -452,9 +445,6 @@ func (p *baseLogicalPlan) findBestTask(prop *property.PhysicalProperty, planCoun
 
 END:
 	p.storeTask(prop, bestTask)
-	if opt != nil {
-		opt.setBest(p.self, bestTask.plan(), prop.String())
-	}
 	return bestTask, cntPlan, nil
 }
 

planner/core/logical_plan_trace_test.go

@@ -23,61 +23,6 @@ import (
 	"github.com/pingcap/tidb/util/testleak"
 )
 
-func (s *testPlanSuite) TestLogicalOptimizeWithTraceEnabled(c *C) {
-	sql := "select * from t where a in (1,2)"
-	defer testleak.AfterTest(c)()
-	tt := []struct {
-		flags []uint64
-		steps int
-	}{
-		{
-			flags: []uint64{
-				flagEliminateAgg,
-				flagPushDownAgg},
-			steps: 2,
-		},
-		{
-			flags: []uint64{
-				flagEliminateAgg,
-				flagPushDownAgg,
-				flagPrunColumns,
-				flagBuildKeyInfo,
-			},
-			steps: 4,
-		},
-		{
-			flags: []uint64{},
-			steps: 0,
-		},
-	}
-
-	for i, tc := range tt {
-		comment := Commentf("case:%v sql:%s", i, sql)
-		stmt, err := s.ParseOneStmt(sql, "", "")
-		c.Assert(err, IsNil, comment)
-		err = Preprocess(s.ctx, stmt, WithPreprocessorReturn(&PreprocessorReturn{InfoSchema: s.is}))
-		c.Assert(err, IsNil, comment)
-		sctx := MockContext()
-		sctx.GetSessionVars().StmtCtx.EnableOptimizeTrace = true
-		builder, _ := NewPlanBuilder().Init(sctx, s.is, &hint.BlockHintProcessor{})
-		domain.GetDomain(sctx).MockInfoCacheAndLoadInfoSchema(s.is)
-		ctx := context.TODO()
-		p, err := builder.Build(ctx, stmt)
-		c.Assert(err, IsNil)
-		flag := uint64(0)
-		for _, f := range tc.flags {
-			flag = flag | f
-		}
-		p, err = logicalOptimize(ctx, flag, p.(LogicalPlan))
-		c.Assert(err, IsNil)
-		_, ok := p.(*LogicalProjection)
-		c.Assert(ok, IsTrue)
-		otrace := sctx.GetSessionVars().StmtCtx.LogicalOptimizeTrace
-		c.Assert(otrace, NotNil)
-		c.Assert(len(otrace.Steps), Equals, tc.steps)
-	}
-}
-
 func (s *testPlanSuite) TestSingleRuleTraceStep(c *C) {
 	defer testleak.AfterTest(c)()
 	tt := []struct {

planner/core/optimizer.go

@@ -108,7 +108,7 @@ func (op *logicalOptimizeOp) appendBeforeRuleOptimize(index int, name string, be
 	if op == nil || op.tracer == nil {
 		return
 	}
-	op.tracer.AppendRuleTracerBeforeRuleOptimize(index, name, before.buildLogicalPlanTrace())
+	op.tracer.AppendRuleTracerBeforeRuleOptimize(index, name, before.buildPlanTrace())
 }
 
 func (op *logicalOptimizeOp) appendStepToCurrent(id int, tp string, reason, action func() string) {
@@ -122,7 +122,7 @@ func (op *logicalOptimizeOp) recordFinalLogicalPlan(final LogicalPlan) {
 	if op == nil || op.tracer == nil {
 		return
 	}
-	op.tracer.RecordFinalLogicalPlan(final.buildLogicalPlanTrace())
+	op.tracer.RecordFinalLogicalPlan(final.buildPlanTrace())
 }
 
 // logicalOptRule means a logical optimizing rule, which contains decorrelate, ppd, column pruning, etc.
@@ -440,7 +440,7 @@ func isLogicalRuleDisabled(r logicalOptRule) bool {
 	return disabled
 }
 
-func physicalOptimize(logic LogicalPlan, planCounter *PlanCounterTp) (PhysicalPlan, float64, error) {
+func physicalOptimize(logic LogicalPlan, planCounter *PlanCounterTp) (plan PhysicalPlan, cost float64, err error) {
 	if _, err := logic.recursiveDeriveStats(nil); err != nil {
 		return nil, 0, err
 	}
@@ -458,7 +458,10 @@ func physicalOptimize(logic LogicalPlan, planCounter *PlanCounterTp) (PhysicalPl
 		tracer := &tracing.PhysicalOptimizeTracer{State: make(map[string]map[string]*tracing.PhysicalOptimizeTraceInfo)}
 		opt = opt.withEnableOptimizeTracer(tracer)
 		defer func() {
-			stmtCtx.PhysicalOptimizeTrace = tracer
+			if err == nil {
+				tracer.RecordFinalPlanTrace(plan.buildPlanTrace())
+				stmtCtx.PhysicalOptimizeTrace = tracer
+			}
 		}()
 	}
 

planner/core/physical_plan_trace_test.go

@@ -57,10 +57,9 @@ func TestPhysicalOptimizeWithTraceEnabled(t *testing.T) {
 	require.NoError(t, err)
 	otrace := sctx.GetSessionVars().StmtCtx.PhysicalOptimizeTrace
 	require.NotNil(t, otrace)
-	logicalList, physicalList, bests := getList(otrace)
+	logicalList, physicalList := getList(otrace)
 	require.True(t, checkList(logicalList, []string{"Projection_3", "Selection_2"}))
 	require.True(t, checkList(physicalList, []string{"Projection_4", "Selection_5"}))
-	require.True(t, checkList(bests, []string{"Projection_4", "Selection_5"}))
 }
 
 func checkList(d []string, s []string) bool {
@@ -75,18 +74,16 @@ func checkList(d []string, s []string) bool {
 	return true
 }
 
-func getList(otrace *tracing.PhysicalOptimizeTracer) (ll []string, pl []string, bests []string) {
+func getList(otrace *tracing.PhysicalOptimizeTracer) (ll []string, pl []string) {
 	for logicalPlan, v := range otrace.State {
 		ll = append(ll, logicalPlan)
 		for _, info := range v {
-			bests = append(bests, tracing.CodecPlanName(info.BestTask.TP, info.BestTask.ID))
 			for _, task := range info.Candidates {
 				pl = append(pl, tracing.CodecPlanName(task.TP, task.ID))
 			}
 		}
 	}
 	sort.Strings(ll)
 	sort.Strings(pl)
-	sort.Strings(bests)
-	return ll, pl, bests
+	return ll, pl
 }

planner/core/plan.go

@@ -66,6 +66,8 @@ type Plan interface {
 	SetOutputNames(names types.NameSlice)
 
 	SelectBlockOffset() int
+
+	buildPlanTrace() *tracing.PlanTrace
 }
 
 func enforceProperty(p *property.PhysicalProperty, tsk task, ctx sessionctx.Context) task {
@@ -306,9 +308,6 @@ type LogicalPlan interface {
 
 	// canPushToCop check if we might push this plan to a specific store.
 	canPushToCop(store kv.StoreType) bool
-
-	// buildLogicalPlanTrace clone necessary information from LogicalPlan
-	buildLogicalPlanTrace() *tracing.LogicalPlanTrace
 }
 
 // PhysicalPlan is a tree of the physical operators.
@@ -381,15 +380,6 @@ func (p *baseLogicalPlan) ExplainInfo() string {
 	return ""
 }
 
-// buildLogicalPlanTrace implements LogicalPlan
-func (p *baseLogicalPlan) buildLogicalPlanTrace() *tracing.LogicalPlanTrace {
-	planTrace := &tracing.LogicalPlanTrace{ID: p.ID(), TP: p.TP(), ExplainInfo: p.self.ExplainInfo()}
-	for _, child := range p.Children() {
-		planTrace.Children = append(planTrace.Children, child.buildLogicalPlanTrace())
-	}
-	return planTrace
-}
-
 type basePhysicalPlan struct {
 	basePlan
 
@@ -711,3 +701,27 @@ func (p *basePhysicalPlan) SetChild(i int, child PhysicalPlan) {
 func (p *basePlan) SCtx() sessionctx.Context {
 	return p.ctx
 }
+
+// buildPlanTrace implements Plan
+func (p *basePhysicalPlan) buildPlanTrace() *tracing.PlanTrace {
+	planTrace := &tracing.PlanTrace{ID: p.ID(), TP: p.TP(), ExplainInfo: p.self.ExplainInfo(), Cost: p.Cost()}
+	for _, child := range p.Children() {
+		planTrace.Children = append(planTrace.Children, child.buildPlanTrace())
+	}
+	return planTrace
+}
+
+// buildPlanTrace implements Plan
+func (p *baseLogicalPlan) buildPlanTrace() *tracing.PlanTrace {
+	planTrace := &tracing.PlanTrace{ID: p.ID(), TP: p.TP(), ExplainInfo: p.self.ExplainInfo()}
+	for _, child := range p.Children() {
+		planTrace.Children = append(planTrace.Children, child.buildPlanTrace())
+	}
+	return planTrace
+}
+
+// buildPlanTrace implements Plan
+func (p *basePlan) buildPlanTrace() *tracing.PlanTrace {
+	planTrace := &tracing.PlanTrace{ID: p.ID(), TP: p.TP()}
+	return planTrace
+}