*: add metadata lock when using the plan cache

pkg/ddl/tests/metadatalock/BUILD.bazel

@@ -8,10 +8,11 @@ go_test(
         "mdl_test.go",
     ],
     flaky = True,
-    shard_count = 34,
+    shard_count = 36,
     deps = [
         "//pkg/config",
         "//pkg/ddl",
+        "//pkg/ddl/ingest/testutil",
         "//pkg/errno",
         "//pkg/server",
         "//pkg/testkit",

pkg/ddl/tests/metadatalock/mdl_test.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/pingcap/failpoint"
+	ingesttestutil "github.com/pingcap/tidb/pkg/ddl/ingest/testutil"
 	mysql "github.com/pingcap/tidb/pkg/errno"
 	"github.com/pingcap/tidb/pkg/server"
 	"github.com/pingcap/tidb/pkg/testkit"
@@ -895,6 +896,103 @@ func TestMDLPreparePlanCacheInvalid(t *testing.T) {
 	tk.MustQuery(`execute stmt_test_1 using @a;`).Check(testkit.Rows("1 <nil>", "2 <nil>", "3 <nil>", "4 <nil>"))
 }
 
+func TestMDLPreparePlanCacheExecute(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	defer ingesttestutil.InjectMockBackendMgr(t, store)()
+
+	sv := server.CreateMockServer(t, store)
+
+	sv.SetDomain(dom)
+	dom.InfoSyncer().SetSessionManager(sv)
+	defer sv.Close()
+
+	conn1 := server.CreateMockConn(t, sv)
+	tk := testkit.NewTestKitWithSession(t, store, conn1.Context().Session)
+	conn2 := server.CreateMockConn(t, sv)
+	tkDDL := testkit.NewTestKitWithSession(t, store, conn2.Context().Session)
+	tk.MustExec("use test")
+	tk.MustExec("set global tidb_enable_metadata_lock=1")
+	tk.MustExec("create table t(a int);")
+	tk.MustExec("create table t2(a int);")
+	tk.MustExec("insert into t values(1), (2), (3), (4);")
+
+	tk.MustExec(`prepare stmt_test_1 from 'update t set a = ? where a = ?';`)
+	tk.MustExec(`set @a = 1, @b = 3;`)
+	tk.MustExec(`execute stmt_test_1 using @a, @b;`)
+
+	tk.MustExec("begin")
+
+	ch := make(chan struct{})
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		<-ch
+		tkDDL.MustExec("alter table test.t add index idx(a);")
+		wg.Done()
+	}()
+
+	tk.MustQuery("select * from t2")
+	tk.MustExec(`set @a = 2, @b=4;`)
+	tk.MustExec(`execute stmt_test_1 using @a, @b;`)
+	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("1"))
+	// The plan is from cache, the metadata lock should be added to block the DDL.
+	ch <- struct{}{}
+
+	time.Sleep(5 * time.Second)
+
+	tk.MustExec("commit")
+
+	wg.Wait()
+
+	tk.MustExec("admin check table t")
+}
+
+func TestMDLPreparePlanCacheExecute2(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	defer ingesttestutil.InjectMockBackendMgr(t, store)()
+
+	sv := server.CreateMockServer(t, store)
+
+	sv.SetDomain(dom)
+	dom.InfoSyncer().SetSessionManager(sv)
+	defer sv.Close()
+
+	conn1 := server.CreateMockConn(t, sv)
+	tk := testkit.NewTestKitWithSession(t, store, conn1.Context().Session)
+	conn2 := server.CreateMockConn(t, sv)
+	tkDDL := testkit.NewTestKitWithSession(t, store, conn2.Context().Session)
+	tk.MustExec("use test")
+	tk.MustExec("set global tidb_enable_metadata_lock=1")
+	tk.MustExec("create table t(a int);")
+	tk.MustExec("create table t2(a int);")
+	tk.MustExec("insert into t values(1), (2), (3), (4);")
+
+	tk.MustExec(`prepare stmt_test_1 from 'select * from t where a = ?';`)
+	tk.MustExec(`set @a = 1;`)
+	tk.MustExec(`execute stmt_test_1 using @a;`)
+
+	tk.MustExec("begin")
+	tk.MustQuery("select * from t2")
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		tkDDL.MustExec("alter table test.t add index idx(a);")
+		wg.Done()
+	}()
+
+	wg.Wait()
+
+	tk.MustExec(`set @a = 2;`)
+	tk.MustExec(`execute stmt_test_1 using @a;`)
+	// The plan should not be from cache because the schema has changed.
+	tk.MustQuery("select @@last_plan_from_cache;").Check(testkit.Rows("0"))
+	tk.MustExec("commit")
+
+	tk.MustExec("admin check table t")
+}
+
 func TestMDLDisable2Enable(t *testing.T) {
 	store, dom := testkit.CreateMockStoreAndDomain(t)
 	sv := server.CreateMockServer(t, store)

pkg/executor/prepared.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	plannercore "github.com/pingcap/tidb/pkg/planner/core"
 	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/sessiontxn"
 	"github.com/pingcap/tidb/pkg/types"
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/chunk"
@@ -119,7 +120,7 @@ func (e *PrepareExec) Next(ctx context.Context, _ *chunk.Chunk) error {
 			return err
 		}
 	}
