Repository: tikv/client-go
Pull request #1364 (2024-06-12 07:31:46)

txn-file: Handle locked keys in split region
### Changes

- Handle locked keys by resolve locks in split region



Patch:
diff --git a/examples/gcworker/go.mod b/examples/gcworker/go.mod
index a47f527..cf28cd2 100644
--- a/examples/gcworker/go.mod
+++ b/examples/gcworker/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/rawkv/go.mod b/examples/rawkv/go.mod
index 7945319..b2289ef 100644
--- a/examples/rawkv/go.mod
+++ b/examples/rawkv/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/1pc_txn/go.mod b/examples/txnkv/1pc_txn/go.mod
index 3820a60..8fc447d 100644
--- a/examples/txnkv/1pc_txn/go.mod
+++ b/examples/txnkv/1pc_txn/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/async_commit/go.mod b/examples/txnkv/async_commit/go.mod
index 6b53262..c372be4 100644
--- a/examples/txnkv/async_commit/go.mod
+++ b/examples/txnkv/async_commit/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/delete_range/go.mod b/examples/txnkv/delete_range/go.mod
index c486ef8..d74d356 100644
--- a/examples/txnkv/delete_range/go.mod
+++ b/examples/txnkv/delete_range/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/go.mod b/examples/txnkv/go.mod
index 6babdd6..e1f00a0 100644
--- a/examples/txnkv/go.mod
+++ b/examples/txnkv/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/pessimistic_txn/go.mod b/examples/txnkv/pessimistic_txn/go.mod
index 41632b9..a00ba12 100644
--- a/examples/txnkv/pessimistic_txn/go.mod
+++ b/examples/txnkv/pessimistic_txn/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/examples/txnkv/unsafedestoryrange/go.mod b/examples/txnkv/unsafedestoryrange/go.mod
index d704071..a36072e 100644
--- a/examples/txnkv/unsafedestoryrange/go.mod
+++ b/examples/txnkv/unsafedestoryrange/go.mod
@@ -20,7 +20,7 @@ require (
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c // indirect
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c // indirect
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 // indirect
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 // indirect
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/prometheus/client_golang v1.18.0 // indirect
diff --git a/go.mod b/go.mod
index 29cb803..6c97c65 100644
--- a/go.mod
+++ b/go.mod
@@ -14,7 +14,7 @@ require (
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.18.0
diff --git a/go.sum b/go.sum
index 21804c4..24809c1 100644
--- a/go.sum
+++ b/go.sum
@@ -74,8 +74,8 @@ github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c h1:CgbKAHto5CQgW
 github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c/go.mod h1:4qGtCB0QK0wBzKtFEGDhxXnSnbQApw1gc9siScUl8ew=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
-github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 h1:tAfzWZi+hVqX8tM+jnltSiC2k3X+305L/aGiblhX6dk=
-github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 h1:6IiXyf4QaivsArGCl19AqhdC2gp0LYx0s0QZI95S96M=
+github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/internal/apicodec/codec_v2.go b/internal/apicodec/codec_v2.go
index d4a89dc..15642fb 100644
--- a/internal/apicodec/codec_v2.go
+++ b/internal/apicodec/codec_v2.go
@@ -598,6 +598,10 @@ func (c *codecV2) DecodeResponse(req *tikvrpc.Request, resp *tikvrpc.Response) (
 		if err != nil {
 			return nil, err
 		}
+		r.Errors, err = c.decodeKeyErrors(r.Errors)
+		if err != nil {
+			return nil, err
+		}
 		r.Regions, err = c.decodeRegions(r.Regions)
 		if err != nil {
 			return nil, err
diff --git a/tikv/split_region.go b/tikv/split_region.go
index 2844b38..20162f5 100644
--- a/tikv/split_region.go
+++ b/tikv/split_region.go
@@ -53,6 +53,7 @@ import (
 	"github.com/tikv/client-go/v2/kv"
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"github.com/tikv/client-go/v2/txnkv/rangetask"
+	"github.com/tikv/client-go/v2/txnkv/txnlock"
 	"github.com/tikv/client-go/v2/util"
 	pd "github.com/tikv/pd/client"
 	"go.uber.org/zap"
@@ -174,6 +175,20 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte
 	}
 
 	spResp := resp.Resp.(*kvrpcpb.SplitRegionResponse)
+
+	keyErrs := spResp.GetErrors()
+	if len(keyErrs) > 0 {
+		err := s.handleSplitRegionKeyErrors(bo, keyErrs)
+		if err != nil {
+			batchResp.Error = err
+			return batchResp
+		}
+		resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID)
+		batchResp.Response = resp
+		batchResp.Error = err
+		return batchResp
+	}
+
 	regions := spResp.GetRegions()
 	if len(regions) > 0 {
 		// Divide a region into n, one of them may not need to be scattered,
@@ -218,6 +233,46 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte
 	return batchResp
 }
 
+func (s *KVStore) handleSplitRegionKeyErrors(bo *Backoffer, keyErrs []*kvrpcpb.KeyError) error {
+	var (
+		locks   []*txnlock.Lock
+		startTS uint64 = math.MaxUint64 // Set as MaxUint64 and check txn status will not push the minCommiTS.
+	)
+	for _, keyErr := range keyErrs {
+		lock, err1 := txnlock.ExtractLockFromKeyErr(keyErr)
+		if err1 != nil {
+			// Split region should return key error of locked only.
+			return err1
+		}
+		logutil.Logger(bo.GetCtx()).Info("split region encounters lock", zap.Stringer("lock", lock))
+		locks = append(locks, lock)
+	}
+
+	token := s.GetLockResolver().RecordResolvingLocks(locks, startTS)
+	defer s.GetLockResolver().ResolveLocksDone(startTS, token)
+
+	resolveLockOpts := txnlock.ResolveLocksOptions{
+		CallerStartTS: startTS,
+		Locks:         locks,
+	}
+	resolveLockRes, err := s.GetLockResolver().ResolveLocksWithOpts(bo, resolveLockOpts)
+	if err != nil {
+		return errors.WithStack(err)
+	}
+	msBeforeExpired := resolveLockRes.TTL
+	if msBeforeExpired > 0 {
+		err = bo.BackoffWithCfgAndMaxSleep(
+			retry.BoTxnLock,
+			int(msBeforeExpired),
+			errors.Errorf("split region lockedKeys: %d", len(locks)),
+		)
+		if err != nil {
+			return errors.WithStack(err)
+		}
+	}
+	return nil
+}
+
 const (
 	splitRegionBackoff     = 20000
 	maxSplitRegionsBackoff = 120000


Test patch:
diff --git a/integration_tests/go.mod b/integration_tests/go.mod
index 6c3a255..2d9ab86 100644
--- a/integration_tests/go.mod
+++ b/integration_tests/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/ninedraft/israce v0.0.3
 	github.com/pingcap/errors v0.11.5-0.20231212100244-799fae176cfb
 	github.com/pingcap/failpoint v0.0.0-20220801062533-2eaa32854a6c
-	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9
+	github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9
 	github.com/pingcap/tidb v1.1.0-beta.0.20231025055903-f3dfd896b5e6
 	github.com/pkg/errors v0.9.1
 	github.com/stretchr/testify v1.8.4
diff --git a/integration_tests/go.sum b/integration_tests/go.sum
index b17b4fd..c8be2da 100644
--- a/integration_tests/go.sum
+++ b/integration_tests/go.sum
@@ -416,8 +416,8 @@ github.com/pingcap/fn v0.0.0-20200306044125-d5540d389059/go.mod h1:fMRU1BA1y+r89
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989 h1:surzm05a8C9dN8dIUmo4Be2+pMRb6f55i+UIYrluu2E=
 github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17XtbryoCJhkKGbT62+L2OlrniwqiGLSqrmdHCMzZw=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9 h1:tAfzWZi+hVqX8tM+jnltSiC2k3X+305L/aGiblhX6dk=
-github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9 h1:6IiXyf4QaivsArGCl19AqhdC2gp0LYx0s0QZI95S96M=
+github.com/pingcap/kvproto v0.0.0-20240613021635-bca367ddc3b9/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.0/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=
 github.com/pingcap/log v1.1.1-0.20230317032135-a0d097d16e22 h1:2SOzvGvE8beiC1Y4g9Onkvu6UmuBBOeWRGQEjJaT/JY=
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index b82e927..99ceea2 100644
--- a/internal/locate/region_request_test.go
+++ b/internal/locate/region_request_test.go
@@ -533,6 +533,10 @@ func (s *mockTikvGrpcServer) CancelDisaggTask(context.Context, *disaggregated.Ca
 	return nil, errors.New("unreachable")
 }
 
+func (s *mockTikvGrpcServer) GetHealthFeedback(ctx context.Context, request *kvrpcpb.GetHealthFeedbackRequest) (*kvrpcpb.GetHealthFeedbackResponse, error) {
+	return nil, errors.New("unreachable")
+}
+
 func (s *testRegionRequestToSingleStoreSuite) TestNoReloadRegionForGrpcWhenCtxCanceled() {
 	// prepare a mock tikv grpc server
 	addr := "localhost:56341"