sessiontxn/staleread: don't overwrite StaleTSOProvider for single statement (#65110)

close #65090

Author: you06

pkg/expression/sessionexpr/sessionctx_test.go

@@ -122,7 +122,7 @@ func TestSessionEvalContextCurrentTime(t *testing.T) {
 	impl := sessionexpr.NewEvalContext(ctx)
 
 	var now atomic.Pointer[time.Time]
-	sc.SetStaleTSOProvider(func() (uint64, error) {
+	sc.SetStaleTSOProviderIfNotExist(func() (uint64, error) {
 		v := time.UnixMilli(123456789)
 		// should only be called once
 		require.True(t, now.CompareAndSwap(nil, &v))
@@ -142,7 +142,7 @@ func TestSessionEvalContextCurrentTime(t *testing.T) {
 	require.Equal(t, v.UnixNano(), tm.UnixNano())
 
 	// now should return the system variable if "timestamp" is set
-	sc.SetStaleTSOProvider(nil)
+	sc.SetStaleTSOProviderIfNotExist(nil)
 	sc.Reset()
 	require.NoError(t, vars.SetSystemVar("timestamp", "7654321.875"))
 	tm, err = impl.CurrentTime()

pkg/sessionctx/stmtctx/stmtctx.go

@@ -1191,10 +1191,13 @@ func (sc *StatementContext) DetachMemDiskTracker() {
 	}
 }
 
-// SetStaleTSOProvider sets the stale TSO provider.
-func (sc *StatementContext) SetStaleTSOProvider(eval func() (uint64, error)) {
+// SetStaleTSOProviderIfNotExist sets the stale TSO provider.
+func (sc *StatementContext) SetStaleTSOProviderIfNotExist(eval func() (uint64, error)) {
 	sc.StaleTSOProvider.Lock()
 	defer sc.StaleTSOProvider.Unlock()
+	if sc.StaleTSOProvider.eval != nil {
+		return
+	}
 	sc.StaleTSOProvider.value = nil
 	sc.StaleTSOProvider.eval = eval
 }

pkg/sessiontxn/staleread/BUILD.bazel

@@ -46,7 +46,7 @@ go_test(
         "provider_test.go",
     ],
     flaky = True,
-    shard_count = 6,
+    shard_count = 7,
     deps = [
         ":staleread",
         "//pkg/domain",
@@ -56,6 +56,7 @@ go_test(
         "//pkg/parser/ast",
         "//pkg/parser/auth",
         "//pkg/sessionctx",
+        "//pkg/sessionctx/stmtctx",
         "//pkg/sessiontxn",
         "//pkg/store/mockstore",
         "//pkg/table/temptable",

pkg/sessiontxn/staleread/processor_test.go

@@ -27,6 +27,7 @@ import (
 	"github.com/pingcap/tidb/pkg/parser"
 	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/sessionctx"
+	"github.com/pingcap/tidb/pkg/sessionctx/stmtctx"
 	"github.com/pingcap/tidb/pkg/sessiontxn/staleread"
 	"github.com/pingcap/tidb/pkg/store/mockstore"
 	"github.com/pingcap/tidb/pkg/table/temptable"
@@ -428,3 +429,35 @@ func createProcessor(t *testing.T, se sessionctx.Context) staleread.Processor {
 	require.Nil(t, processor.GetStalenessInfoSchema())
 	return processor
 }
+
+func TestConsistentCalculateAsOfTsExpr(t *testing.T) {
+	store := testkit.CreateMockStore(t, mockstore.WithStoreType(mockstore.EmbedUnistore))
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("set time_zone = '+00:00'")
+	tk.Session().GetSessionVars().TimeZone = time.UTC
+
+	p := parser.New()
+	stmt, err := p.ParseOneStmt(`select now(3) - interval 1 second`, "", "")
+	require.NoError(t, err)
+	secondTsExpr := stmt.(*ast.SelectStmt).Fields.Fields[0].Expr
+	stmt, err = p.ParseOneStmt(`select now(3) - interval 3 second`, "", "")
+	require.NoError(t, err)
+	threeSecondTsExpr := stmt.(*ast.SelectStmt).Fields.Fields[0].Expr
+
+	se := tk.Session()
+	se.GetSessionVars().StmtCtx = stmtctx.NewStmtCtxWithTimeZone(time.UTC)
+
+	ts1, err := staleread.CalculateAsOfTsExpr(context.Background(), se.GetPlanCtx(), secondTsExpr)
+	require.NoError(t, err)
+
+	time.Sleep(10 * time.Millisecond)
+
+	ts2, err := staleread.CalculateAsOfTsExpr(context.Background(), se.GetPlanCtx(), secondTsExpr)
+	require.NoError(t, err)
+	require.Equal(t, ts1, ts2)
+
+	ts3, err := staleread.CalculateAsOfTsExpr(context.Background(), se.GetPlanCtx(), threeSecondTsExpr)
+	require.NoError(t, err)
+	require.True(t, ts3 < ts1)
+	require.Equal(t, ts1-ts3, uint64(2000<<18))
+}

pkg/sessiontxn/staleread/util.go

@@ -34,7 +34,7 @@ import (
 
 // CalculateAsOfTsExpr calculates the TsExpr of AsOfClause to get a StartTS.
 func CalculateAsOfTsExpr(ctx context.Context, sctx planctx.PlanContext, tsExpr ast.ExprNode) (uint64, error) {
-	sctx.GetSessionVars().StmtCtx.SetStaleTSOProvider(func() (uint64, error) {
+	sctx.GetSessionVars().StmtCtx.SetStaleTSOProviderIfNotExist(func() (uint64, error) {
 		failpoint.Inject("mockStaleReadTSO", func(val failpoint.Value) (uint64, error) {
 			return uint64(val.(int)), nil
 		})