Repository: tikv/client-go
Pull request #1695 (2025-07-03 02:22:41)

txnkv: prevent some actions from being interrupted by kill (#1665)
I have no permission to edit #1690 so create this PR to resolve the conflicts.

![2025-07-03_101139](https://github.com/user-attachments/assets/2db6ed9e-815c-44a9-a6f9-68005957e6bd)

ti-chi-bot[bot]: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by:
**Once this PR has been reviewed and has the lgtm label**, please assign [iosmanthus](https://github.com/iosmanthus) for approval. For more information see [the Code Review Process](https://book.prow.tidb.net/#/workflows/pr).
**Please ensure that each of them provides their approval before proceeding.**

The full list of commands accepted by this bot can be found [here](https://prow.tidb.net/command-help?repo=tikv%2Fclient-go).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/tikv/client-go/blob/tidb-8.5/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["iosmanthus"]} -->

Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 4158911..d9f9917 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -1329,7 +1329,16 @@ func (c *RegionCache) BatchLocateKeyRanges(bo *retry.Backoffer, keyRanges []kv.K
 
 	// 2. load remaining regions from pd client
 	for len(uncachedRanges) > 0 {
-		regions, err := c.BatchLoadRegionsWithKeyRanges(bo, uncachedRanges, defaultRegionsPerBatch, opts...)
+		// If we send too many ranges to PD, it may exceed the size of grpc limitation or cause a timeout error.
+		// So we limit the number of ranges per batch to avoid this issue.
+		maxRangesPerBatch := 16 * defaultRegionsPerBatch
+		var toBeSentRanges []pd.KeyRange
+		if len(uncachedRanges) > maxRangesPerBatch {
+			toBeSentRanges = uncachedRanges[:maxRangesPerBatch]
+		} else {
+			toBeSentRanges = uncachedRanges
+		}
+		regions, err := c.BatchLoadRegionsWithKeyRanges(bo, toBeSentRanges, defaultRegionsPerBatch, opts...)
 		if err != nil {
 			return nil, err
 		}
@@ -1791,7 +1800,7 @@ func (c *RegionCache) BatchLoadRegionsWithKeyRange(bo *retry.Backoffer, startKey
 	// TODO(youjiali1995): scanRegions always fetch regions from PD and these regions don't contain buckets information
 	// for less traffic, so newly inserted regions in region cache don't have buckets information. We should improve it.
 	for _, region := range regions {
-		c.insertRegionToCache(region, true, false)
+		c.insertRegionToCache(region, false, false)
 	}
 
 	return
@@ -1816,7 +1825,7 @@ func (c *RegionCache) BatchLoadRegionsWithKeyRanges(bo *retry.Backoffer, keyRang
 	defer c.mu.Unlock()
 
 	for _, region := range regions {
-		c.insertRegionToCache(region, true, false)
+		c.insertRegionToCache(region, false, false)
 	}
 	return
 }
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 3b7482a..3571b0e 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -893,8 +893,10 @@ func (s *RegionRequestSender) SendReqCtx(
 		}
 
 		// recheck whether the session/query is killed during the Next()
-		if err2 := bo.CheckKilled(); err2 != nil {
-			return nil, nil, retryTimes, err2
+		if req.IsInterruptible() {
+			if err2 := bo.CheckKilled(); err2 != nil {
+				return nil, nil, retryTimes, err2
+			}
 		}
 		if val, err := util.EvalFailpoint("mockRetrySendReqToRegion"); err == nil {
 			if val.(bool) {
diff --git a/metrics/metrics.go b/metrics/metrics.go
index ce24760..e3e01b0 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -47,7 +47,6 @@ var (
 	TiKVSendReqCounter                             *prometheus.CounterVec
 	TiKVSendReqTimeCounter                         *prometheus.CounterVec
 	TiKVRPCNetLatencyHistogram                     *prometheus.HistogramVec
-	TiKVCoprocessorHistogram                       *prometheus.HistogramVec
 	TiKVLockResolverCounter                        *prometheus.CounterVec
 	TiKVRegionErrorCounter                         *prometheus.CounterVec
 	TiKVRPCErrorCounter                            *prometheus.CounterVec
@@ -171,7 +170,7 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) {
 			Subsystem:   subsystem,
 			Name:        "request_seconds",
 			Help:        "Bucketed histogram of sending request duration.",
-			Buckets:     prometheus.ExponentialBuckets(0.0005, 2, 29), // 0.5ms ~ 1.5days
+			Buckets:     prometheus.ExponentialBuckets(0.0005, 2, 24), // 0.5ms ~ 1.2h
 			ConstLabels: constLabels,
 		}, []string{LblType, LblStore, LblStaleRead, LblScope})
 
@@ -199,20 +198,10 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) {
 			Subsystem:   subsystem,
 			Name:        "rpc_net_latency_seconds",
 			Help:        "Bucketed histogram of time difference between TiDB and TiKV.",
-			Buckets:     prometheus.ExponentialBuckets(5e-5, 2, 22), // 50us ~ 105s
+			Buckets:     prometheus.ExponentialBuckets(0.0001, 2, 20), // 0.1ms ~ 52s
 			ConstLabels: constLabels,
 		}, []string{LblStore, LblScope})
 
-	TiKVCoprocessorHistogram = prometheus.NewHistogramVec(
-		prometheus.HistogramOpts{
-			Namespace:   namespace,
-			Subsystem:   subsystem,
-			Name:        "cop_duration_seconds",
-			Help:        "Run duration of a single coprocessor task, includes backoff time.",
-			Buckets:     prometheus.ExponentialBuckets(0.0005, 2, 29), // 0.5ms ~ 1.5days
-			ConstLabels: constLabels,
-		}, []string{LblStore, LblStaleRead, LblScope})
-
 	TiKVLockResolverCounter = prometheus.NewCounterVec(
 		prometheus.CounterOpts{
 			Namespace:   namespace,
@@ -878,7 +867,6 @@ func RegisterMetrics() {
 	prometheus.MustRegister(TiKVSendReqCounter)
 	prometheus.MustRegister(TiKVSendReqTimeCounter)
 	prometheus.MustRegister(TiKVRPCNetLatencyHistogram)
-	prometheus.MustRegister(TiKVCoprocessorHistogram)
 	prometheus.MustRegister(TiKVLockResolverCounter)
 	prometheus.MustRegister(TiKVRegionErrorCounter)
 	prometheus.MustRegister(TiKVRPCErrorCounter)
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 46e952b..f395edf 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -339,6 +339,16 @@ func (req *Request) IsDebugReq() bool {
 	return false
 }
 
+// IsInterruptible checks if the request can be interrupted when the query is killed.
+func (req *Request) IsInterruptible() bool {
+	switch req.Type {
+	case CmdPessimisticRollback, CmdBatchRollback, CmdCommit:
+		return false
+	default:
+		return true
+	}
+}
+
 // Get returns GetRequest in request.
 func (req *Request) Get() *kvrpcpb.GetRequest {
 	return req.Req.(*kvrpcpb.GetRequest)
diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go
index d4ea080..73a3132 100644
--- a/txnkv/transaction/2pc.go
+++ b/txnkv/transaction/2pc.go
@@ -75,6 +75,7 @@ const slowRequestThreshold = time.Minute
 type twoPhaseCommitAction interface {
 	handleSingleBatch(*twoPhaseCommitter, *retry.Backoffer, batchMutations) error
 	tiKVTxnRegionsNumHistogram() prometheus.Observer
+	isInterruptible() bool
 	String() string
 }
 
@@ -1064,9 +1065,10 @@ func (c *twoPhaseCommitter) doActionOnBatches(
 	// killSignal should never be nil for TiDB
 	if c.txn != nil && c.txn.vars != nil && c.txn.vars.Killed != nil {
 		// Do not reset the killed flag here. Let the upper layer reset the flag.
-		// Before it resets, any request is considered valid to be killed.
+		// Before it resets, any request is considered valid to be killed if the
+		// corresponding action is interruptible.
 		status := atomic.LoadUint32(c.txn.vars.Killed)
-		if status != 0 {
+		if status != 0 && action.isInterruptible() {
 			logutil.BgLogger().Info(
 				"query is killed", zap.Uint32(
 					"signal",
diff --git a/txnkv/transaction/cleanup.go b/txnkv/transaction/cleanup.go
index eaf0668..5991b94 100644
--- a/txnkv/transaction/cleanup.go
+++ b/txnkv/transaction/cleanup.go
@@ -103,6 +103,10 @@ func (action actionCleanup) handleSingleBatch(c *twoPhaseCommitter, bo *retry.Ba
 	return nil
 }
 
+func (actionCleanup) isInterruptible() bool {
+	return false
+}
+
 func (c *twoPhaseCommitter) cleanupMutations(bo *retry.Backoffer, mutations CommitterMutations) error {
 	return c.doActionOnMutations(bo, actionCleanup{isInternal: c.txn.isInternal()}, mutations)
 }
diff --git a/txnkv/transaction/commit.go b/txnkv/transaction/commit.go
index b3ad8c1..b34f8b5 100644
--- a/txnkv/transaction/commit.go
+++ b/txnkv/transaction/commit.go
@@ -232,6 +232,10 @@ func (action actionCommit) handleSingleBatch(c *twoPhaseCommitter, bo *retry.Bac
 	return nil
 }
 
+func (actionCommit) isInterruptible() bool {
+	return false
+}
+
 func (c *twoPhaseCommitter) commitMutations(bo *retry.Backoffer, mutations CommitterMutations) error {
 	if span := opentracing.SpanFromContext(bo.GetCtx()); span != nil && span.Tracer() != nil {
 		span1 := span.Tracer().StartSpan("twoPhaseCommitter.commitMutations", opentracing.ChildOf(span.Context()))
diff --git a/txnkv/transaction/pessimistic.go b/txnkv/transaction/pessimistic.go
index fee43b3..3459c63 100644
--- a/txnkv/transaction/pessimistic.go
+++ b/txnkv/transaction/pessimistic.go
@@ -537,6 +537,10 @@ func (action actionPessimisticLock) handlePessimisticLockResponseForceLockMode(
 	return true, nil
 }
 
+func (actionPessimisticLock) isInterruptible() bool {
+	return true
+}
+
 func (actionPessimisticRollback) handleSingleBatch(
 	c *twoPhaseCommitter, bo *retry.Backoffer, batch batchMutations,
 ) error {
@@ -572,6 +576,10 @@ func (actionPessimisticRollback) handleSingleBatch(
 	return nil
 }
 
+func (actionPessimisticRollback) isInterruptible() bool {
+	return false
+}
+
 func (c *twoPhaseCommitter) pessimisticLockMutations(
 	bo *retry.Backoffer, lockCtx *kv.LockCtx, lockWaitMode kvrpcpb.PessimisticLockWakeUpMode,
 	mutations CommitterMutations,
diff --git a/txnkv/transaction/pipelined_flush.go b/txnkv/transaction/pipelined_flush.go
index 5797d46..1471214 100644
--- a/txnkv/transaction/pipelined_flush.go
+++ b/txnkv/transaction/pipelined_flush.go
@@ -290,6 +290,10 @@ func (action actionPipelinedFlush) handleSingleBatch(
 	}
 }
 
+func (actionPipelinedFlush) isInterruptible() bool {
+	return true
+}
+
 func (c *twoPhaseCommitter) pipelinedFlushMutations(bo *retry.Backoffer, mutations CommitterMutations, generation uint64) error {
 	if span := opentracing.SpanFromContext(bo.GetCtx()); span != nil && span.Tracer() != nil {
 		span1 := span.Tracer().StartSpan("twoPhaseCommitter.pipelinedFlushMutations", opentracing.ChildOf(span.Context()))
diff --git a/txnkv/transaction/prewrite.go b/txnkv/transaction/prewrite.go
index c0e25a6..16fc05f 100644
--- a/txnkv/transaction/prewrite.go
+++ b/txnkv/transaction/prewrite.go
@@ -499,6 +499,10 @@ func (action actionPrewrite) handleSingleBatch(
 	}
 }
 
+func (actionPrewrite) isInterruptible() bool {
+	return true
+}
+
 func (c *twoPhaseCommitter) prewriteMutations(bo *retry.Backoffer, mutations CommitterMutations) error {
 	if span := opentracing.SpanFromContext(bo.GetCtx()); span != nil && span.Tracer() != nil {
 		span1 := span.Tracer().StartSpan("twoPhaseCommitter.prewriteMutations", opentracing.ChildOf(span.Context()))


Test patch:
diff --git a/integration_tests/2pc_test.go b/integration_tests/2pc_test.go
index 08de42e..9d215cf 100644
--- a/integration_tests/2pc_test.go
+++ b/integration_tests/2pc_test.go
@@ -2513,6 +2513,50 @@ func (s *testCommitterSuite) TestKillSignal() {
 	s.ErrorContains(err, "query interrupted")
 }
 
+func (s *testCommitterSuite) TestUninterruptibleAction() {
+	s.Run("Cleanup", func() {
+		var killed uint32 = 0
+		txn := s.begin()
+		txn.SetVars(kv.NewVariables(&killed))
+		err := txn.Set([]byte("k1"), []byte("v1"))
+		s.NoError(err)
+		committer, err := txn.NewCommitter(0)
+		s.NoError(err)
+		err = committer.PrewriteAllMutations(context.Background())
+		s.NoError(err)
+		atomic.StoreUint32(&killed, 2)
+		s.NoError(committer.CleanupMutations(context.Background()))
+	})
+	s.Run("PessimisticRollback", func() {
+		var killed uint32 = 0
+		txn := s.begin()
+		txn.SetVars(kv.NewVariables(&killed))
+		txn.SetPessimistic(true)
+		err := txn.LockKeys(context.Background(), kv.NewLockCtx(txn.StartTS(), kv.LockNoWait, time.Now()), []byte("k2"))
+		s.NoError(err)
+		atomic.StoreUint32(&killed, 2)
+		committer, err := txn.NewCommitter(0)
+		s.NoError(err)
+		s.NoError(committer.PessimisticRollbackMutations(context.Background(), committer.GetMutations()))
+	})
+	s.Run("Commit", func() {
+		var killed uint32 = 0
+		txn := s.begin()
+		txn.SetVars(kv.NewVariables(&killed))
+		err := txn.Set([]byte("k1"), []byte("v1"))
+		s.NoError(err)
+		committer, err := txn.NewCommitter(0)
+		s.NoError(err)
+		err = committer.PrewriteAllMutations(context.Background())
+		s.NoError(err)
+		atomic.StoreUint32(&killed, 2)
+		commitTS, err := s.store.GetOracle().GetTimestamp(context.Background(), &oracle.Option{})
+		s.NoError(err)
+		committer.SetCommitTS(commitTS)
+		s.NoError(committer.CommitMutations(context.Background()))
+	})
+}
+
 func (s *testCommitterSuite) Test2PCLifecycleHooks() {
 	reachedPre := atomic.Bool{}
 	reachedPost := atomic.Bool{}
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index 37a17d1..a50615f 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -2994,3 +2994,115 @@ func BenchmarkBatchLocateKeyRangesFromCache(t *testing.B) {
 	}
 	s.TearDownTest()
 }
+
+func (s *testRegionCacheSuite) TestRegionCacheValidAfterLoading() {
+	s.cache.clear()
+
+	// Split regions at "a", "b", "c", ..., "j"
+	regions := make([]uint64, 0, 11)
+	region1 := s.region1
+	regions = append(regions, region1)
+	for i := 0; i < 10; i++ {
+		region2 := s.cluster.AllocID()
+		newPeers := s.cluster.AllocIDs(2)
+		s.cluster.Split(region1, region2, []byte{'a' + byte(i)}, newPeers, newPeers[0])
+		region1 = region2
+		regions = append(regions, region1)
+	}
+
+	fns := []func(){
+		func() {
+			_, err := s.cache.LocateKeyRange(s.bo, []byte("a"), []byte("e"))
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.BatchLocateKeyRanges(s.bo, []kv.KeyRange{{StartKey: []byte("a"), EndKey: []byte("e")}})
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.LocateKey(s.bo, []byte("b"))
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.LocateEndKey(s.bo, []byte("c"))
+			s.Nil(err)
+		},
+		func() {
+			for _, regionID := range regions {
+				_, err := s.cache.LocateRegionByID(s.bo, regionID)
+				s.Nil(err)
+			}
+		},
+		func() {
+			_, _, err := s.cache.GroupKeysByRegion(s.bo, [][]byte{[]byte("a"), []byte("b"), []byte("c")}, nil)
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.ListRegionIDsInKeyRange(s.bo, []byte("a"), []byte("e"))
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.LoadRegionsInKeyRange(s.bo, []byte("a"), []byte("e"))
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.BatchLoadRegionsWithKeyRange(s.bo, []byte("a"), []byte("e"), 10)
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.BatchLoadRegionsWithKeyRanges(s.bo, []pd.KeyRange{{StartKey: []byte("a"), EndKey: []byte("e")}}, 10)
+			s.Nil(err)
+		},
+		func() {
+			_, err := s.cache.BatchLoadRegionsFromKey(s.bo, []byte("a"), 10)
+			s.Nil(err)
+		},
+	}
+
+	// Whether the region is loaded from PD(bypass the region cache) or from the region cache,
+	// the existing valid region should not be invalidated.
+	for _, fn := range fns {
+		loc, err := s.cache.LocateKey(s.bo, []byte("b"))
+		s.Nil(err)
+		region := s.cache.GetCachedRegionWithRLock(loc.Region)
+		fn()
+		s.True(region.isValid())
+	}
+
+	// If the region is invalidated already, it should be reloaded from PD and inserted into region cache anyway.
+	for _, fn := range fns {
+		loc, err := s.cache.LocateKey(s.bo, []byte("b"))
+		s.Nil(err)
+		region := s.cache.GetCachedRegionWithRLock(loc.Region)
+		region.invalidate(Other)
+		s.False(region.isValid())
+		fn()
+		s.False(region.isValid())
+		newLoc := s.cache.TryLocateKey([]byte("b"))
+		s.NotNil(newLoc)
+		region = s.cache.GetCachedRegionWithRLock(newLoc.Region)
+		s.True(region.isValid())
+	}
+}
+
+func (s *testRegionCacheSuite) TestBatchLoadLimitRanges() {
+	ranges := make([]kv.KeyRange, 0, 100000)
+	for i := 0; i < 100000; i++ {
+		startKey := make([]byte, 8)
+		endKey := make([]byte, 8)
+		binary.BigEndian.PutUint64(startKey, uint64(i*2))
+		binary.BigEndian.PutUint64(endKey, uint64(i*2+1))
+		ranges = append(ranges, kv.KeyRange{StartKey: startKey, EndKey: endKey})
+	}
+
+	originalBatchScanRegions := s.cache.pdClient.BatchScanRegions
+	s.cache.pdClient = &inspectedPDClient{
+		Client: s.cache.pdClient,
+		batchScanRegions: func(ctx context.Context, keyRanges []pd.KeyRange, limit int, opts ...pd.GetRegionOption) ([]*pd.Region, error) {
+			s.LessOrEqual(len(keyRanges), 16*defaultRegionsPerBatch)
+			return originalBatchScanRegions(ctx, keyRanges, limit, opts...)
+		},
+	}
+	_, err := s.cache.BatchLocateKeyRanges(s.bo, ranges)
+	s.Nil(err)
+}