txn: manage the initialization of RCCheckTS by transaction context provider

executor/adapter.go

@@ -222,6 +222,10 @@ type ExecStmt struct {
 	Ti          *TelemetryInfo
 }
 
+func (a ExecStmt) GetStmtNode() ast.StmtNode {
+	return a.StmtNode
+}
+
 // PointGet short path for point exec directly from plan, keep only necessary steps
 func (a *ExecStmt) PointGet(ctx context.Context, is infoschema.InfoSchema) (*recordSet, error) {
 	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {

executor/executor.go

@@ -1904,11 +1904,6 @@ func ResetContextOfStmt(ctx sessionctx.Context, s ast.StmtNode) (err error) {
 			sc.NotFillCache = !opts.SQLCache
 		}
 		sc.WeakConsistency = isWeakConsistencyRead(ctx, stmt)
-		// Try to mark the `RCCheckTS` flag for the first time execution of in-transaction read requests
-		// using read-consistency isolation level.
-		if NeedSetRCCheckTSFlag(ctx, stmt) {
-			sc.RCCheckTS = true
-		}
 	case *ast.SetOprStmt:
 		sc.InSelectStmt = true
 		sc.OverflowAsWarning = true
@@ -2042,13 +2037,3 @@ func isWeakConsistencyRead(ctx sessionctx.Context, node ast.Node) bool {
 	return sessionVars.ConnectionID > 0 && sessionVars.ReadConsistency.IsWeak() &&
 		plannercore.IsAutoCommitTxn(ctx) && plannercore.IsReadOnly(node, sessionVars)
 }
-
-// NeedSetRCCheckTSFlag checks whether it's needed to set `RCCheckTS` flag in current stmtctx.
-func NeedSetRCCheckTSFlag(ctx sessionctx.Context, node ast.Node) bool {
-	sessionVars := ctx.GetSessionVars()
-	if sessionVars.ConnectionID > 0 && sessionVars.RcReadCheckTS && sessionVars.InTxn() &&
-		sessionVars.IsPessimisticReadConsistency() && !sessionVars.RetryInfo.Retrying && plannercore.IsReadOnly(node, sessionVars) {
-		return true
-	}
-	return false
-}

executor/executor_test.go

@@ -3470,7 +3470,7 @@ func TestUnreasonablyClose(t *testing.T) {
 		err = sessiontxn.NewTxn(context.Background(), tk.Session())
 		require.NoError(t, err, comment)
 
-		err = sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO())
+		err = sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO(), stmt)
 		require.NoError(t, err, comment)
 
 		executorBuilder := executor.NewMockExecutorBuilderForTest(tk.Session(), is, nil, oracle.GlobalTxnScope)

planner/core/preprocess.go

@@ -1653,7 +1653,7 @@ func (p *preprocessor) updateStateFromStaleReadProcessor() error {
 				if err := txnManager.EnterNewTxn(context.TODO(), newTxnRequest); err != nil {
 					return err
 				}
-				if err := txnManager.OnStmtStart(context.TODO()); err != nil {
+				if err := txnManager.OnStmtStart(context.TODO(), nil); err != nil {
 					return err
 				}
 			}

planner/funcdep/extract_fd_test.go

@@ -214,7 +214,7 @@ func TestFDSet_ExtractFD(t *testing.T) {
 	for i, tt := range tests {
 		comment := fmt.Sprintf("case:%v sql:%s", i, tt.sql)
 		require.NoError(t, tk.Session().PrepareTxnCtx(context.TODO()))
-		require.NoError(t, sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO()))
+		require.NoError(t, sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO(), nil))
 		stmt, err := par.ParseOneStmt(tt.sql, "", "")
 		require.NoError(t, err, comment)
 		tk.Session().GetSessionVars().PlanID = 0
@@ -312,7 +312,7 @@ func TestFDSet_ExtractFDForApply(t *testing.T) {
 	is := testGetIS(t, tk.Session())
 	for i, tt := range tests {
 		require.NoError(t, tk.Session().PrepareTxnCtx(context.TODO()))
-		require.NoError(t, sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO()))
+		require.NoError(t, sessiontxn.GetTxnManager(tk.Session()).OnStmtStart(context.TODO(), nil))
 		comment := fmt.Sprintf("case:%v sql:%s", i, tt.sql)
 		stmt, err := par.ParseOneStmt(tt.sql, "", "")
 		require.NoError(t, err, comment)

session/session.go

@@ -1098,7 +1098,7 @@ func (s *session) retry(ctx context.Context, maxCnt uint) (err error) {
 			}
 			_, digest := s.sessionVars.StmtCtx.SQLDigest()
 			s.txn.onStmtStart(digest.String())
-			if err = sessiontxn.GetTxnManager(s).OnStmtStart(ctx); err == nil {
+			if err = sessiontxn.GetTxnManager(s).OnStmtStart(ctx, st.GetStmtNode()); err == nil {
 				_, err = st.Exec(ctx)
 			}
 			s.txn.onStmtEnd()
@@ -1904,7 +1904,7 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 	s.txn.onStmtStart(digest.String())
 	defer s.txn.onStmtEnd()
 
-	if err := s.onTxnManagerStmtStartOrRetry(ctx); err != nil {
+	if err := s.onTxnManagerStmtStartOrRetry(ctx, stmtNode); err != nil {
 		return nil, err
 	}
 
@@ -1967,11 +1967,11 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 	return recordSet, nil
 }
 
-func (s *session) onTxnManagerStmtStartOrRetry(ctx context.Context) error {
+func (s *session) onTxnManagerStmtStartOrRetry(ctx context.Context, node ast.StmtNode) error {
 	if s.sessionVars.RetryInfo.Retrying {
 		return sessiontxn.GetTxnManager(s).OnStmtRetry(ctx)
 	}
-	return sessiontxn.GetTxnManager(s).OnStmtStart(ctx)
+	return sessiontxn.GetTxnManager(s).OnStmtStart(ctx, node)
 }
 
 func (s *session) validateStatementReadOnlyInStaleness(stmtNode ast.StmtNode) error {
@@ -2185,7 +2185,7 @@ func (s *session) PrepareStmt(sql string) (stmtID uint32, paramCount int, fields
 		return
 	}
 
-	if err = s.onTxnManagerStmtStartOrRetry(ctx); err != nil {
+	if err = s.onTxnManagerStmtStartOrRetry(ctx, nil); err != nil {
 		return
 	}
 
@@ -2406,7 +2406,8 @@ func (s *session) ExecutePreparedStmt(ctx context.Context, stmtID uint32, args [
 	s.txn.onStmtStart(preparedStmt.SQLDigest.String())
 	defer s.txn.onStmtEnd()
 
-	if err = s.onTxnManagerStmtStartOrRetry(ctx); err != nil {
+	execAst := &ast.ExecuteStmt{ExecID: stmtID, BinaryArgs: args}
+	if err = s.onTxnManagerStmtStartOrRetry(ctx, execAst); err != nil {
 		return nil, err
 	}
 

session/txnmanager.go

@@ -46,6 +46,7 @@ func getTxnManager(sctx sessionctx.Context) sessiontxn.TxnManager {
 type txnManager struct {
 	sctx sessionctx.Context
 
+	stmtNode    ast.StmtNode
 	ctxProvider sessiontxn.TxnContextProvider
 }
 
@@ -116,14 +117,18 @@ func (m *txnManager) EnterNewTxn(ctx context.Context, r *sessiontxn.EnterNewTxnR
 
 func (m *txnManager) OnTxnEnd() {
 	m.ctxProvider = nil
+	m.stmtNode = nil
 }
 
 // OnStmtStart is the hook that should be called when a new statement started
-func (m *txnManager) OnStmtStart(ctx context.Context) error {
+func (m *txnManager) OnStmtStart(ctx context.Context, node ast.StmtNode) error {
+	if node != nil {
+		m.stmtNode = node
+	}
 	if m.ctxProvider == nil {
 		return errors.New("context provider not set")
 	}
-	return m.ctxProvider.OnStmtStart(ctx)
+	return m.ctxProvider.OnStmtStart(ctx, m.stmtNode)
 }
 
 // OnStmtErrorForNextAction is the hook that should be called when a new statement get an error

sessiontxn/interface.go

@@ -124,7 +124,7 @@ type TxnContextProvider interface {
 	// OnInitialize is the hook that should be called when enter a new txn with this provider
 	OnInitialize(ctx context.Context, enterNewTxnType EnterNewTxnType) error
 	// OnStmtStart is the hook that should be called when a new statement started
-	OnStmtStart(ctx context.Context) error
+	OnStmtStart(ctx context.Context, node ast.StmtNode) error
 	// OnStmtErrorForNextAction is the hook that should be called when a new statement get an error
 	OnStmtErrorForNextAction(point StmtErrorHandlePoint, err error) (StmtErrorAction, error)
 	// OnStmtRetry is the hook that should be called when a statement is retried internally.
@@ -150,7 +150,7 @@ type TxnManager interface {
 	// OnTxnEnd is the hook that should be called after transaction commit or rollback
 	OnTxnEnd()
 	// OnStmtStart is the hook that should be called when a new statement started
-	OnStmtStart(ctx context.Context) error
+	OnStmtStart(ctx context.Context, node ast.StmtNode) error
 	// OnStmtErrorForNextAction is the hook that should be called when a new statement get an error
 	// This method is not required to be called for every error in the statement,
 	// it is only required to be called for some errors handled in some specified points given by the parameter `point`.
@@ -178,7 +178,7 @@ func NewTxnInStmt(ctx context.Context, sctx sessionctx.Context) error {
 	if err := NewTxn(ctx, sctx); err != nil {
 		return err
 	}
-	return GetTxnManager(sctx).OnStmtStart(ctx)
+	return GetTxnManager(sctx).OnStmtStart(ctx, nil)
 }
 
 // GetTxnManager returns the TxnManager object from session context

sessiontxn/isolation/base.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/infoschema"
 	"github.com/pingcap/tidb/kv"
+	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/sessiontxn"
@@ -138,7 +139,7 @@ func (p *baseTxnContextProvider) GetStmtForUpdateTS() (uint64, error) {
 	return p.getStmtForUpdateTSFunc()
 }
 
-func (p *baseTxnContextProvider) OnStmtStart(ctx context.Context) error {
+func (p *baseTxnContextProvider) OnStmtStart(ctx context.Context, _ ast.StmtNode) error {
 	p.ctx = ctx
 	return nil
 }

sessiontxn/isolation/optimistic_test.go

@@ -49,7 +49,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 	se := tk.Session()
 	compareTS := getOracleTS(t, se)
 	provider := initializeOptimisticProvider(t, tk, true)
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 	readTS, err := provider.GetStmtReadTS()
 	require.NoError(t, err)
 	updateTS, err := provider.GetStmtForUpdateTS()
@@ -59,7 +59,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 	compareTS = readTS
 
 	// for optimistic mode ts, ts should be the same for all statements
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 	readTS, err = provider.GetStmtReadTS()
 	require.NoError(t, err)
 	updateTS, err = provider.GetStmtForUpdateTS()
@@ -72,7 +72,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 	require.NoError(t, err)
 	stmt := stmts[0]
 	provider = initializeOptimisticProvider(t, tk, false)
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), stmt))
 	plan, _, err := planner.Optimize(context.TODO(), tk.Session(), stmt, provider.GetTxnInfoSchema())
 	require.NoError(t, err)
 	require.NoError(t, provider.AdviseOptimizeWithPlan(plan))
@@ -85,7 +85,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 
 	// if the oracle future is prepared fist, `math.MaxUint64` should still be used after plan
 	provider = initializeOptimisticProvider(t, tk, false)
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), stmt))
 	require.NoError(t, provider.AdviseWarmup())
 	plan, _, err = planner.Optimize(context.TODO(), tk.Session(), stmt, provider.GetTxnInfoSchema())
 	require.NoError(t, err)
@@ -100,7 +100,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 	// when it is in explicit txn, we should not use `math.MaxUint64`
 	compareTS = getOracleTS(t, se)
 	provider = initializeOptimisticProvider(t, tk, true)
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), stmt))
 	plan, _, err = planner.Optimize(context.TODO(), tk.Session(), stmt, provider.GetTxnInfoSchema())
 	require.NoError(t, err)
 	require.NoError(t, provider.AdviseOptimizeWithPlan(plan))
@@ -115,7 +115,7 @@ func TestOptimisticTxnContextProviderTS(t *testing.T) {
 	tk.MustExec("set @@autocommit=0")
 	compareTS = getOracleTS(t, se)
 	provider = initializeOptimisticProvider(t, tk, false)
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), stmt))
 	plan, _, err = planner.Optimize(context.TODO(), tk.Session(), stmt, provider.GetTxnInfoSchema())
 	require.NoError(t, err)
 	require.NoError(t, provider.AdviseOptimizeWithPlan(plan))
@@ -175,15 +175,15 @@ func TestOptimisticHandleError(t *testing.T) {
 	}
 
 	for _, c := range cases {
-		require.NoError(t, provider.OnStmtStart(context.TODO()))
+		require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 		action, err := provider.OnStmtErrorForNextAction(c.point, c.err)
 		if c.point == sessiontxn.StmtErrAfterPessimisticLock {
 			require.Error(t, err)
 			require.Same(t, c.err, err)
 			require.Equal(t, sessiontxn.StmtActionError, action)
 
 			// next statement should not update ts
-			require.NoError(t, provider.OnStmtStart(context.TODO()))
+			require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 			checkTS()
 		} else {
 			require.NoError(t, err)
@@ -194,13 +194,13 @@ func TestOptimisticHandleError(t *testing.T) {
 			checkTS()
 
 			// OnStmtErrorForNextAction again
-			require.NoError(t, provider.OnStmtStart(context.TODO()))
+			require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 			action, err = provider.OnStmtErrorForNextAction(c.point, c.err)
 			require.NoError(t, err)
 			require.Equal(t, sessiontxn.StmtActionNoIdea, action)
 
 			// next statement should not update ts
-			require.NoError(t, provider.OnStmtStart(context.TODO()))
+			require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 			checkTS()
 		}
 	}
@@ -280,7 +280,7 @@ func TestOptimisticProviderInitialize(t *testing.T) {
 			assertAfterActive.couldRetry = c.autocommit || !c.disableTxnAutoRetry
 			require.NoError(t, se.PrepareTxnCtx(context.TODO()))
 			provider := assert.CheckAndGetProvider(t)
-			require.NoError(t, provider.OnStmtStart(context.TODO()))
+			require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 			ts, err := provider.GetStmtReadTS()
 			require.NoError(t, err)
 			assertAfterActive.Check(t)
@@ -343,12 +343,12 @@ func TestTidbSnapshotVarInOptimisticTxn(t *testing.T) {
 	}
 
 	// information schema and ts should equal to snapshot when tidb_snapshot is set
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 	checkUseSnapshot()
 
 	// information schema and ts will restore when set tidb_snapshot to empty
 	tk.MustExec("set @@tidb_snapshot=''")
-	require.NoError(t, provider.OnStmtStart(context.TODO()))
+	require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 	checkUseTxn()
 
 	// txn will not be active after `GetStmtReadTS` or `GetStmtForUpdateTS` when `tidb_snapshot` is set
@@ -368,7 +368,7 @@ func TestTidbSnapshotVarInOptimisticTxn(t *testing.T) {
 			assertAfterUseSnapshot := activeSnapshotTxnAssert(se, se.GetSessionVars().SnapshotTS, "")
 			require.NoError(t, se.PrepareTxnCtx(context.TODO()))
 			provider = assert.CheckAndGetProvider(t)
-			require.NoError(t, provider.OnStmtStart(context.TODO()))
+			require.NoError(t, provider.OnStmtStart(context.TODO(), nil))
 			checkUseSnapshot()
 			assertAfterUseSnapshot.Check(t)
 		}()

sessiontxn/isolation/readcommitted.go

@@ -21,6 +21,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/parser/ast"
 	"github.com/pingcap/tidb/parser/terror"
+	plannercore "github.com/pingcap/tidb/planner/core"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
 	"github.com/pingcap/tidb/sessiontxn"
@@ -77,13 +78,30 @@ func NewPessimisticRCTxnContextProvider(sctx sessionctx.Context, causalConsisten
 }
 
 // OnStmtStart is the hook that should be called when a new statement started
-func (p *PessimisticRCTxnContextProvider) OnStmtStart(ctx context.Context) error {
-	if err := p.baseTxnContextProvider.OnStmtStart(ctx); err != nil {
+func (p *PessimisticRCTxnContextProvider) OnStmtStart(ctx context.Context, node ast.StmtNode) error {
+	if err := p.baseTxnContextProvider.OnStmtStart(ctx, node); err != nil {
 		return err
 	}
+
+	// Try to mark the `RCCheckTS` flag for the first time execution of in-transaction read requests
+	// using read-consistency isolation level.
+	if node != nil && NeedSetRCCheckTSFlag(p.sctx, node) {
+		p.sctx.GetSessionVars().StmtCtx.RCCheckTS = true
+	}
+
 	return p.prepareStmt(!p.isTxnPrepared)
 }
 
+// NeedSetRCCheckTSFlag checks whether it's needed to set `RCCheckTS` flag in current stmtctx.
+func NeedSetRCCheckTSFlag(ctx sessionctx.Context, node ast.Node) bool {
+	sessionVars := ctx.GetSessionVars()
+	if sessionVars.ConnectionID > 0 && sessionVars.RcReadCheckTS && sessionVars.InTxn() &&
+		!sessionVars.RetryInfo.Retrying && plannercore.IsReadOnly(node, sessionVars) {
+		return true
+	}
+	return false
+}
+
 // OnStmtErrorForNextAction is the hook that should be called when a new statement get an error
 func (p *PessimisticRCTxnContextProvider) OnStmtErrorForNextAction(point sessiontxn.StmtErrorHandlePoint, err error) (sessiontxn.StmtErrorAction, error) {
 	// Invalid rc check for next statement or retry when error occurs