Repository: tikv/client-go
Pull request #1694 (2025-07-03 02:14:02)

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


zyguan: /retest

Patch:
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index 126f4bf..5038c97 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -1408,7 +1408,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)
+		c.insertRegionToCache(region, false)
 	}
 
 	return
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 1bd97fa..14bb050 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -1679,8 +1679,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 d1138ac..7a6d4b7 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -160,7 +160,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})
 
@@ -188,7 +188,7 @@ 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})
 
diff --git a/tikvrpc/tikvrpc.go b/tikvrpc/tikvrpc.go
index 9d8cab6..881b329 100644
--- a/tikvrpc/tikvrpc.go
+++ b/tikvrpc/tikvrpc.go
@@ -326,6 +326,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 b7471f8..04fd629 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
 }
 
@@ -1055,9 +1056,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 b080319..92e87f7 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 3818e39..46a0e97 100644
--- a/txnkv/transaction/commit.go
+++ b/txnkv/transaction/commit.go
@@ -221,6 +221,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 5855a08..c621b83 100644
--- a/txnkv/transaction/pessimistic.go
+++ b/txnkv/transaction/pessimistic.go
@@ -535,6 +535,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 {
@@ -570,6 +574,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/prewrite.go b/txnkv/transaction/prewrite.go
index d74b5fe..7c1742d 100644
--- a/txnkv/transaction/prewrite.go
+++ b/txnkv/transaction/prewrite.go
@@ -493,6 +493,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 9369bfb..01faf54 100644
--- a/integration_tests/2pc_test.go
+++ b/integration_tests/2pc_test.go
@@ -2512,3 +2512,47 @@ func (s *testCommitterSuite) TestKillSignal() {
 	err = txn.Commit(context.Background())
 	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()))
+	})
+}
diff --git a/internal/locate/region_cache_test.go b/internal/locate/region_cache_test.go
index 9c1aff6..54d485d 100644
--- a/internal/locate/region_cache_test.go
+++ b/internal/locate/region_cache_test.go
@@ -2025,3 +2025,81 @@ func (s *testRegionCacheSuite) TestIssue1401() {
 	}, 3*time.Second, time.Second)
 	s.Require().True(isStoreContainLabel(newStore1.labels, "host", "0.0.0.0:20161"))
 }
+
+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.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.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())
+	}
+}