planner: add CloneForPlanCache method to Plan interface (#54285)

ref #54057

Author: qw4990

pkg/planner/core/BUILD.bazel

@@ -230,6 +230,7 @@ go_test(
         "plan_cache_instance_test.go",
         "plan_cache_lru_test.go",
         "plan_cache_param_test.go",
+        "plan_cache_rebuild_test.go",
         "plan_cache_test.go",
         "plan_cacheable_checker_test.go",
         "plan_cost_detail_test.go",

pkg/planner/core/base/plan_base.go

@@ -79,6 +79,12 @@ type Plan interface {
 	QueryBlockOffset() int
 
 	BuildPlanTrace() *tracing.PlanTrace
+
+	// CloneForPlanCache clones this physical plan specially for instance level plan cache.
+	// A cached plan might be shared across multiple sessions, so for safety we have to clone it to make it thread-safe.
+	// Compared with the prior PhysicalPlan.Clone(), CloneForPlanCache() doesn't deep clones all fields instead it only
+	// deep clones fields that might be modified during reusing it and shallow clones all other fields for performance.
+	CloneForPlanCache() (cloned Plan, ok bool)
 }
 
 // PhysicalPlan is a tree of the physical operators.

pkg/planner/core/operator/baseimpl/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/expression",
         "//pkg/planner/context",
+        "//pkg/planner/core/base",
         "//pkg/planner/property",
         "//pkg/types",
         "//pkg/util/stringutil",

pkg/planner/core/operator/baseimpl/plan.go

@@ -21,6 +21,7 @@ import (
 
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/planner/context"
+	"github.com/pingcap/tidb/pkg/planner/core/base"
 	"github.com/pingcap/tidb/pkg/planner/property"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util/stringutil"
@@ -136,3 +137,8 @@ func (p *Plan) BuildPlanTrace() *tracing.PlanTrace {
 	planTrace := &tracing.PlanTrace{ID: p.ID(), TP: p.TP()}
 	return planTrace
 }
+
+// CloneForPlanCache clones this Plan for instance plan cache.
+func (*Plan) CloneForPlanCache() (cloned base.Plan, ok bool) {
+	return nil, false
+}

pkg/planner/core/physical_plans.go

@@ -280,6 +280,19 @@ func (p *PhysicalTableReader) Clone() (base.PhysicalPlan, error) {
 	return cloned, nil
 }
 
+// CloneForPlanCache implements Plan.CloneForPlanCache method.
+func (p *PhysicalTableReader) CloneForPlanCache() (base.Plan, bool) {
+	cloned := new(PhysicalTableReader)
+	*cloned = *p
+	t, ok := p.tablePlan.CloneForPlanCache()
+	if !ok {
+		return nil, false
+	}
+	cloned.tablePlan = t.(base.PhysicalPlan)
+	cloned.TablePlans = flattenPushDownPlan(cloned.tablePlan)
+	return cloned, true
+}
+
 // SetChildren overrides op.PhysicalPlan SetChildren interface.
 func (p *PhysicalTableReader) SetChildren(children ...base.PhysicalPlan) {
 	p.tablePlan = children[0]
@@ -342,6 +355,19 @@ func (p *PhysicalIndexReader) Clone() (base.PhysicalPlan, error) {
 	return cloned, err
 }
 
+// CloneForPlanCache implements Plan.CloneForPlanCache method.
+func (p *PhysicalIndexReader) CloneForPlanCache() (base.Plan, bool) {
+	cloned := new(PhysicalIndexReader)
+	*cloned = *p
+	t, ok := p.indexPlan.CloneForPlanCache()
+	if !ok {
+		return nil, false
+	}
+	cloned.indexPlan = t.(base.PhysicalPlan)
+	cloned.IndexPlans = flattenPushDownPlan(cloned.indexPlan)
+	return cloned, true
+}
+
 // SetSchema overrides op.PhysicalPlan SetSchema interface.
 func (p *PhysicalIndexReader) SetSchema(_ *expression.Schema) {
 	if p.indexPlan != nil {
@@ -768,6 +794,14 @@ func (p *PhysicalIndexScan) Clone() (base.PhysicalPlan, error) {
 	return cloned, nil
 }
 
+// CloneForPlanCache implements op.CloneForPlanCache interface.
+func (p *PhysicalIndexScan) CloneForPlanCache() (base.Plan, bool) {
+	cloned := new(PhysicalIndexScan)
+	*cloned = *p
+	cloned.Ranges = util.CloneRanges(p.Ranges)
+	return cloned, true
+}
+
 // ExtractCorrelatedCols implements op.PhysicalPlan interface.
 func (p *PhysicalIndexScan) ExtractCorrelatedCols() []*expression.CorrelatedColumn {
 	corCols := make([]*expression.CorrelatedColumn, 0, len(p.AccessCondition))
@@ -949,6 +983,14 @@ func (ts *PhysicalTableScan) Clone() (base.PhysicalPlan, error) {
 	return clonedScan, nil
 }
 
+// CloneForPlanCache implements op.CloneForPlanCache interface.
+func (ts *PhysicalTableScan) CloneForPlanCache() (base.Plan, bool) {
+	cloned := new(PhysicalTableScan)
+	*cloned = *ts
+	cloned.Ranges = util.CloneRanges(ts.Ranges)
+	return cloned, true
+}
+
 // ExtractCorrelatedCols implements op.PhysicalPlan interface.
 func (ts *PhysicalTableScan) ExtractCorrelatedCols() []*expression.CorrelatedColumn {
 	corCols := make([]*expression.CorrelatedColumn, 0, len(ts.AccessCondition)+len(ts.LateMaterializationFilterCondition))
@@ -2223,6 +2265,20 @@ func (p *PhysicalSelection) Clone() (base.PhysicalPlan, error) {
 	return cloned, nil
 }
 
+// CloneForPlanCache implements base.Plan.CloneForPlanCache method.
+func (p *PhysicalSelection) CloneForPlanCache() (base.Plan, bool) {
+	cloned := new(PhysicalSelection)
+	*cloned = *p
+	for i, c := range p.children {
+		clonedChild, ok := c.CloneForPlanCache()
+		if !ok {
+			return nil, false
+		}
+		cloned.children[i] = clonedChild.(base.PhysicalPlan)
+	}
+	return cloned, true
+}
+
 // ExtractCorrelatedCols implements op.PhysicalPlan interface.
 func (p *PhysicalSelection) ExtractCorrelatedCols() []*expression.CorrelatedColumn {
 	corCols := make([]*expression.CorrelatedColumn, 0, len(p.Conditions))

pkg/planner/core/plan_cache.go

@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/pkg/util/chunk"
 	contextutil "github.com/pingcap/tidb/pkg/util/context"
 	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
+	"github.com/pingcap/tidb/pkg/util/intest"
 	"github.com/pingcap/tidb/pkg/util/kvcache"
 )
 
@@ -48,6 +49,12 @@ type PlanCacheKeyTestIssue46760 struct{}
 // PlanCacheKeyTestIssue47133 is only for test.
 type PlanCacheKeyTestIssue47133 struct{}
 
+// PlanCacheKeyTestBeforeAdjust is only for test.
+type PlanCacheKeyTestBeforeAdjust struct{}
+
+// PlanCacheKeyTestAfterAdjust is only for test.
+type PlanCacheKeyTestAfterAdjust struct{}
+
 // SetParameterValuesIntoSCtx sets these parameters into session context.
 func SetParameterValuesIntoSCtx(sctx base.PlanContext, isNonPrep bool, markers []ast.ParamMarkerExpr, params []expression.Expression) error {
 	vars := sctx.GetSessionVars()
@@ -239,7 +246,13 @@ func GetPlanFromPlanCache(ctx context.Context, sctx sessionctx.Context,
 			cacheVal, hit = sctx.GetSessionPlanCache().Get(cacheKey, matchOpts)
 		}
 		if hit {
+			if intest.InTest && ctx.Value(PlanCacheKeyTestBeforeAdjust{}) != nil {
+				ctx.Value(PlanCacheKeyTestBeforeAdjust{}).(func(cachedVal *PlanCacheValue))(cacheVal.(*PlanCacheValue))
+			}
 			if plan, names, ok, err := adjustCachedPlan(sctx, cacheVal.(*PlanCacheValue), isNonPrepared, isPointPlan, cacheKey, bindSQL, is, stmt); err != nil || ok {
+				if intest.InTest && ctx.Value(PlanCacheKeyTestAfterAdjust{}) != nil {
+					ctx.Value(PlanCacheKeyTestAfterAdjust{}).(func(cachedVal *PlanCacheValue))(cacheVal.(*PlanCacheValue))
+				}
 				return plan, names, err
 			}
 		}

pkg/planner/core/plan_cache_rebuild_test.go

@@ -0,0 +1,88 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//	http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package core_test
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/planner/core"
+	"github.com/pingcap/tidb/pkg/planner/core/base"
+	"github.com/pingcap/tidb/pkg/testkit"
+	"github.com/stretchr/testify/require"
+)
+
+func TestPlanCacheClone(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec(`use test`)
+	tk.MustExec(`create table t (a int, b int, c int, primary key(a), key(b))`)
+
+	// TableScan
+	testCachedPlanClone(t, tk, `prepare st from 'select * from t where a<?'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	testCachedPlanClone(t, tk, `prepare st from 'select * from t where a>=?'`,
+		`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
+	//// IndexScan
+	//testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(b) where b<=?'`,
+	//	`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+	//testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(b) where b>?'`,
+	//	`set @a1=1, @a2=2`, `execute st using @a1`, `execute st using @a2`)
+
+	// TableScan + Selection
+	testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(primary) where a<? and b<?'`,
+		`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+	testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(primary) where a<? and b+?=10'`,
+		`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+
+	//// IndexScan + Selection
+	//testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(b) where a<? and b<?'`,
+	//	`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+	//testCachedPlanClone(t, tk, `prepare st from 'select * from t use index(b) where a<? and b+?=10'`,
+	//	`set @a1=1, @b1=1, @a2=2, @b2=2`, `execute st using @a1,@b1`, `execute st using @a2,@b2`)
+}
+
+func testCachedPlanClone(t *testing.T, tk *testkit.TestKit, prep, set, exec1, exec2 string) {
+	tk.MustExec(prep)
+	tk.MustExec(set)
+	tk.MustQuery(exec1) // generate the first cached plan
+
+	var original base.Plan
+	var originalFingerprint string
+	before := func(cachedVal *core.PlanCacheValue) {
+		// get the current cached plan and its fingerprint
+		original, originalFingerprint = cachedVal.Plan, planFingerprint(t, cachedVal.Plan)
+		// replace the cached plan with a cloned one
+		cloned, ok := original.CloneForPlanCache()
+		require.True(t, ok)
+		cachedVal.Plan = cloned
+	}
+	after := func(cachedVal *core.PlanCacheValue) {
+		cloned := cachedVal.Plan
+		require.True(t, originalFingerprint != planFingerprint(t, cloned))   // this cloned one have been adjusted by the optimizer
+		require.True(t, originalFingerprint == planFingerprint(t, original)) // the prior one should keep unchanged
+	}
+	ctx := context.WithValue(context.Background(), core.PlanCacheKeyTestBeforeAdjust{}, before)
+	ctx = context.WithValue(ctx, core.PlanCacheKeyTestAfterAdjust{}, after)
+	tk.MustQueryWithContext(ctx, exec2)
+}
+
+func planFingerprint(t *testing.T, p base.Plan) string {
+	v, err := json.Marshal(p)
+	require.NoError(t, err)
+	return string(v)
+}