txn: txnManager providers GetReadSnapshot and GetForUpdateSnapshot.

executor/adapter.go

@@ -259,7 +259,7 @@ func (a *ExecStmt) PointGet(ctx context.Context, is infoschema.InfoSchema) (*rec
 		} else {
 			// CachedPlan type is already checked in last step
 			pointGetPlan := a.PsStmt.PreparedAst.CachedPlan.(*plannercore.PointGetPlan)
-			exec.Init(pointGetPlan, startTs)
+			exec.Init(pointGetPlan, startTs, nil)
 			a.PsStmt.Executor = exec
 		}
 	}

executor/batch_point_get.go

@@ -78,9 +78,8 @@ type BatchPointGetExec struct {
 	// virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns.
 	virtualColumnRetFieldTypes []*types.FieldType
 
-	snapshot   kv.Snapshot
-	stats      *runtimeStatsWithSnapshot
-	cacheTable kv.MemBuffer
+	snapshot kv.Snapshot
+	stats    *runtimeStatsWithSnapshot
 }
 
 // buildVirtualColumnInfo saves virtual column indices and sort them in definition order
@@ -104,35 +103,21 @@ func (e *BatchPointGetExec) Open(context.Context) error {
 		return err
 	}
 	e.txn = txn
-	var snapshot kv.Snapshot
-	if txn.Valid() && txnCtx.StartTS == txnCtx.GetForUpdateTS() && txnCtx.StartTS == e.snapshotTS {
-		// We can safely reuse the transaction snapshot if snapshotTS is equal to forUpdateTS.
-		// The snapshot may contain cache that can reduce RPC call.
-		snapshot = txn.GetSnapshot()
-	} else {
-		snapshot = e.ctx.GetSnapshotWithTS(e.snapshotTS)
-	}
-	if e.ctx.GetSessionVars().StmtCtx.RCCheckTS {
-		snapshot.SetOption(kv.IsolationLevel, kv.RCCheckTS)
-	}
-	if e.cacheTable != nil {
-		snapshot = cacheTableSnapshot{snapshot, e.cacheTable}
-	}
 	if e.runtimeStats != nil {
 		snapshotStats := &txnsnapshot.SnapshotRuntimeStats{}
 		e.stats = &runtimeStatsWithSnapshot{
 			SnapshotRuntimeStats: snapshotStats,
 		}
-		snapshot.SetOption(kv.CollectRuntimeStats, snapshotStats)
+		e.snapshot.SetOption(kv.CollectRuntimeStats, snapshotStats)
 		stmtCtx.RuntimeStatsColl.RegisterStats(e.id, e.stats)
 	}
 	replicaReadType := e.ctx.GetSessionVars().GetReplicaRead()
 	if replicaReadType.IsFollowerRead() && !e.ctx.GetSessionVars().StmtCtx.RCCheckTS {
-		snapshot.SetOption(kv.ReplicaRead, replicaReadType)
+		e.snapshot.SetOption(kv.ReplicaRead, replicaReadType)
 	}
