*: Update client-go and verify all read ts (pingcap#58054)

MyonKeminta · ekexium · commit 4d05a0519c7e · 2025-01-15T13:01:12.000+08:00
ref pingcap#57786
diff --git a/DEPS.bzl b/DEPS.bzl
@@ -4179,8 +4179,8 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sum = "h1:E+JsLmxpa4rsT0rS9m52+FhiiRbQxZYn079T4JD12jU=",
-        version = "v2.0.8-0.20241204085508-80a6b021f0f6",
+        sum = "h1:zrgHZRnd7Veio1Q+jV9LVa/JAfBi2Y7fGx91s7Uw+fo=",
+        version = "v2.0.8-0.20250108151910-7b4409c11946",
     )
     go_repository(
         name = "com_github_tikv_pd",
diff --git a/ddl/column_change_test.go b/ddl/column_change_test.go
@@ -35,7 +35,6 @@ import (
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/testkit/external"
 	"github.com/pingcap/tidb/types"
-	"github.com/pingcap/tidb/util/mock"
 	"github.com/stretchr/testify/require"
 )
 
diff --git a/executor/set.go b/executor/set.go
@@ -198,7 +198,8 @@ func (e *SetExecutor) setSysVariable(ctx context.Context, name string, v *expres
 	newSnapshotTS := getSnapshotTSByName()
 	newSnapshotIsSet := newSnapshotTS > 0 && newSnapshotTS != oldSnapshotTS
 	if newSnapshotIsSet {
-		err = sessionctx.ValidateSnapshotReadTS(ctx, e.ctx.GetStore(), newSnapshotTS)
+		isStaleRead := name == variable.TiDBTxnReadTS
+		err = sessionctx.ValidateSnapshotReadTS(ctx, e.ctx.GetStore(), newSnapshotTS, isStaleRead)
 		if name != variable.TiDBTxnReadTS {
 			// Also check gc safe point for snapshot read.
 			// We don't check snapshot with gc safe point for read_ts
diff --git a/go.mod b/go.mod
@@ -94,7 +94,7 @@ require (
 	github.com/stretchr/testify v1.8.2
 	github.com/tdakkota/asciicheck v0.2.0
 	github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
-	github.com/tikv/client-go/v2 v2.0.8-0.20241204085508-80a6b021f0f6
+	github.com/tikv/client-go/v2 v2.0.8-0.20250108151910-7b4409c11946
 	github.com/tikv/pd/client v0.0.0-20240725070735-fb162bf0aa3f
 	github.com/timakin/bodyclose v0.0.0-20221125081123-e39cf3fc478e
 	github.com/twmb/murmur3 v1.1.6
diff --git a/go.sum b/go.sum
@@ -968,8 +968,8 @@ github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJf
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW93SG+q0F8KI+yFrcIDT4c/RNoc4=
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
-github.com/tikv/client-go/v2 v2.0.8-0.20241204085508-80a6b021f0f6 h1:E+JsLmxpa4rsT0rS9m52+FhiiRbQxZYn079T4JD12jU=
-github.com/tikv/client-go/v2 v2.0.8-0.20241204085508-80a6b021f0f6/go.mod h1:45NuHB8x+VAoztMIjF6hEgXvPQXhXWPfMxDg0N8CoRY=
+github.com/tikv/client-go/v2 v2.0.8-0.20250108151910-7b4409c11946 h1:zrgHZRnd7Veio1Q+jV9LVa/JAfBi2Y7fGx91s7Uw+fo=
+github.com/tikv/client-go/v2 v2.0.8-0.20250108151910-7b4409c11946/go.mod h1:45NuHB8x+VAoztMIjF6hEgXvPQXhXWPfMxDg0N8CoRY=
 github.com/tikv/pd/client v0.0.0-20240725070735-fb162bf0aa3f h1:Szw9YxqGGEneSniBd4ep09jgB77cKUy+AuhKOmdGPdE=
 github.com/tikv/pd/client v0.0.0-20240725070735-fb162bf0aa3f/go.mod h1:QCBn54O5lhfkYfxj8Tyiqaxue/mthHEMyi7AqJP/+n4=
 github.com/timakin/bodyclose v0.0.0-20221125081123-e39cf3fc478e h1:MV6KaVu/hzByHP0UvJ4HcMGE/8a6A4Rggc/0wx2AvJo=
diff --git a/planner/core/planbuilder.go b/planner/core/planbuilder.go
@@ -3568,7 +3568,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,
 			if err != nil {
 				return nil, err
 			}
-			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {
+			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {
 				return nil, err
 			}
 			p.StaleTxnStartTS = startTS
@@ -3582,7 +3582,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,
 			if err != nil {
 				return nil, err
 			}
-			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {
+			if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {
 				return nil, err
 			}
 			p.StaleTxnStartTS = startTS
diff --git a/sessionctx/context.go b/sessionctx/context.go
@@ -219,9 +219,12 @@ const (
 	LastExecuteDDL basicCtxType = 3
 )
 
-// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp
-func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64) error {
-	return store.GetOracle().ValidateSnapshotReadTS(ctx, readTS, &oracle.Option{TxnScope: oracle.GlobalTxnScope})
+// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp.
+// For read requests to the storage, the check can be implicitly performed when sending the RPC request. So this
+// function is only needed when it's not proper to delay the check to when RPC requests are being sent (e.g., `BEGIN`
+// statements that don't make reading operation immediately).
+func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64, isStaleRead bool) error {
+	return store.GetOracle().ValidateReadTS(ctx, readTS, isStaleRead, &oracle.Option{TxnScope: oracle.GlobalTxnScope})
 }
 
 // SysProcTracker is used to track background sys processes
diff --git a/sessiontxn/staleread/processor.go b/sessiontxn/staleread/processor.go
@@ -285,7 +285,7 @@ func parseAndValidateAsOf(ctx context.Context, sctx sessionctx.Context, asOf *as
 		return 0, err
 	}
 
-	if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts); err != nil {
+	if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts, true); err != nil {
 		return 0, err
 	}
 
diff --git a/sessiontxn/staleread/util.go b/sessiontxn/staleread/util.go
@@ -84,7 +84,7 @@ func CalculateTsWithReadStaleness(ctx context.Context, sctx sessionctx.Context,
 		// If the final calculated exceeds the min safe ts, we are not sure whether the ts is safe to read (note that
 		// reading with a ts larger than PD's max allocated ts + 1 is unsafe and may break linearizability).
 		// So in this case, do an extra check on it.
-		err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS)
+		err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS, true)
 		if err != nil {
 			return 0, err
 		}
diff --git a/store/copr/BUILD.bazel b/store/copr/BUILD.bazel
@@ -51,6 +51,7 @@ go_library(
         "@com_github_tikv_client_go_v2//config",
         "@com_github_tikv_client_go_v2//error",
         "@com_github_tikv_client_go_v2//metrics",
+        "@com_github_tikv_client_go_v2//oracle",
         "@com_github_tikv_client_go_v2//tikv",
         "@com_github_tikv_client_go_v2//tikvrpc",
         "@com_github_tikv_client_go_v2//txnkv/txnlock",
diff --git a/store/copr/batch_coprocessor.go b/store/copr/batch_coprocessor.go
@@ -1150,7 +1150,7 @@ func (b *batchCopIterator) retryBatchCopTask(ctx context.Context, bo *backoff.Ba
 const readTimeoutUltraLong = 3600 * time.Second // For requests that may scan many regions for tiflash.
 
 func (b *batchCopIterator) handleTaskOnce(ctx context.Context, bo *backoff.Backoffer, task *batchCopTask) ([]*batchCopTask, error) {
-	sender := NewRegionBatchRequestSender(b.store.GetRegionCache(), b.store.GetTiKVClient(), b.enableCollectExecutionInfo)
+	sender := NewRegionBatchRequestSender(b.store.GetRegionCache(), b.store.GetTiKVClient(), b.store.store.GetOracle(), b.enableCollectExecutionInfo)
 	var regionInfos = make([]*coprocessor.RegionInfo, 0, len(task.regionInfos))
 	for _, ri := range task.regionInfos {
 		regionInfos = append(regionInfos, ri.toCoprocessorRegionInfo())
diff --git a/store/copr/batch_request_sender.go b/store/copr/batch_request_sender.go
@@ -16,6 +16,7 @@ package copr
 
 import (
 	"context"
+	"github.com/tikv/client-go/v2/oracle"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -56,9 +57,9 @@ type RegionBatchRequestSender struct {
 }
 
 // NewRegionBatchRequestSender creates a RegionBatchRequestSender object.
-func NewRegionBatchRequestSender(cache *RegionCache, client tikv.Client, enableCollectExecutionInfo bool) *RegionBatchRequestSender {
+func NewRegionBatchRequestSender(cache *RegionCache, client tikv.Client, oracle oracle.Oracle, enableCollectExecutionInfo bool) *RegionBatchRequestSender {
 	return &RegionBatchRequestSender{
-		RegionRequestSender:        tikv.NewRegionRequestSender(cache.RegionCache, client),
+		RegionRequestSender:        tikv.NewRegionRequestSender(cache.RegionCache, client, oracle),
 		enableCollectExecutionInfo: enableCollectExecutionInfo,
 	}
 }
diff --git a/store/copr/mpp.go b/store/copr/mpp.go
@@ -275,7 +275,7 @@ func (m *mppIterator) handleDispatchReq(ctx context.Context, bo *Backoffer, req
 	// Or else it's the task without region, which always happens in high layer task without table.
 	// In that case
 	if originalTask != nil {
-		sender := NewRegionBatchRequestSender(m.store.GetRegionCache(), m.store.GetTiKVClient(), m.enableCollectExecutionInfo)
+		sender := NewRegionBatchRequestSender(m.store.GetRegionCache(), m.store.GetTiKVClient(), m.store.store.GetOracle(), m.enableCollectExecutionInfo)
 		rpcResp, retry, _, err = sender.SendReqToAddr(bo, originalTask.ctx, originalTask.regionInfos, wrappedReq, tikv.ReadTimeoutMedium)
 		// No matter what the rpc error is, we won't retry the mpp dispatch tasks.
 		// TODO: If we want to retry, we must redo the plan fragment cutting and task scheduling.
diff --git a/util/mock/BUILD.bazel b/util/mock/BUILD.bazel
@@ -21,6 +21,7 @@ go_library(
         "//sessionctx/variable",
         "//util",
         "//util/disk",
+        "//util/logutil",
         "//util/memory",
         "//util/sli",
         "//util/sqlexec",
diff --git a/util/mock/context.go b/util/mock/context.go
@@ -18,6 +18,7 @@ package mock
 import (
 	"context"
 	"fmt"
+	"github.com/pingcap/tidb/util/logutil"
 	"time"
 
 	"github.com/pingcap/errors"
@@ -67,7 +68,7 @@ type wrapTxn struct {
 }
 
 func (txn *wrapTxn) validOrPending() bool {
-	return txn.tsFuture != nil || txn.Transaction.Valid()
+	return txn.tsFuture != nil || (txn.Transaction != nil && txn.Transaction.Valid())
 }
 
 func (txn *wrapTxn) pending() bool {
@@ -173,7 +174,15 @@ func (c *Context) GetSessionVars() *variable.SessionVars {
 }
 
 // Txn implements sessionctx.Context Txn interface.
-func (c *Context) Txn(bool) (kv.Transaction, error) {
+func (c *Context) Txn(active bool) (kv.Transaction, error) {
+	if active {
+		if !c.txn.validOrPending() {
+			err := c.newTxn(context.Background())
+			if err != nil {
+				return nil, err
+			}
+		}
+	}
 	return &c.txn, nil
 }
 
@@ -253,10 +262,12 @@ func (c *Context) GetSessionPlanCache() sessionctx.PlanCache {
 	return c.pcache
 }
 
-// NewTxn implements the sessionctx.Context interface.
-func (c *Context) NewTxn(context.Context) error {
+// newTxn Creates new transaction on the session context.
+func (c *Context) newTxn(ctx context.Context) error {
 	if c.Store == nil {
-		return errors.New("store is not set")
+		logutil.Logger(ctx).Warn("mock.Context: No store is specified when trying to create new transaction. A fake transaction will be created. Note that this is unrecommended usage.")
+		c.fakeTxn()
+		return nil
 	}
 	if c.txn.Valid() {
 		err := c.txn.Commit(c.ctx)
@@ -273,14 +284,46 @@ func (c *Context) NewTxn(context.Context) error {
 	return nil
 }
 
+// fakeTxn is used to let some tests pass in the context without an available kv.Storage. Once usages to access
+// transactions without a kv.Storage are removed, this type should also be removed.
+// New code should never use this.
+type fakeTxn struct {
+	// The inner should always be nil.
+	kv.Transaction
+	startTS uint64
+}
+
+func (t *fakeTxn) StartTS() uint64 {
+	return t.startTS
+}
+
+func (*fakeTxn) SetDiskFullOpt(_ kvrpcpb.DiskFullOpt) {}
+
+func (*fakeTxn) SetOption(_ int, _ any) {}
+
+func (*fakeTxn) Get(ctx context.Context, _ kv.Key) ([]byte, error) {
+	// Check your implementation if you meet this error. It's dangerous if some calculation relies on the data but the
+	// read result is faked.
+	logutil.Logger(ctx).Warn("mock.Context: No store is specified but trying to access data from a transaction.")
+	return nil, nil
+}
+
+func (*fakeTxn) Valid() bool { return true }
+
+func (c *Context) fakeTxn() {
+	c.txn.Transaction = &fakeTxn{
+		startTS: 1,
+	}
+}
+
 // NewStaleTxnWithStartTS implements the sessionctx.Context interface.
 func (c *Context) NewStaleTxnWithStartTS(ctx context.Context, _ uint64) error {
-	return c.NewTxn(ctx)
+	return c.newTxn(ctx)
 }
 
 // RefreshTxnCtx implements the sessionctx.Context interface.
 func (c *Context) RefreshTxnCtx(ctx context.Context) error {
-	return errors.Trace(c.NewTxn(ctx))
+	return errors.Trace(c.newTxn(ctx))
 }
 
 // RefreshVars implements the sessionctx.Context interface.

Original file line number	Diff line number	Diff line change
`@@ -35,7 +35,6 @@ import (`
`35`	`35`	`"github.com/pingcap/tidb/testkit"`
`36`	`36`	`"github.com/pingcap/tidb/testkit/external"`
`37`	`37`	`"github.com/pingcap/tidb/types"`
`38`		`- "github.com/pingcap/tidb/util/mock"`
`39`	`38`	`"github.com/stretchr/testify/require"`
`40`	`39`	`)`
`41`	`40`
Original file line number	Diff line number	Diff line change
`@@ -3568,7 +3568,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,`
`3568`	`3568`	`if err != nil {`
`3569`	`3569`	`return nil, err`
`3570`	`3570`	`}`
`3571`		`- if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {`
	`3571`	`+ if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {`
`3572`	`3572`	`return nil, err`
`3573`	`3573`	`}`
`3574`	`3574`	`p.StaleTxnStartTS = startTS`
`@@ -3582,7 +3582,7 @@ func (b *PlanBuilder) buildSimple(ctx context.Context, node ast.StmtNode) (Plan,`
`3582`	`3582`	`if err != nil {`
`3583`	`3583`	`return nil, err`
`3584`	`3584`	`}`
`3585`		`- if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS); err != nil {`
	`3585`	`+ if err := sessionctx.ValidateSnapshotReadTS(ctx, b.ctx.GetStore(), startTS, true); err != nil {`
`3586`	`3586`	`return nil, err`
`3587`	`3587`	`}`
`3588`	`3588`	`p.StaleTxnStartTS = startTS`
Original file line number	Diff line number	Diff line change
`@@ -219,9 +219,12 @@ const (`
`219`	`219`	`LastExecuteDDL basicCtxType = 3`
`220`	`220`	`)`
`221`	`221`
`222`		`-// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp`
`223`		`-func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64) error {`
`224`		`- return store.GetOracle().ValidateSnapshotReadTS(ctx, readTS, &oracle.Option{TxnScope: oracle.GlobalTxnScope})`
	`222`	`+// ValidateSnapshotReadTS strictly validates that readTS does not exceed the PD timestamp.`
	`223`	`+// For read requests to the storage, the check can be implicitly performed when sending the RPC request. So this`
	`224`	+// function is only needed when it's not proper to delay the check to when RPC requests are being sent (e.g., `BEGIN`
	`225`	`+// statements that don't make reading operation immediately).`
	`226`	`+func ValidateSnapshotReadTS(ctx context.Context, store kv.Storage, readTS uint64, isStaleRead bool) error {`
	`227`	`+ return store.GetOracle().ValidateReadTS(ctx, readTS, isStaleRead, &oracle.Option{TxnScope: oracle.GlobalTxnScope})`
`225`	`228`	`}`
`226`	`229`
`227`	`230`	`// SysProcTracker is used to track background sys processes`
Original file line number	Diff line number	Diff line change
`@@ -285,7 +285,7 @@ func parseAndValidateAsOf(ctx context.Context, sctx sessionctx.Context, asOf *as`
`285`	`285`	`return 0, err`
`286`	`286`	`}`
`287`	`287`
`288`		`- if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts); err != nil {`
	`288`	`+ if err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), ts, true); err != nil {`
`289`	`289`	`return 0, err`
`290`	`290`	`}`
`291`	`291`
Original file line number	Diff line number	Diff line change
`@@ -84,7 +84,7 @@ func CalculateTsWithReadStaleness(ctx context.Context, sctx sessionctx.Context,`
`84`	`84`	`// If the final calculated exceeds the min safe ts, we are not sure whether the ts is safe to read (note that`
`85`	`85`	`// reading with a ts larger than PD's max allocated ts + 1 is unsafe and may break linearizability).`
`86`	`86`	`// So in this case, do an extra check on it.`
`87`		`- err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS)`
	`87`	`+ err = sessionctx.ValidateSnapshotReadTS(ctx, sctx.GetStore(), readTS, true)`
`88`	`88`	`if err != nil {`
`89`	`89`	`return 0, err`
`90`	`90`	`}`