infoschema_v2: remove schemaTableInfos call of DeadTableLockChecker

br/pkg/task/BUILD.bazel

@@ -55,6 +55,7 @@ go_library(
         "//pkg/ddl",
         "//pkg/domain",
         "//pkg/infoschema",
+        "//pkg/infoschema/context",
         "//pkg/kv",
         "//pkg/meta/model",
         "//pkg/parser/model",

br/pkg/task/restore_data.go

@@ -22,7 +22,7 @@ import (
 	"github.com/pingcap/tidb/br/pkg/summary"
 	"github.com/pingcap/tidb/br/pkg/utils"
 	tidbconfig "github.com/pingcap/tidb/pkg/config"
-	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/kv"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
@@ -181,7 +181,7 @@ func resetTiFlashReplicas(ctx context.Context, g glue.Glue, storage kv.Storage,
 
 	expectTiFlashStoreCount := uint64(0)
 	needTiFlash := false
-	tableInfoRes := info.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	tableInfoRes := info.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, s := range tableInfoRes {
 		for _, t := range s.TableInfos {
 			if t.TiFlashReplica != nil {

pkg/ddl/ddl_tiflash_api.go

@@ -34,6 +34,7 @@ import (
 	ddlutil "github.com/pingcap/tidb/pkg/ddl/util"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/table"
@@ -451,7 +452,7 @@ func (d *ddl) refreshTiFlashTicker(ctx sessionctx.Context, pollTiFlashContext *T
 	var tableList = make([]TiFlashReplicaStatus, 0)
 
 	// Collect TiFlash Replica info, for every table.
-	ch := schema.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	ch := schema.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, v := range ch {
 		for _, tblInfo := range v.TableInfos {
 			LoadTiFlashReplicaInfo(tblInfo, &tableList)

pkg/ddl/executor.go

@@ -39,6 +39,7 @@ import (
 	"github.com/pingcap/tidb/pkg/errctx"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
@@ -421,7 +422,7 @@ func (e *executor) getPendingTiFlashTableCount(originVersion int64, pendingCount
 		return originVersion, pendingCount
 	}
 	cnt := uint32(0)
-	dbs := is.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	dbs := is.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, db := range dbs {
 		if util.IsMemOrSysDB(db.DBName.L) {
 			continue

pkg/ddl/placement_policy.go

@@ -23,6 +23,7 @@ import (
 	"github.com/pingcap/tidb/pkg/ddl/placement"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/parser/ast"
@@ -377,7 +378,7 @@ func CheckPlacementPolicyNotInUseFromInfoSchema(is infoschema.InfoSchema, policy
 		}
 	}
 
-	schemaTables := is.ListTablesWithSpecialAttribute(infoschema.AllPlacementPolicyAttribute)
+	schemaTables := is.ListTablesWithSpecialAttribute(infoschemacontext.AllPlacementPolicyAttribute)
 	for _, schemaTable := range schemaTables {
 		for _, tblInfo := range schemaTable.TableInfos {
 			if err := checkPlacementPolicyNotUsedByTable(tblInfo, policy); err != nil {

pkg/ddl/util/dead_table_lock_checker.go

@@ -83,24 +83,19 @@ func (d *DeadTableLockChecker) GetDeadLockedTables(ctx context.Context, is infos
 		return nil, err
 	}
 	deadLockTables := make(map[model.SessionInfo][]model.TableLockTpInfo)
-	for _, schema := range is.AllSchemas() {
-		select {
-		case <-ctx.Done():
-			return nil, ctx.Err()
-		default:
-		}
-		tblInfos, err := is.SchemaTableInfos(ctx, schema.Name)
-		if err != nil {
-			return nil, errors.Trace(err)
-		}
-		for _, tbl := range tblInfos {
+
+	tbls := is.ListTablesWithSpecialAttribute(func(t *model.TableInfo) bool {
+		return t.Lock != nil
+	})
+	for _, db := range tbls {
+		for _, tbl := range db.TableInfos {
 			if tbl.Lock == nil {
 				continue
 			}
 			for _, se := range tbl.Lock.Sessions {
 				if _, ok := aliveServers[se.ServerID]; !ok {
 					deadLockTables[se] = append(deadLockTables[se], model.TableLockTpInfo{
-						SchemaID: schema.ID,
+						SchemaID: tbl.DBID,
 						TableID:  tbl.ID,
 						Tp:       tbl.Lock.Tp,
 					})

pkg/executor/infoschema_reader.go

@@ -40,6 +40,7 @@ import (
 	"github.com/pingcap/tidb/pkg/executor/internal/pdhelper"
 	"github.com/pingcap/tidb/pkg/expression"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/keyspace"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/meta/autoid"
@@ -2656,7 +2657,7 @@ func (e *memtableRetriever) dataForTableTiFlashReplica(_ context.Context, sctx s
 		rows          [][]types.Datum
 		tiFlashStores map[int64]pd.StoreInfo
 	)
-	rs := e.is.ListTablesWithSpecialAttribute(infoschema.TiFlashAttribute)
+	rs := e.is.ListTablesWithSpecialAttribute(infoschemacontext.TiFlashAttribute)
 	for _, schema := range rs {
 		for _, tbl := range schema.TableInfos {
 			if checker != nil && !checker.RequestVerification(sctx.GetSessionVars().ActiveRoles, schema.DBName.L, tbl.Name.L, "", mysql.AllPrivMask) {

pkg/executor/show_placement.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/tidb/pkg/ddl/placement"
 	"github.com/pingcap/tidb/pkg/domain/infosync"
 	"github.com/pingcap/tidb/pkg/infoschema"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta/model"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
@@ -361,7 +362,7 @@ func (e *ShowExec) fetchAllTablePlacements(ctx context.Context, scheduleState ma
 	dbs := e.is.AllSchemaNames()
 	slices.SortFunc(dbs, func(i, j pmodel.CIStr) int { return cmp.Compare(i.O, j.O) })
 
-	tbls := e.is.ListTablesWithSpecialAttribute(infoschema.AllSpecialAttribute)
+	tbls := e.is.ListTablesWithSpecialAttribute(infoschemacontext.AllSpecialAttribute)
 	for _, db := range tbls {
 		tableRowSets := make([]tableRowSet, 0)
 		for _, tblInfo := range db.TableInfos {

pkg/infoschema/BUILD.bazel

@@ -96,6 +96,7 @@ go_test(
     deps = [
         "//pkg/ddl/placement",
         "//pkg/domain",
+        "//pkg/infoschema/context",
         "//pkg/infoschema/internal",
         "//pkg/kv",
         "//pkg/meta",

pkg/infoschema/builder_misc.go

@@ -18,6 +18,7 @@ import (
 	"fmt"
 
 	"github.com/pingcap/errors"
+	infoschemacontext "github.com/pingcap/tidb/pkg/infoschema/context"
 	"github.com/pingcap/tidb/pkg/meta"
 	"github.com/pingcap/tidb/pkg/meta/model"
 )
@@ -115,7 +116,7 @@ func (b *Builder) initMisc(dbInfos []*model.DBInfo, policies []*model.PolicyInfo
 
 	// Maintain foreign key reference information.
 	if b.enableV2 {
-		rs := b.ListTablesWithSpecialAttribute(ForeignKeysAttribute)
+		rs := b.ListTablesWithSpecialAttribute(infoschemacontext.ForeignKeysAttribute)
 		for _, db := range rs {
 			for _, tbl := range db.TableInfos {
 				info.addReferredForeignKeys(db.DBName, tbl)

pkg/infoschema/context/infoschema.go

@@ -22,6 +22,79 @@ import (
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
 )
 
+// SpecialAttributeFilter is used to filter tables with special attributes.
+type SpecialAttributeFilter func(*model.TableInfo) bool
+
+// TTLAttribute is the TTL attribute filter used by ListTablesWithSpecialAttribute.
+var TTLAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.State == model.StatePublic && t.TTLInfo != nil
+}
+
+// TiFlashAttribute is the TiFlashReplica attribute filter used by ListTablesWithSpecialAttribute.
+var TiFlashAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.TiFlashReplica != nil
+}
+
+// PlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
+var PlacementPolicyAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	if t.PlacementPolicyRef != nil {
+		return true
+	}
+	if parInfo := t.GetPartitionInfo(); parInfo != nil {
+		for _, def := range parInfo.Definitions {
+			if def.PlacementPolicyRef != nil {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// AllPlacementPolicyAttribute is the Placement Policy attribute filter used by ListTablesWithSpecialAttribute.
+// Different from PlacementPolicyAttribute, Partition.Enable flag will be ignored.
+var AllPlacementPolicyAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	if t.PlacementPolicyRef != nil {
+		return true
+	}
+	if t.Partition != nil {
+		for _, def := range t.Partition.Definitions {
+			if def.PlacementPolicyRef != nil {
+				return true
+			}
+		}
+	}
+	return false
+}
+
+// TableLockAttribute is the Table Lock attribute filter used by ListTablesWithSpecialAttribute.
+var TableLockAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.Lock != nil
+}
+
+// ForeignKeysAttribute is the ForeignKeys attribute filter used by ListTablesWithSpecialAttribute.
+var ForeignKeysAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return len(t.ForeignKeys) > 0
+}
+
+// PartitionAttribute is the Partition attribute filter used by ListTablesWithSpecialAttribute.
+var PartitionAttribute SpecialAttributeFilter = func(t *model.TableInfo) bool {
+	return t.GetPartitionInfo() != nil
+}
+
+// HasSpecialAttributes checks if a table has any special attributes.
+func HasSpecialAttributes(t *model.TableInfo) bool {
+	return TTLAttribute(t) || TiFlashAttribute(t) || PlacementPolicyAttribute(t) || PartitionAttribute(t) || TableLockAttribute(t) || ForeignKeysAttribute(t)
+}
+
+// AllSpecialAttribute marks a model.TableInfo with any special attributes.
+var AllSpecialAttribute SpecialAttributeFilter = HasSpecialAttributes
+
+// TableInfoResult is used to store the result of ListTablesWithSpecialAttribute.
+type TableInfoResult struct {
+	DBName     pmodel.CIStr
+	TableInfos []*model.TableInfo
+}
+
 // MetaOnlyInfoSchema is a workaround.
 // Due to circular dependency cannot return the complete interface.
 // But MetaOnlyInfoSchema is widely used for scenes that require meta only, so we give a convenience for that.
@@ -37,6 +110,7 @@ type MetaOnlyInfoSchema interface {
 	SchemaAndTable
 	AllSchemaNames() []pmodel.CIStr
 	SchemaSimpleTableInfos(ctx stdctx.Context, schema pmodel.CIStr) ([]*model.TableNameInfo, error)
+	ListTablesWithSpecialAttribute(filter SpecialAttributeFilter) []TableInfoResult
 	Misc
 }
 

pkg/infoschema/infoschema.go

@@ -364,15 +364,10 @@ func (is *infoSchema) SchemaSimpleTableInfos(ctx stdctx.Context, schema pmodel.C
 	return ret, nil
 }
 
-type tableInfoResult struct {
-	DBName     pmodel.CIStr
-	TableInfos []*model.TableInfo
-}
-
-func (is *infoSchema) ListTablesWithSpecialAttribute(filter specialAttributeFilter) []tableInfoResult {
-	ret := make([]tableInfoResult, 0, 10)
+func (is *infoSchema) ListTablesWithSpecialAttribute(filter context.SpecialAttributeFilter) []context.TableInfoResult {
+	ret := make([]context.TableInfoResult, 0, 10)
 	for _, dbName := range is.AllSchemaNames() {
-		res := tableInfoResult{DBName: dbName}
+		res := context.TableInfoResult{DBName: dbName}
 		tblInfos, err := is.SchemaTableInfos(stdctx.Background(), dbName)
 		terror.Log(err)
 		for _, tblInfo := range tblInfos {