planner, session: add isolation read with engine type (#12997)

kv/kv.go

@@ -237,6 +237,14 @@ const (
 	TiFlash
 )
 
+// Name returns the name of store type.
+func (t StoreType) Name() string {
+	if t == TiFlash {
+		return "tiflash"
+	}
+	return "tikv"
+}
+
 // Request represents a kv request.
 type Request struct {
 	// Tp is the request type.

planner/core/common_plans.go

@@ -734,9 +734,9 @@ func (e *Explain) explainPlanInRowFormat(p Plan, taskType, indent string, isLast
 		var storeType string
 		switch x.StoreType {
 		case kv.TiKV:
-			storeType = "tikv"
+			storeType = kv.TiKV.Name()
 		case kv.TiFlash:
-			storeType = "tiflash"
+			storeType = kv.TiFlash.Name()
 		default:
 			err = errors.Errorf("the store type %v is unknown", x.StoreType)
 			return

planner/core/integration_test.go

@@ -238,3 +238,20 @@ func (s *testIntegrationSuite) TestSimplifyOuterJoinWithCast(c *C) {
 		tk.MustQuery(tt).Check(testkit.Rows(output[i].Plan...))
 	}
 }
+
+func (s *testIntegrationSuite) TestNoneAccessPathsFoundByIsolationRead(c *C) {
+	tk := testkit.NewTestKit(c, s.store)
+
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(a int primary key)")
+
+	_, err := tk.Exec("select * from t")
+	c.Assert(err, IsNil)
+
+	tk.MustExec("set @@session.tidb_isolation_read_engines = 'tiflash'")
+
+	_, err = tk.Exec("select * from t")
+	c.Assert(err, NotNil)
+	c.Assert(err.Error(), Equals, "[planner:1815]Internal : Can not find access path matching 'tidb_isolation_read_engines'(value: 'tiflash'). Available values are 'tikv'.")
+}

planner/core/logical_plan_builder.go

@@ -2438,6 +2438,10 @@ func (b *PlanBuilder) buildDataSource(ctx context.Context, tn *ast.TableName, as
 	if err != nil {
 		return nil, err
 	}
+	possiblePaths, err = b.filterPathByIsolationRead(possiblePaths)
+	if err != nil {
+		return nil, err
+	}
 
 	var columns []*table.Column
 	if b.inUpdateStmt {

planner/core/planbuilder.go

@@ -35,6 +35,7 @@ import (
 	"github.com/pingcap/tidb/planner/property"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
+	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/statistics"
 	"github.com/pingcap/tidb/store/tikv"
 	"github.com/pingcap/tidb/table"
@@ -662,6 +663,32 @@ func (b *PlanBuilder) getPossibleAccessPaths(indexHints []*ast.IndexHint, tblInf
 	return available, nil
 }
 
+func (b *PlanBuilder) filterPathByIsolationRead(paths []*accessPath) ([]*accessPath, error) {
+	// TODO: filter paths with isolation read locations.
+	isolationReadEngines := b.ctx.GetSessionVars().GetIsolationReadEngines()
+	availableEngine := map[kv.StoreType]struct{}{}
+	var availableEngineStr string
+	for i := len(paths) - 1; i >= 0; i-- {
+		if _, ok := availableEngine[paths[i].storeType]; !ok {
+			availableEngine[paths[i].storeType] = struct{}{}
+			if availableEngineStr != "" {
+				availableEngineStr += ", "
+			}
+			availableEngineStr += paths[i].storeType.Name()
+		}
+		if _, ok := isolationReadEngines[paths[i].storeType]; !ok {
+			paths = append(paths[:i], paths[i+1:]...)
+		}
+	}
+	var err error
+	if len(paths) == 0 {
+		engineVals, _ := b.ctx.GetSessionVars().GetSystemVar(variable.TiDBIsolationReadEngines)
+		err = ErrInternal.GenWithStackByArgs(fmt.Sprintf("Can not find access path matching '%v'(value: '%v'). Available values are '%v'.",
+			variable.TiDBIsolationReadEngines, engineVals, availableEngineStr))
+	}
+	return paths, err
+}
+
 func removeIgnoredPaths(paths, ignoredPaths []*accessPath, tblInfo *model.TableInfo) []*accessPath {
 	if len(ignoredPaths) == 0 {
 		return paths

session/session.go

@@ -1838,6 +1838,7 @@ var builtinGlobalVariable = []string{
 	variable.TiDBMaxDeltaSchemaCount,
 	variable.TiDBCapturePlanBaseline,
 	variable.TiDBUsePlanBaselines,
+	variable.TiDBIsolationReadEngines,
 }
 
 var (

session/session_test.go

@@ -2815,6 +2815,21 @@ func (s *testSessionSuite) TestReplicaRead(c *C) {
 	c.Assert(tk.Se.GetSessionVars().GetReplicaRead(), Equals, kv.ReplicaReadLeader)
 }
 
+func (s *testSessionSuite) TestIsolationRead(c *C) {
+	var err error
+	tk := testkit.NewTestKit(c, s.store)
+	tk.Se, err = session.CreateSession4Test(s.store)
+	c.Assert(err, IsNil)
+	c.Assert(len(tk.Se.GetSessionVars().GetIsolationReadEngines()), Equals, 2)
+	tk.MustExec("set @@tidb_isolation_read_engines = 'tiflash';")
+	engines := tk.Se.GetSessionVars().GetIsolationReadEngines()
+	c.Assert(len(engines), Equals, 1)
+	_, hasTiFlash := engines[kv.TiFlash]
+	_, hasTiKV := engines[kv.TiKV]
+	c.Assert(hasTiFlash, Equals, true)
+	c.Assert(hasTiKV, Equals, false)
+}
+
 func (s *testSessionSuite) TestStmtHints(c *C) {
 	var err error
 	tk := testkit.NewTestKit(c, s.store)

sessionctx/variable/session.go

@@ -443,6 +443,9 @@ type SessionVars struct {
 	// replicaRead is used for reading data from replicas, only follower is supported at this time.
 	replicaRead kv.ReplicaReadType
 
+	// isolationReadEngines is used to isolation read, tidb only read from the stores whose engine type is in the engines.
+	isolationReadEngines map[kv.StoreType]struct{}
+
 	PlannerSelectBlockAsName []ast.HintTable
 }
 
@@ -520,6 +523,7 @@ func NewSessionVars() *SessionVars {
 		replicaRead:                 kv.ReplicaReadLeader,
 		AllowRemoveAutoInc:          DefTiDBAllowRemoveAutoInc,
 		UsePlanBaselines:            DefTiDBUsePlanBaselines,
+		isolationReadEngines:        map[kv.StoreType]struct{}{kv.TiKV: {}, kv.TiFlash: {}},
 	}
 	vars.Concurrency = Concurrency{
 		IndexLookupConcurrency:     DefIndexLookupConcurrency,
@@ -616,6 +620,11 @@ func (s *SessionVars) GetSplitRegionTimeout() time.Duration {
 	return time.Duration(s.WaitSplitRegionTimeout) * time.Second
 }
 
+// GetIsolationReadEngines gets isolation read engines.
+func (s *SessionVars) GetIsolationReadEngines() map[kv.StoreType]struct{} {
+	return s.isolationReadEngines
+}
+
 // CleanBuffers cleans the temporary bufs
 func (s *SessionVars) CleanBuffers() {
 	s.GetWriteStmtBufs().clean()
@@ -956,6 +965,16 @@ func (s *SessionVars) SetSystemVar(name string, val string) error {
 		SetMaxDeltaSchemaCount(tidbOptInt64(val, DefTiDBMaxDeltaSchemaCount))
 	case TiDBUsePlanBaselines:
 		s.UsePlanBaselines = TiDBOptOn(val)
+	case TiDBIsolationReadEngines:
+		s.isolationReadEngines = make(map[kv.StoreType]struct{})
+		for _, engine := range strings.Split(val, ",") {
+			switch engine {
+			case kv.TiKV.Name():
+				s.isolationReadEngines[kv.TiKV] = struct{}{}
+			case kv.TiFlash.Name():
+				s.isolationReadEngines[kv.TiFlash] = struct{}{}
+			}
+		}
 	}
 	s.systems[name] = val
 	return nil

sessionctx/variable/sysvar.go

@@ -730,6 +730,7 @@ var defaultSysVars = []*SysVar{
 	{ScopeGlobal | ScopeSession, TiDBEnableStmtSummary, "0"},
 	{ScopeGlobal | ScopeSession, TiDBCapturePlanBaseline, "0"},
 	{ScopeGlobal | ScopeSession, TiDBUsePlanBaselines, BoolToIntStr(DefTiDBUsePlanBaselines)},
+	{ScopeGlobal | ScopeSession, TiDBIsolationReadEngines, "tikv,tiflash"},
 }
 
 // SynonymsSysVariables is synonyms of system variables.

sessionctx/variable/tidb_vars.go

@@ -333,6 +333,10 @@ const (
 
 	// TiDBUsePlanBaselines indicates whether the use of plan baselines is enabled.
 	TiDBUsePlanBaselines = "tidb_use_plan_baselines"
+
+	// TiDBIsolationReadEngines indicates the tidb only read from the stores whose engine type is involved in IsolationReadEngines.
+	// Now, only support TiKV and TiFlash.
+	TiDBIsolationReadEngines = "tidb_isolation_read_engines"
 )
 
 // Default TiDB system variable values.

sessionctx/variable/varsutil.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/parser/ast"
 	"github.com/pingcap/parser/mysql"
 	"github.com/pingcap/tidb/config"
+	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/timeutil"
 )
@@ -625,6 +626,24 @@ func ValidateSetSystemVar(vars *SessionVars, name string, value string) (string,
 			return "", nil
 		}
 		return value, ErrWrongValueForVar.GenWithStackByArgs(name, value)
+	case TiDBIsolationReadEngines:
+		engines := strings.Split(value, ",")
+		var formatVal string
+		for i, engine := range engines {
+			engine = strings.TrimSpace(engine)
+			if i != 0 {
+				formatVal += ","
+			}
+			switch {
+			case strings.EqualFold(engine, kv.TiKV.Name()):
+				formatVal += kv.TiKV.Name()
+			case strings.EqualFold(engine, kv.TiFlash.Name()):
+				formatVal += kv.TiFlash.Name()
+			default:
+				return value, ErrWrongValueForVar.GenWithStackByArgs(name, value)
+			}
+		}
+		return formatVal, nil
 	}
 	return value, nil
 }

sessionctx/variable/varsutil_test.go

@@ -459,6 +459,10 @@ func (s *testVarsutilSuite) TestValidate(c *C) {
 		{TiDBTxnMode, "pessimistic", false},
 		{TiDBTxnMode, "optimistic", false},
 		{TiDBTxnMode, "", false},
+		{TiDBIsolationReadEngines, "", true},
+		{TiDBIsolationReadEngines, "tikv", false},
+		{TiDBIsolationReadEngines, "TiKV,tiflash", false},
+		{TiDBIsolationReadEngines, "   tikv,   tiflash  ", false},
 	}
 
 	for _, t := range tests {

store/tikv/region_cache.go

@@ -1220,7 +1220,7 @@ func (s *Store) initResolve(bo *Backoffer, c *RegionCache) (addr string, err err
 		s.storeType = kv.TiKV
 		for _, label := range store.Labels {
 			if label.Key == "engine" {
-				if label.Value == "tiflash" {
+				if label.Value == kv.TiFlash.Name() {
 					s.storeType = kv.TiFlash
 				}
 				break
@@ -1265,7 +1265,7 @@ func (s *Store) reResolve(c *RegionCache) {
 	storeType := kv.TiKV
 	for _, label := range store.Labels {
 		if label.Key == "engine" {
-			if label.Value == "tiflash" {
+			if label.Value == kv.TiFlash.Name() {
 				storeType = kv.TiFlash
 			}
 			break