-	stmt, p, paramCnt, err := plannercore.GeneratePlanCacheStmtWithAST(ctx, e.Ctx(), true, stmt0.Text(), stmt0, nil)
+	stmt, p, paramCnt, err := plannercore.GeneratePlanCacheStmtWithAST(ctx, e.Ctx(), true, stmt0.Text(), stmt0, sessiontxn.GetTxnManager(e.Ctx()).GetTxnInfoSchema())
 	if err != nil {
 		return err
 	}

pkg/meta/meta.go

@@ -942,6 +942,8 @@ func (m *Meta) UpdateTable(dbID int64, tableInfo *model.TableInfo) error {
 		return errors.Trace(err)
 	}
 
+	tableInfo.TableVersion++
+
 	data, err := json.Marshal(tableInfo)
 	if err != nil {
 		return errors.Trace(err)

pkg/parser/model/model.go

@@ -542,6 +542,8 @@ type TableInfo struct {
 
 	TTLInfo *TTLInfo `json:"ttl_info"`
 
+	TableVersion uint64 `json:"table_version"`
+
 	DBID int64 `json:"-"`
 }
 

pkg/planner/core/plan_cache.go

@@ -138,6 +138,7 @@ func planCachePreprocess(ctx context.Context, sctx sessionctx.Context, isNonPrep
 		stmtAst.CachedPlan = nil
 		vars.LastUpdateTime4PC = expiredTimeStamp4PC
 	}
+
 	return nil
 }
 
@@ -169,6 +170,22 @@ func GetPlanFromSessionPlanCache(ctx context.Context, sctx sessionctx.Context,
 		stmtCtx.ForceSetSkipPlanCache(errors.NewNoStackError(stmt.UncacheableReason))
 	}
 
+	if stmtCtx.UseCache {
+		// Acquire the metadata lock and update the schema version.
+		for i := 0; i < len(stmt.dbName); i++ {
+			prevVersion := stmt.tbls[i].Meta().TableVersion
+			tbl, err := tryLockMDLAndUpdateSchemaIfNecessary(sctx.GetPlanCtx(), stmt.dbName[i], stmt.tbls[i], is)
+			if err != nil {
+				return nil, nil, err
+			}
+			if tbl.Meta().TableVersion != prevVersion {
+				// Handle the case that the schema is updated for select statement.
+				sctx.GetSessionVars().StmtCtx.ForceSetSkipPlanCache(errors.NewNoStackError("schema changed before adding the metadata lock"))
+				break
+			}
+		}
+	}
+
 	var bindSQL string
 	if stmtCtx.UseCache {
 		var ignoreByBinding bool

pkg/planner/core/plan_cache_utils.go

@@ -36,6 +36,7 @@ import (
 	"github.com/pingcap/tidb/pkg/planner/util/fixcontrol"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessionctx/variable"
+	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/types"
 	driver "github.com/pingcap/tidb/pkg/types/parser_driver"
 	"github.com/pingcap/tidb/pkg/util/codec"
@@ -182,6 +183,19 @@ func GeneratePlanCacheStmtWithAST(ctx context.Context, sctx sessionctx.Context,
 	features := new(PlanCacheQueryFeatures)
 	paramStmt.Accept(features)
 
+	// Collect table information for metadata lock.
+	names := bindinfo.CollectTableNames(paramStmt)
+	dbName := make([]model.CIStr, 0, len(names))
+	tbls := make([]table.Table, 0, len(names))
+	for _, name := range names {
+		dbName = append(dbName, name.Schema)
+		tbl, err := is.TableByName(name.Schema, name.Name)
+		if err != nil {
+			return nil, nil, 0, err
+		}
+		tbls = append(tbls, tbl)
+	}
+
 	preparedObj := &PlanCacheStmt{
 		PreparedAst:         prepared,
 		StmtDB:              vars.CurrentDB,
@@ -194,6 +208,8 @@ func GeneratePlanCacheStmtWithAST(ctx context.Context, sctx sessionctx.Context,
 		StmtCacheable:       cacheable,
 		UncacheableReason:   reason,
 		QueryFeatures:       features,
+		dbName:              dbName,
+		tbls:                tbls,
 	}
 	if err = CheckPreparedPriv(sctx, preparedObj, ret.InfoSchema); err != nil {
 		return nil, nil, 0, err
@@ -471,6 +487,10 @@ type PlanCacheStmt struct {
 	//  NormalizedSQL4PC: select * from `test` . `t` where `a` > ? and `b` < ? --> schema name is added,
 	//  StmtText: select * from t where a>1 and b <? --> just format the original query;
 	StmtText string
+
+	// dbName and tbls are used to add metadata lock.
+	dbName []model.CIStr
+	tbls   []table.Table
 }
 
 // GetPreparedStmt extract the prepared statement from the execute statement.