infoschema: add WithRefillOption for TableByName and TableByID (#55511)

ref #50959

pkg/ddl/foreign_key.go

@@ -161,8 +161,9 @@ func checkTableForeignKeysValid(sctx sessionctx.Context, is infoschema.InfoSchem
 	}
 
 	referredFKInfos := is.GetTableReferredForeignKeys(schema, tbInfo.Name.L)
+	ctx := infoschema.WithRefillOption(context.Background(), false)
 	for _, referredFK := range referredFKInfos {
-		childTable, err := is.TableByName(context.Background(), referredFK.ChildSchema, referredFK.ChildTable)
+		childTable, err := is.TableByName(ctx, referredFK.ChildSchema, referredFK.ChildTable)
 		if err != nil {
 			return err
 		}

pkg/domain/plan_replayer_dump.go

@@ -440,16 +440,17 @@ func dumpMeta(zw *zip.Writer) error {
 	return nil
 }
 
-func dumpTiFlashReplica(ctx sessionctx.Context, zw *zip.Writer, pairs map[tableNamePair]struct{}) error {
+func dumpTiFlashReplica(sctx sessionctx.Context, zw *zip.Writer, pairs map[tableNamePair]struct{}) error {
 	bf, err := zw.Create(PlanReplayerTiFlashReplicasFile)
 	if err != nil {
 		return errors.AddStack(err)
 	}
-	is := GetDomain(ctx).InfoSchema()
+	is := GetDomain(sctx).InfoSchema()
+	ctx := infoschema.WithRefillOption(context.Background(), false)
 	for pair := range pairs {
 		dbName := model.NewCIStr(pair.DBName)
 		tableName := model.NewCIStr(pair.TableName)
-		t, err := is.TableByName(context.Background(), dbName, tableName)
+		t, err := is.TableByName(ctx, dbName, tableName)
 		if err != nil {
 			logutil.BgLogger().Warn("failed to find table info", zap.Error(err),
 				zap.String("dbName", dbName.L), zap.String("tableName", tableName.L))
@@ -496,11 +497,12 @@ func dumpSchemaMeta(zw *zip.Writer, tables map[tableNamePair]struct{}) error {
 func dumpStatsMemStatus(zw *zip.Writer, pairs map[tableNamePair]struct{}, do *Domain) error {
 	statsHandle := do.StatsHandle()
 	is := do.InfoSchema()
+	ctx := infoschema.WithRefillOption(context.Background(), false)
 	for pair := range pairs {
 		if pair.IsView {
 			continue
 		}
-		tbl, err := is.TableByName(context.Background(), model.NewCIStr(pair.DBName), model.NewCIStr(pair.TableName))
+		tbl, err := is.TableByName(ctx, model.NewCIStr(pair.DBName), model.NewCIStr(pair.TableName))
 		if err != nil {
 			return err
 		}

pkg/infoschema/BUILD.bazel

@@ -91,7 +91,7 @@ go_test(
     ],
     embed = [":infoschema"],
     flaky = True,
-    shard_count = 21,
+    shard_count = 22,
     deps = [
         "//pkg/ddl/placement",
         "//pkg/domain",

pkg/infoschema/infoschema_nokit_test.go

@@ -21,6 +21,13 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// GetCache is exported for testing.
+func (is *infoschemaV2) HasCache(tableID int64, schemaVersion int64) bool {
+	key := tableCacheKey{tableID, schemaVersion}
+	_, found := is.tableCache.Get(key)
+	return found
+}
+
 func TestInfoSchemaAddDel(t *testing.T) {
 	is := newInfoSchema()
 	is.addSchema(&schemaTables{

pkg/infoschema/infoschema_test.go

@@ -17,6 +17,7 @@ package infoschema_test
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"math"
 	"strings"
 	"testing"
@@ -598,6 +599,55 @@ func TestBuildBundle(t *testing.T) {
 	assertBundle(is2, p1.ID, p1Bundle)
 }
 
+func TestWithRefillOption(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("set @@global.tidb_schema_cache_size = 512 * 1024 * 1024")
+
+	tk.MustExec("create table t1 (id int)")
+	tk.MustQuery("select * from t1").Check(testkit.Rows())
+	is := dom.InfoSchema()
+	tbl, err := is.TableByName(context.Background(), model.NewCIStr("test"), model.NewCIStr("t1"))
+	require.NoError(t, err)
+	tblInfo := tbl.Meta()
+	ok, v2 := infoschema.IsV2(is)
+	require.True(t, ok)
+
+	hit := true
+	miss := false
+	testCases := []struct {
+		OP     string
+		ctx    context.Context
+		expect bool
+	}{
+		{"TableByName", context.Background(), hit},
+		{"TableByName", infoschema.WithRefillOption(context.Background(), true), hit},
+		{"TableByName", infoschema.WithRefillOption(context.Background(), false), miss},
+		{"TableByID", context.Background(), miss},
+		{"TableByID", infoschema.WithRefillOption(context.Background(), true), hit},
+		{"TableByID", infoschema.WithRefillOption(context.Background(), false), miss},
+	}
+
+	for i, testCase := range testCases {
+		// Mock t1 schema cache been evicted.
+		v2.EvictTable(model.NewCIStr("test"), model.NewCIStr("t1"))
+
+		// Test the API
+		switch testCase.OP {
+		case "TableByID":
+			_, found := is.TableByID(testCase.ctx, tblInfo.ID)
+			require.True(t, found)
+		case "TableByName":
+			_, err := is.TableByName(testCase.ctx, model.NewCIStr("test"), model.NewCIStr("t1"))
+			require.NoError(t, err)
+		}
+
+		got := v2.HasCache(tblInfo.ID, is.SchemaMetaVersion())
+		require.Equal(t, testCase.expect, got, fmt.Sprintf("case %d failed", i))
+	}
+}
+
 func TestLocalTemporaryTables(t *testing.T) {
 	re := internal.CreateAutoIDRequirement(t)
 	var err error

pkg/infoschema/infoschema_v2.go

@@ -588,11 +588,10 @@ func (is *infoschemaV2) CloneAndUpdateTS(startTS uint64) *infoschemaV2 {
 	return &tmp
 }
 
+// TableByID implements the InfoSchema interface.
+// As opposed to TableByName, TableByID will not refill cache when schema cache miss,
+// unless the caller changes the behavior by passing a context use WithRefillOption.
 func (is *infoschemaV2) TableByID(ctx context.Context, id int64) (val table.Table, ok bool) {
-	return is.tableByID(ctx, id, true)
-}
-
-func (is *infoschemaV2) tableByID(ctx context.Context, id int64, noRefill bool) (val table.Table, ok bool) {
 	if !tableIDIsValid(id) {
 		return
 	}
@@ -618,11 +617,17 @@ func (is *infoschemaV2) tableByID(ctx context.Context, id int64, noRefill bool)
 		}
 		return nil, false
 	}
+
+	refill := false
+	if opt := ctx.Value(refillOptionKey); opt != nil {
+		refill = opt.(bool)
+	}
+
 	// get cache with old key
 	oldKey := tableCacheKey{itm.tableID, itm.schemaVersion}
 	tbl, found = is.tableCache.Get(oldKey)
 	if found && tbl != nil {
-		if !noRefill {
+		if refill {
 			is.tableCache.Set(key, tbl)
 		}
 		return tbl, true
@@ -634,7 +639,7 @@ func (is *infoschemaV2) tableByID(ctx context.Context, id int64, noRefill bool)
 		return nil, false
 	}
 
-	if !noRefill {
+	if refill {
 		is.tableCache.Set(oldKey, ret)
 	}
 	return ret, true
@@ -680,6 +685,8 @@ func (h *tableByNameHelper) onItem(item tableItem) bool {
 	return true
 }
 
+// TableByName implements the InfoSchema interface.
+// When schema cache miss, it will fetch the TableInfo from TikV and refill cache.
 func (is *infoschemaV2) TableByName(ctx context.Context, schema, tbl model.CIStr) (t table.Table, err error) {
 	if IsSpecialDB(schema.L) {
 		if raw, ok := is.specials.Load(schema.L); ok {
@@ -716,7 +723,15 @@ func (is *infoschemaV2) TableByName(ctx context.Context, schema, tbl model.CIStr
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
-	is.tableCache.Set(oldKey, ret)
+
+	refill := true
+	if opt := ctx.Value(refillOptionKey); opt != nil {
+		refill = opt.(bool)
+	}
+	if refill {
+		is.tableCache.Set(oldKey, ret)
+	}
+
 	metrics.TableByNameMissDuration.Observe(float64(time.Since(start)))
 	return ret, nil
 }
@@ -1034,8 +1049,7 @@ func (is *infoschemaV2) loadTableInfo(ctx context.Context, tblID, dbID int64, ts
 				goto retry
 			}
 
-			// TODO load table panic!!!
-			panic(err)
+			return nil, errors.Trace(err)
 		}
 
 		// table removed.
@@ -1456,3 +1470,14 @@ func (is *infoschemaV2) ListTablesWithSpecialAttribute(filter specialAttributeFi
 	}
 	return ret
 }
+
+type refillOption struct{}
+
+var refillOptionKey refillOption
+
+// WithRefillOption controls the infoschema v2 cache refill operation.
+// By default, TableByID does not refill schema cache, and TableByName does.
+// The behavior can be changed by providing the context.Context.
+func WithRefillOption(ctx context.Context, refill bool) context.Context {
+	return context.WithValue(ctx, refillOptionKey, refill)
+}

pkg/planner/core/memtable_infoschema_extractor.go

@@ -617,6 +617,7 @@ func findTableAndSchemaByName(
 		table  *model.TableInfo
 	}
 	tableMap := make(map[int64]schemaAndTable, len(tableNames))
+	ctx = infoschema.WithRefillOption(ctx, false)
 	for _, n := range tableNames {
 		for _, s := range schemas {
 			tbl, err := is.TableByName(ctx, s, n)