*: Introduce PessimisticRRTxnContextProvider for pessimistic repeatable read txn

planner/optimize.go

@@ -33,15 +33,13 @@ import (
 	"github.com/pingcap/tidb/metrics"
 	"github.com/pingcap/tidb/parser"
 	"github.com/pingcap/tidb/parser/ast"
-	"github.com/pingcap/tidb/parser/model"
 	"github.com/pingcap/tidb/planner/cascades"
 	"github.com/pingcap/tidb/planner/core"
 	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/privilege"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
-	"github.com/pingcap/tidb/sessiontxn"
 	"github.com/pingcap/tidb/types"
 	"github.com/pingcap/tidb/util/hint"
 	"github.com/pingcap/tidb/util/logutil"
@@ -123,17 +121,9 @@ func Optimize(ctx context.Context, sctx sessionctx.Context, node ast.Node, is in
 			fp = plannercore.TryFastPlan(sctx, node)
 		}
 		if fp != nil {
-			if !useMaxTS(sctx, fp) {
-				if err := sessiontxn.WarmUpTxn(sctx); err != nil {
-					return nil, nil, err
-				}
-			}
 			return fp, fp.OutputNames(), nil
 		}
 	}
-	if err := sessiontxn.WarmUpTxn(sctx); err != nil {
-		return nil, nil, err
-	}
 
 	useBinding := sessVars.UsePlanBaselines
 	stmtNode, ok := node.(ast.StmtNode)
@@ -486,29 +476,6 @@ func handleEvolveTasks(ctx context.Context, sctx sessionctx.Context, br *bindinf
 	globalHandle.AddEvolvePlanTask(br.OriginalSQL, br.Db, binding)
 }
 
-// useMaxTS returns true when meets following conditions:
-//  1. ctx is auto commit tagged.
-//  2. plan is point get by pk.
-//  3. not a cache table.
-func useMaxTS(ctx sessionctx.Context, p plannercore.Plan) bool {
-	if !plannercore.IsAutoCommitTxn(ctx) {
-		return false
-	}
-	v, ok := p.(*plannercore.PointGetPlan)
-	if !ok {
-		return false
-	}
-	noSecondRead := v.IndexInfo == nil || (v.IndexInfo.Primary && v.TblInfo.IsCommonHandle)
-	if !noSecondRead {
-		return false
-	}
-
-	if v.TblInfo != nil && (v.TblInfo.TableCacheStatusType != model.TableCacheStatusDisable) {
-		return false
-	}
-	return true
-}
-
 // OptimizeExecStmt to optimize prepare statement protocol "execute" statement
 // this is a short path ONLY does things filling prepare related params
 // for point select like plan which does not need extra things

session/session.go

@@ -1911,6 +1911,9 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 		}
 	})
 
+	if err := sessiontxn.WarmUpTxn(s); err != nil {
+		return nil, err
+	}
 	// Transform abstract syntax tree to a physical plan(stored in executor.ExecStmt).
 	compiler := executor.Compiler{Ctx: s}
 	stmt, err := compiler.Compile(ctx, stmtNode)
@@ -1924,6 +1927,11 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 		}
 		return nil, err
 	}