-	snapshot.SetOption(kv.TaskID, stmtCtx.TaskID)
-	snapshot.SetOption(kv.ReadReplicaScope, e.readReplicaScope)
-	snapshot.SetOption(kv.IsStalenessReadOnly, e.isStaleness)
+	e.snapshot.SetOption(kv.TaskID, stmtCtx.TaskID)
+	e.snapshot.SetOption(kv.ReadReplicaScope, e.readReplicaScope)
+	e.snapshot.SetOption(kv.IsStalenessReadOnly, e.isStaleness)
 	failpoint.Inject("assertBatchPointReplicaOption", func(val failpoint.Value) {
 		assertScope := val.(string)
 		if replicaReadType.IsClosestRead() && assertScope != e.readReplicaScope {
@@ -141,26 +126,25 @@ func (e *BatchPointGetExec) Open(context.Context) error {
 	})
 
 	if replicaReadType.IsClosestRead() && e.readReplicaScope != kv.GlobalTxnScope {
-		snapshot.SetOption(kv.MatchStoreLabels, []*metapb.StoreLabel{
+		e.snapshot.SetOption(kv.MatchStoreLabels, []*metapb.StoreLabel{
 			{
 				Key:   placement.DCLabelKey,
 				Value: e.readReplicaScope,
 			},
 		})
 	}
-	setOptionForTopSQL(stmtCtx, snapshot)
-	var batchGetter kv.BatchGetter = snapshot
+	setOptionForTopSQL(stmtCtx, e.snapshot)
+	var batchGetter kv.BatchGetter = e.snapshot
 	if txn.Valid() {
 		lock := e.tblInfo.Lock
 		if e.lock {
-			batchGetter = driver.NewBufferBatchGetter(txn.GetMemBuffer(), &PessimisticLockCacheGetter{txnCtx: txnCtx}, snapshot)
+			batchGetter = driver.NewBufferBatchGetter(txn.GetMemBuffer(), &PessimisticLockCacheGetter{txnCtx: txnCtx}, e.snapshot)
 		} else if lock != nil && (lock.Tp == model.TableLockRead || lock.Tp == model.TableLockReadOnly) && e.ctx.GetSessionVars().EnablePointGetCache {
-			batchGetter = newCacheBatchGetter(e.ctx, e.tblInfo.ID, snapshot)
+			batchGetter = newCacheBatchGetter(e.ctx, e.tblInfo.ID, e.snapshot)
 		} else {
-			batchGetter = driver.NewBufferBatchGetter(txn.GetMemBuffer(), nil, snapshot)
+			batchGetter = driver.NewBufferBatchGetter(txn.GetMemBuffer(), nil, e.snapshot)
 		}
 	}
-	e.snapshot = snapshot
 	e.batchGetter = batchGetter
 	return nil
 }

executor/builder.go

@@ -1561,6 +1561,14 @@ func (b *executorBuilder) getSnapshotTS() (uint64, error) {
 	return txnManager.GetStmtReadTS()
 }
 
+func (b *executorBuilder) getSnapshot() (kv.Snapshot, error) {
+	txnManager := sessiontxn.GetTxnManager(b.ctx)
+	if b.inInsertStmt || b.inUpdateStmt || b.inDeleteStmt || b.inSelectLockStmt {
+		return txnManager.GetForUpdateSnapshot()
+	}
+	return txnManager.GetReadSnapshot()
+}
+
 func (b *executorBuilder) buildMemTable(v *plannercore.PhysicalMemTable) Executor {
 	switch v.DBName.L {
 	case util.MetricSchemaName.L:
@@ -4600,6 +4608,12 @@ func (b *executorBuilder) buildBatchPointGet(plan *plannercore.BatchPointGetPlan
 		return nil
 	}
 
+	snapshot, err := b.getSnapshot()
+	if err != nil {
+		b.err = err
+		return nil
+	}
+
 	decoder := NewRowDecoder(b.ctx, plan.Schema(), plan.TblInfo)
 	e := &BatchPointGetExec{
 		baseExecutor:     newBaseExecutor(b.ctx, plan.Schema(), plan.ID()),
@@ -4618,10 +4632,13 @@ func (b *executorBuilder) buildBatchPointGet(plan *plannercore.BatchPointGetPlan
 		singlePart:       plan.SinglePart,
 		partTblID:        plan.PartTblID,
 		columns:          plan.Columns,
+		snapshot:         snapshot,
 	}
 
 	if plan.TblInfo.TableCacheStatusType == model.TableCacheStatusEnable {
-		e.cacheTable = b.getCacheTable(plan.TblInfo, snapshotTS)
+		if cacheTable := b.getCacheTable(plan.TblInfo, snapshotTS); cacheTable != nil {
+			e.snapshot = cacheTableSnapshot{e.snapshot, cacheTable}
+		}
 	}
 
 	if plan.TblInfo.TempTableType != model.TempTableNone {

executor/point_get.go

@@ -62,19 +62,27 @@ func (b *executorBuilder) buildPointGet(p *plannercore.PointGetPlan) Executor {
 		return nil
 	}
 
+	snapshot, err := b.getSnapshot()
+	if err != nil {
+		b.err = err
+		return nil
+	}
+
 	e := &PointGetExecutor{
 		baseExecutor:     newBaseExecutor(b.ctx, p.Schema(), p.ID()),
 		readReplicaScope: b.readReplicaScope,
 		isStaleness:      b.isStaleness,
 	}
 
-	if p.TblInfo.TableCacheStatusType == model.TableCacheStatusEnable {
-		e.cacheTable = b.getCacheTable(p.TblInfo, snapshotTS)
-	}
-
 	e.base().initCap = 1
 	e.base().maxChunkSize = 1
-	e.Init(p, snapshotTS)
+	e.Init(p, snapshotTS, snapshot)
+
+	if p.TblInfo.TableCacheStatusType == model.TableCacheStatusEnable {
+		if cacheTable := b.getCacheTable(p.TblInfo, snapshotTS); cacheTable != nil {
+			e.snapshot = cacheTableSnapshot{e.snapshot, cacheTable}
+		}
+	}
 
 	if e.lock {
 		b.hasLock = true
@@ -112,18 +120,18 @@ type PointGetExecutor struct {
 	// virtualColumnRetFieldTypes records the RetFieldTypes of virtual columns.
 	virtualColumnRetFieldTypes []*types.FieldType
 
-	stats      *runtimeStatsWithSnapshot
-	cacheTable kv.MemBuffer
+	stats *runtimeStatsWithSnapshot
 }
 
 // Init set fields needed for PointGetExecutor reuse, this does NOT change baseExecutor field
-func (e *PointGetExecutor) Init(p *plannercore.PointGetPlan, snapshotTS uint64) {
+func (e *PointGetExecutor) Init(p *plannercore.PointGetPlan, snapshotTS uint64, snapshot kv.Snapshot) {
 	decoder := NewRowDecoder(e.ctx, p.Schema(), p.TblInfo)
 	e.tblInfo = p.TblInfo
 	e.handle = p.Handle
 	e.idxInfo = p.IndexInfo
 	e.idxVals = p.IndexValues
 	e.snapshotTS = snapshotTS
+	e.snapshot = snapshot
 	e.done = false
 	if e.tblInfo.TempTableType == model.TempTableNone {
 		e.lock = p.Lock
@@ -152,24 +160,11 @@ func (e *PointGetExecutor) buildVirtualColumnInfo() {
 
 // Open implements the Executor interface.
 func (e *PointGetExecutor) Open(context.Context) error {
-	txnCtx := e.ctx.GetSessionVars().TxnCtx
-	snapshotTS := e.snapshotTS
 	var err error
 	e.txn, err = e.ctx.Txn(false)
 	if err != nil {
 		return err
 	}
-	if e.txn.Valid() && txnCtx.StartTS == txnCtx.GetForUpdateTS() && txnCtx.StartTS == snapshotTS {
-		e.snapshot = e.txn.GetSnapshot()
-	} else {
-		e.snapshot = e.ctx.GetSnapshotWithTS(snapshotTS)
-	}
-	if e.ctx.GetSessionVars().StmtCtx.RCCheckTS {
-		e.snapshot.SetOption(kv.IsolationLevel, kv.RCCheckTS)
-	}
-	if e.cacheTable != nil {
-		e.snapshot = cacheTableSnapshot{e.snapshot, e.cacheTable}
-	}
 	if err := e.verifyTxnScope(); err != nil {
 		return err
 	}

session/session.go

@@ -3156,13 +3156,6 @@ func (s *session) RefreshTxnCtx(ctx context.Context) error {
 	return sessiontxn.NewTxn(ctx, s)
 }
 
-// GetSnapshotWithTS returns a snapshot with ts.
-func (s *session) GetSnapshotWithTS(ts uint64) kv.Snapshot {
-	snap := s.GetStore().GetSnapshot(kv.Version{Ver: ts})
-	snap.SetOption(kv.SnapInterceptor, s.getSnapshotInterceptor())
-	return snap
-}
-
 // GetStore gets the store of session.
 func (s *session) GetStore() kv.Storage {
 	return s.store

session/txnmanager.go

@@ -94,6 +94,20 @@ func (m *txnManager) GetStmtForUpdateTS() (uint64, error) {
 	return ts, nil
 }
 
+func (m *txnManager) GetReadSnapshot() (kv.Snapshot, error) {
+	if m.ctxProvider == nil {
+		return nil, errors.New("context provider not set")
+	}
+	return m.ctxProvider.GetReadSnapshot()
+}
+
+func (m *txnManager) GetForUpdateSnapshot() (kv.Snapshot, error) {
+	if m.ctxProvider == nil {
+		return nil, errors.New("context provider not set")
+	}
+	return m.ctxProvider.GetForUpdateSnapshot()
+}
+
 func (m *txnManager) GetContextProvider() sessiontxn.TxnContextProvider {
 	return m.ctxProvider
 }

sessionctx/context.go

@@ -109,9 +109,6 @@ type Context interface {
 	// only used to daemon session like `statsHandle` to detect global variable change.
 	RefreshVars(context.Context) error
 
-	// GetSnapshotWithTS returns a snapshot with start ts
-	GetSnapshotWithTS(ts uint64) kv.Snapshot
-
 	// GetStore returns the store of session.
 	GetStore() kv.Storage
 

sessiontxn/interface.go

@@ -121,6 +121,10 @@ type TxnContextProvider interface {
 	GetStmtReadTS() (uint64, error)
 	// GetStmtForUpdateTS returns the read timestamp used by update/insert/delete or select ... for update
 	GetStmtForUpdateTS() (uint64, error)
+	// GetReadSnapshot get snapshot with read ts
+	GetReadSnapshot() (kv.Snapshot, error)
+	// GetForUpdateSnapshot get snapshot with for update ts
+	GetForUpdateSnapshot() (kv.Snapshot, error)
 
 	// OnInitialize is the hook that should be called when enter a new txn with this provider
 	OnInitialize(ctx context.Context, enterNewTxnType EnterNewTxnType) error
@@ -147,6 +151,10 @@ type TxnManager interface {
 	GetStmtForUpdateTS() (uint64, error)
 	// GetContextProvider returns the current TxnContextProvider
 	GetContextProvider() TxnContextProvider
+	// GetReadSnapshot get snapshot with read ts
+	GetReadSnapshot() (kv.Snapshot, error)
+	// GetForUpdateSnapshot get snapshot with for update ts
+	GetForUpdateSnapshot() (kv.Snapshot, error)
 
 	// EnterNewTxn enters a new transaction.
 	EnterNewTxn(ctx context.Context, req *EnterNewTxnRequest) error

sessiontxn/isolation/base.go

@@ -279,3 +279,37 @@ func (p *baseTxnContextProvider) AdviseWarmup() error {
 func (p *baseTxnContextProvider) AdviseOptimizeWithPlan(_ interface{}) error {
 	return nil
 }
+
+// GetReadSnapshot get snapshot with read ts
+func (p *baseTxnContextProvider) GetReadSnapshot() (kv.Snapshot, error) {
+	ts, err := p.GetStmtReadTS()
+	if err != nil {
+		return nil, err
+	}
+
+	return p.getSnapshotByTS(ts)
+}
+
+// GetForUpdateSnapshot get snapshot with for update ts
+func (p *baseTxnContextProvider) GetForUpdateSnapshot() (kv.Snapshot, error) {
+	ts, err := p.GetStmtForUpdateTS()
+	if err != nil {
+		return nil, err
+	}
+
+	return p.getSnapshotByTS(ts)
+}
+
+func (p *baseTxnContextProvider) getSnapshotByTS(snapshotTS uint64) (kv.Snapshot, error) {
+	txn, err := p.sctx.Txn(false)
+	if err != nil {
+		return nil, err
+	}
+
+	txnCtx := p.sctx.GetSessionVars().TxnCtx
+	if txn.Valid() && txnCtx.StartTS == txnCtx.GetForUpdateTS() && txnCtx.StartTS == snapshotTS {
+		return txn.GetSnapshot(), nil
+	}
+
+	return sessiontxn.GetSnapshotWithTS(p.sctx, snapshotTS), nil
+}

sessiontxn/isolation/readcommitted.go

@@ -257,3 +257,17 @@ func (p *PessimisticRCTxnContextProvider) AdviseOptimizeWithPlan(val interface{}
 
 	return nil
 }
+
+// GetReadSnapshot get snapshot with read ts
+func (p *PessimisticRCTxnContextProvider) GetReadSnapshot() (kv.Snapshot, error) {
+	snapshot, err := p.baseTxnContextProvider.GetForUpdateSnapshot()
+	if err != nil {
+		return nil, err
+	}
+
+	if p.sctx.GetSessionVars().StmtCtx.RCCheckTS {
+		snapshot.SetOption(kv.IsolationLevel, kv.RCCheckTS)
+	}
+
+	return snapshot, nil
+}

sessiontxn/staleread/provider.go

@@ -149,3 +149,28 @@ func (p *StalenessTxnContextProvider) AdviseWarmup() error {
 func (p *StalenessTxnContextProvider) AdviseOptimizeWithPlan(_ interface{}) error {
 	return nil
 }
+
+// GetReadSnapshot get snapshot with read ts
+func (p *StalenessTxnContextProvider) GetReadSnapshot() (kv.Snapshot, error) {
+	return p.getStalenessSnapshot()
+}
+
+// GetForUpdateSnapshot get snapshot with for update ts
+func (p *StalenessTxnContextProvider) GetForUpdateSnapshot() (kv.Snapshot, error) {
+	return p.getStalenessSnapshot()
+}
+
+func (p *StalenessTxnContextProvider) getStalenessSnapshot() (kv.Snapshot, error) {
+	snapshotTS := p.ts
+	txn, err := p.sctx.Txn(false)
+	if err != nil {
+		return nil, err
+	}
+
+	txnCtx := p.sctx.GetSessionVars().TxnCtx
+	if txn.Valid() && txnCtx.StartTS == txnCtx.GetForUpdateTS() && txnCtx.StartTS == snapshotTS {
+		return txn.GetSnapshot(), nil
+	}
+
+	return sessiontxn.GetSnapshotWithTS(p.sctx, snapshotTS), nil
+}

sessiontxn/txn.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/tidb/kv"
 	"github.com/pingcap/tidb/sessionctx"
 	"github.com/pingcap/tidb/sessionctx/variable"
+	"github.com/pingcap/tidb/table/temptable"
 	"github.com/tikv/client-go/v2/oracle"
 )
 
@@ -71,6 +72,13 @@ func CanReuseTxnWhenExplicitBegin(sctx sessionctx.Context) bool {
 	return txnCtx.History == nil && !txnCtx.IsStaleness && sessVars.SnapshotTS == 0
 }
 
+// GetSnapshotWithTS returns a snapshot with ts.
+func GetSnapshotWithTS(s sessionctx.Context, ts uint64) kv.Snapshot {
+	snap := s.GetStore().GetSnapshot(kv.Version{Ver: ts})
+	snap.SetOption(kv.SnapInterceptor, temptable.SessionSnapshotInterceptor(s))
+	return snap
+}
+
 // SetTxnAssertionLevel sets assertion level of a transactin. Note that assertion level should be set only once just
 // after creating a new transaction.
 func SetTxnAssertionLevel(txn kv.Transaction, assertionLevel variable.AssertionLevel) {