split implementation of checks for stale read and snapshot read

Signed-off-by: Yilin Chen <sticnarf@gmail.com>

executor/set.go

@@ -161,7 +161,12 @@ func (e *SetExecutor) setSysVariable(ctx context.Context, name string, v *expres
 	newSnapshotTS := getSnapshotTSByName()
 	newSnapshotIsSet := newSnapshotTS > 0 && newSnapshotTS != oldSnapshotTS
 	if newSnapshotIsSet {
-		if err = sessionctx.ValidateReadTS(ctx, e.ctx, newSnapshotTS); err != nil {
+		if name == variable.TiDBTxnReadTS {
+			err = sessionctx.ValidateStaleReadTS(ctx, e.ctx, newSnapshotTS)
+		} else {
+			err = sessionctx.ValidateSnapshotReadTS(ctx, e.ctx, newSnapshotTS)
+		}
+		if err != nil {
 			fallbackOldSnapshotTS()
 			return err
 		}

planner/core/planbuilder.go

@@ -2453,7 +2453,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,
 			if err != nil {
 				return nil, err
 			}
-			if err := sessionctx.ValidateReadTS(ctx, b.ctx, startTS); err != nil {
+			if err := sessionctx.ValidateStaleReadTS(ctx, b.ctx, startTS); err != nil {
 				return nil, err
 			}
 			p.StaleTxnStartTS = startTS

planner/core/preprocess.go

@@ -1526,7 +1526,7 @@ func (p *preprocessor) handleAsOfAndReadTS(node *ast.AsOfClause) {
 		if p.err != nil {
 			return
 		}
-		if err := sessionctx.ValidateReadTS(context.Background(), p.ctx, ts); err != nil {
+		if err := sessionctx.ValidateStaleReadTS(context.Background(), p.ctx, ts); err != nil {
 			p.err = errors.Trace(err)
 			return
 		}

sessionctx/context.go

@@ -153,8 +153,8 @@ const (
 	LastExecuteDDL basicCtxType = 3
 )
 
-// ValidateReadTS validates that readTS does not exceed the current timestamp
-func ValidateReadTS(ctx context.Context, sctx Context, readTS uint64) error {
+// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp
+func ValidateSnapshotReadTS(ctx context.Context, sctx Context, readTS uint64) error {
 	txnScope := sctx.GetSessionVars().CheckAndGetTxnScope()
 	latestTS, err := sctx.GetStore().GetOracle().GetLowResolutionTimestamp(ctx, &oracle.Option{TxnScope: txnScope})
 	// If we fail to get latestTS or the readTS exceeds it, get a timestamp from PD to double check
@@ -170,3 +170,22 @@ func ValidateReadTS(ctx context.Context, sctx Context, readTS uint64) error {
 	}
 	return nil
 }
+
+// ValidateStaleReadTS validates that readTS does not exceed the current time not strictly.
+func ValidateStaleReadTS(ctx context.Context, sctx Context, readTS uint64) error {
+	txnScope := sctx.GetSessionVars().CheckAndGetTxnScope()
+	currentTS, err := sctx.GetStore().GetOracle().GetStaleTimestamp(ctx, txnScope, 0)
+	// If we fail to calculate currentTS from local time, fallback to get a timestamp from PD
+	if err != nil {
+		metrics.ValidateReadTSFromPDCount.Inc()
+		currentVer, err := sctx.GetStore().CurrentVersion(txnScope)
+		if err != nil {
+			return errors.Errorf("fail to validate read timestamp: %v", err)
+		}
+		currentTS = currentVer.Ver
+	}
+	if readTS > currentTS {
+		return errors.Errorf("cannot set read timestamp to a future time")
+	}
+	return nil
+}