+
+	if err = sessiontxn.OptimizeWithPlan(s, stmt.Plan); err != nil {
+		return nil, err
+	}
+
 	durCompile := time.Since(s.sessionVars.StartTime)
 	s.GetSessionVars().DurationCompile = durCompile
 	if s.isInternal() {

session/txnmanager.go

@@ -114,9 +114,9 @@ func (m *txnManager) OnStmtRetry(ctx context.Context) error {
 	return m.ctxProvider.OnStmtRetry(ctx)
 }
 
-func (m *txnManager) Advise(tp sessiontxn.AdviceType) error {
+func (m *txnManager) Advise(tp sessiontxn.AdviceType, val ...any) error {
 	if m.ctxProvider != nil {
-		return m.ctxProvider.Advise(tp)
+		return m.ctxProvider.Advise(tp, val)
 	}
 	return nil
 }
@@ -139,6 +139,12 @@ func (m *txnManager) newProviderWithRequest(r *sessiontxn.EnterNewTxnRequest) se
 		switch m.sctx.GetSessionVars().IsolationLevelForNewTxn() {
 		case ast.ReadCommitted:
 			return isolation.NewPessimisticRCTxnContextProvider(m.sctx, r.CausalConsistencyOnly)
+		case ast.Serializable:
+			// todo: Add pessimistic serializable transaction context provider
+			break
+		default:
+			// We use Repeatable read for all other cases.
+			return isolation.NewPessimisticRRTxnContextProvider(m.sctx, r.CausalConsistencyOnly)
 		}
 	}
 

sessiontxn/interface.go

@@ -100,6 +100,8 @@ const (
 	// AdviceWarmUp indicates to warm up the provider's inner state.
 	// For example, the provider can prefetch tso when it is advised.
 	AdviceWarmUp AdviceType = iota
+	// AdviceOptimizeWithPlan indicates to do some optimizations with the given plan
+	AdviceOptimizeWithPlan
 )
 
 // TxnContextProvider provides txn context
@@ -120,7 +122,7 @@ type TxnContextProvider interface {
 	// OnStmtRetry is the hook that should be called when a statement is retried internally.
 	OnStmtRetry(ctx context.Context) error
 	// Advise is used to give advice to provider
-	Advise(tp AdviceType) error
+	Advise(tp AdviceType, val []any) error
 }
 
 // TxnManager is an interface providing txn context management in session
@@ -147,7 +149,7 @@ type TxnManager interface {
 	OnStmtRetry(ctx context.Context) error
 	// Advise is used to give advice to provider.
 	// For example, `AdviceWarmUp` can tell the provider to warm up its inner state.
-	Advise(tp AdviceType) error
+	Advise(tp AdviceType, val ...any) error
 }
 
 // NewTxn starts a new optimistic and active txn, it can be used for the below scenes:
@@ -176,5 +178,10 @@ func WarmUpTxn(sctx sessionctx.Context) error {
 	return GetTxnManager(sctx).Advise(AdviceWarmUp)
 }
 
+// OptimizeWithPlan gives the `AdviceOptimizeWithPlan` advise to the provider
+func OptimizeWithPlan(sctx sessionctx.Context, plan any) error {
+	return GetTxnManager(sctx).Advise(AdviceOptimizeWithPlan, plan)
+}
+
 // GetTxnManager returns the TxnManager object from session context
 var GetTxnManager func(sctx sessionctx.Context) TxnManager

sessiontxn/isolation/base.go

@@ -25,6 +25,7 @@ import (
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/sessiontxn"
+	"github.com/pingcap/tidb/sessiontxn/staleread"
 	"github.com/pingcap/tidb/table/temptable"
 )
 
@@ -118,7 +119,7 @@ func (p *baseTxnContextProvider) GetStmtForUpdateTS() (uint64, error) {
 	return p.getStmtForUpdateTSFunc()
 }
 
-func (p *baseTxnContextProvider) Advise(tp sessiontxn.AdviceType) error {
+func (p *baseTxnContextProvider) Advise(tp sessiontxn.AdviceType, _ []any) error {
 	switch tp {
 	case sessiontxn.AdviceWarmUp:
 		return p.warmUp()
@@ -197,8 +198,16 @@ func (p *baseTxnContextProvider) isTidbSnapshotEnabled() bool {
 	return p.sctx.GetSessionVars().SnapshotTS != 0
 }
 
+// isBeginStmtWithStaleRead indicates whether the current statement is `BeginStmt` type with stale read
+// Because stale read will use `staleread.StalenessTxnContextProvider` for query, so if `staleread.IsStmtStaleness()`
+// returns true in other providers, it means the current statement is `BeginStmt` with stale read
+func (p *baseTxnContextProvider) isBeginStmtWithStaleRead() bool {
+	return staleread.IsStmtStaleness(p.sctx)
+}
+
 func (p *baseTxnContextProvider) warmUp() error {
-	if p.isTidbSnapshotEnabled() {
+	if p.isBeginStmtWithStaleRead() {
+		// When executing `START TRANSACTION READ ONLY AS OF ...` no need to warmUp
 		return nil
 	}
 	return p.prepareTxn()

sessiontxn/isolation/readcommitted.go

@@ -108,12 +108,12 @@ func (p *PessimisticRCTxnContextProvider) OnStmtRetry(ctx context.Context) error
 }
 
 // Advise is used to give advice to provider
-func (p *PessimisticRCTxnContextProvider) Advise(tp sessiontxn.AdviceType) error {
+func (p *PessimisticRCTxnContextProvider) Advise(tp sessiontxn.AdviceType, val []any) error {
 	switch tp {
 	case sessiontxn.AdviceWarmUp:
 		return p.warmUp()
 	default:
-		return p.baseTxnContextProvider.Advise(tp)
+		return p.baseTxnContextProvider.Advise(tp, val)
 	}
 }
 

sessiontxn/isolation/readcommitted_test.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/executor"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
@@ -274,9 +275,22 @@ func TestRCProviderInitialize(t *testing.T) {
 	assertAfterActive.Check(t)
 	require.Equal(t, ts, se.GetSessionVars().TxnCtx.StartTS)
 	tk.MustExec("rollback")
+
+	// Case Pessimistic Autocommit
+	config.GetGlobalConfig().PessimisticTxn.PessimisticAutoCommit.Store(true)
+	assert = inActiveRCTxnAssert(se)
+	assertAfterActive = activeRCTxnAssert(t, se, true)
+	require.NoError(t, se.PrepareTxnCtx(context.TODO()))
+	provider = assert.CheckAndGetProvider(t)
+	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	ts, err = provider.GetStmtReadTS()
+	require.NoError(t, err)
+	assertAfterActive.Check(t)
+	require.Equal(t, ts, se.GetSessionVars().TxnCtx.StartTS)
+	tk.MustExec("rollback")
 }
 
-func TestTidbSnapshotVar(t *testing.T) {
+func TestTidbSnapshotVarInRC(t *testing.T) {
 	store, dom, clean := testkit.CreateMockStoreAndDomain(t)
 	defer clean()