txn-file: Use LocateKey to get overlapping regions for txn chunks

examples/gcworker/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/rawkv/go.mod

@@ -2,7 +2,10 @@ module rawkv
 
 go 1.21
 
-require github.com/tikv/client-go/v2 v2.0.0
+require (
+	github.com/pingcap/kvproto v0.0.0-20240314010430-d552d2f77da9
+	github.com/tikv/client-go/v2 v2.0.0
+)
 
 require (
 	github.com/beorn7/perks v1.0.1 // indirect
@@ -20,7 +23,6 @@ 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/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
@@ -29,7 +31,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/1pc_txn/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/async_commit/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/delete_range/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/pessimistic_txn/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

examples/txnkv/unsafedestoryrange/go.mod

@@ -29,7 +29,7 @@ require (
 	github.com/prometheus/procfs v0.12.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a // indirect
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 // indirect
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 // indirect
 	github.com/twmb/murmur3 v1.1.3 // indirect
 	go.etcd.io/etcd/api/v3 v3.5.10 // indirect
 	go.etcd.io/etcd/client/pkg/v3 v3.5.10 // indirect

go.mod

@@ -21,7 +21,7 @@ require (
 	github.com/prometheus/client_model v0.5.0
 	github.com/stretchr/testify v1.8.2
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877
 	github.com/twmb/murmur3 v1.1.3
 	go.etcd.io/etcd/api/v3 v3.5.10
 	go.etcd.io/etcd/client/v3 v3.5.10

go.sum

@@ -111,8 +111,8 @@ github.com/stretchr/testify v1.8.2 h1:+h33VjcLVPDHtOdpUCuF+7gSuG3yGIftsP1YvFihtJ
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 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/pd/client v0.0.0-20240407053619-a173a08acd81 h1:MCfYSh3oZIkHjXf+H65LI1abtdnRh6zdNYPwYDSodms=
-github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81/go.mod h1:Qdzax77EzM8h6NbzwKymp96AB0DUvtkHCcC1jwxoiGU=
+github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 h1:8Poo/LLJo9sMO2Yak5+0yWEiN04NenJvvBGBafSERGg=
+github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877/go.mod h1:Qdzax77EzM8h6NbzwKymp96AB0DUvtkHCcC1jwxoiGU=
 github.com/twmb/murmur3 v1.1.3 h1:D83U0XYKcHRYwYIpBKf3Pks91Z0Byda/9SJ8B6EMRcA=
 github.com/twmb/murmur3 v1.1.3/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=

integration_tests/go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/stretchr/testify v1.8.4
 	github.com/tidwall/gjson v1.14.1
 	github.com/tikv/client-go/v2 v2.0.8-0.20240304062027-038a52519d39
-	github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81
+	github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877
 	go.uber.org/goleak v1.3.0
 )
 

integration_tests/go.sum

@@ -530,8 +530,8 @@ github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
 github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
 github.com/tikv/pd v0.0.0-20230420042919-0e0313adef1f h1:1BBlbLSCua+oK3xngEjPO2dj1oLcItJxhCzqYG6XANA=
 github.com/tikv/pd v0.0.0-20230420042919-0e0313adef1f/go.mod h1:2hGdvUNGLaCKLbEXC8JMiQz2JsbxEDCKGDOBGJTt0NQ=
-github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81 h1:MCfYSh3oZIkHjXf+H65LI1abtdnRh6zdNYPwYDSodms=
-github.com/tikv/pd/client v0.0.0-20240407053619-a173a08acd81/go.mod h1:Qdzax77EzM8h6NbzwKymp96AB0DUvtkHCcC1jwxoiGU=
+github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877 h1:8Poo/LLJo9sMO2Yak5+0yWEiN04NenJvvBGBafSERGg=
+github.com/tikv/pd/client v0.0.0-20240410074522-b9240a065877/go.mod h1:Qdzax77EzM8h6NbzwKymp96AB0DUvtkHCcC1jwxoiGU=
 github.com/tklauser/go-sysconf v0.3.9/go.mod h1:11DU/5sG7UexIrp/O6g35hrWzu0JxlwQ3LSFUzyeuhs=
 github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
 github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=

txnkv/transaction/txn_file.go

@@ -23,6 +23,7 @@ import (
 	"hash/crc32"
 	"io"
 	"net/http"
+	"strings"
 	"time"
 
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
@@ -54,15 +55,41 @@ type txnFileCtx struct {
 
 type chunkBatch struct {
 	txnChunkSlice
-	region    *locate.Region
+	region    *locate.KeyLocation
 	isPrimary bool
 }
 
+func (b chunkBatch) String() string {
+	return fmt.Sprintf("chunkBatch{region: %s, isPrimary: %t, txnChunkSlice: %s}", b.region.String(), b.isPrimary, b.txnChunkSlice.String())
+}
+
 type txnChunkSlice struct {
 	chunkIDs    []uint64
 	chunkRanges []txnChunkRange
 }
 
+func (s txnChunkSlice) String() string {
+	slice := make([]string, len(s.chunkRanges))
+	for i, ran := range s.chunkRanges {
+		slice[i] = fmt.Sprintf("txnChunkSlice{%v: [%s, %s]}", s.chunkIDs[i], kv.StrKey(ran.smallest), kv.StrKey(ran.biggest))
+	}
+	return fmt.Sprintf("[%s]", strings.Join(slice, ", "))
+}
+
+func (s *txnChunkSlice) Smallest() []byte {
+	if len(s.chunkRanges) == 0 {
+		return nil
+	}
+	return s.chunkRanges[0].smallest
+}
+
+func (s *txnChunkSlice) Biggest() []byte {
+	if len(s.chunkRanges) == 0 {
+		return nil
+	}
+	return s.chunkRanges[len(s.chunkRanges)-1].biggest
+}
+
 func (cs *txnChunkSlice) appendSlice(other *txnChunkSlice) {
 	if cs.len() == 0 {
 		cs.chunkIDs = append(cs.chunkIDs, other.chunkIDs...)
@@ -89,25 +116,38 @@ func (cs *txnChunkSlice) len() int {
 }
 
 func (cs *txnChunkSlice) groupToBatches(c *locate.RegionCache, bo *retry.Backoffer) ([]chunkBatch, error) {
-	var batches []chunkBatch
-	endKey := kv.NextKey(cs.chunkRanges[cs.len()-1].biggest)
-	regions, err := c.LoadRegionsInKeyRange(bo, cs.chunkRanges[0].smallest, endKey)
-	if err != nil {
-		return nil, err
-	}
-	for _, region := range regions {
-		batch := chunkBatch{
-			region: region,
+	batch_map := make(map[locate.RegionVerID]*chunkBatch)
+	smallest := cs.Smallest()
+
+	for i, chunkRange := range cs.chunkRanges {
+		regions, err := chunkRange.getOverlapRegions(c, bo)
+		if err != nil {
+			return nil, errors.WithStack(err)
 		}
-		for i, chunkRange := range cs.chunkRanges {
-			if chunkRange.overlapRegion(region) {
-				batch.append(cs.chunkIDs[i], chunkRange)
+
+		for _, r := range regions {
+			if batch_map[r.Region] == nil {
+				batch_map[r.Region] = &chunkBatch{
+					region: r,
+				}
 			}
+			batch_map[r.Region].append(cs.chunkIDs[i], chunkRange)
 		}
-		if batch.len() > 0 {
-			batches = append(batches, batch)
+	}
+
+	batches := make([]chunkBatch, 0, len(batch_map))
+	picked := false // Pick the batch with primary and put it at the first.
+	for _, batch := range batch_map {
+		if !picked && bytes.Equal(batch.Smallest(), smallest) {
+			if len(batches) > 0 {
+				batches[0], *batch = *batch, batches[0]
+			}
+			picked = true
 		}
+
+		batches = append(batches, *batch)
 	}
+	logutil.Logger(bo.GetCtx()).Debug("txn file group to batches", zap.Stringers("batches", batches))
 	return batches, nil
 }
 
@@ -123,8 +163,21 @@ func newTxnChunkRange(smallest []byte, biggest []byte) txnChunkRange {
 	}
 }
 
-func (r *txnChunkRange) overlapRegion(region *locate.Region) bool {
-	return bytes.Compare(r.smallest, region.EndKey()) < 0 && bytes.Compare(r.biggest, region.StartKey()) >= 0
+func (r *txnChunkRange) getOverlapRegions(c *locate.RegionCache, bo *retry.Backoffer) ([]*locate.KeyLocation, error) {
+	regions := make([]*locate.KeyLocation, 0)
+	startKey := r.smallest
+	for bytes.Compare(startKey, r.biggest) <= 0 {
+		loc, err := c.LocateKey(bo, startKey)
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		regions = append(regions, loc)
+		if len(loc.EndKey) == 0 {
+			break
+		}
+		startKey = loc.EndKey
+	}
+	return regions, nil
 }
 
 type txnFileAction interface {
@@ -159,7 +212,7 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back
 	sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient())
 	var resolvingRecordToken *int
 	for {
-		resp, _, err := sender.SendReq(bo, req, batch.region.VerID(), client.ReadTimeoutMedium)
+		resp, _, err := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutMedium)
 		if err != nil {
 			return nil, err
 		}
@@ -272,7 +325,7 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof
 	})
 	sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient())
 	for {
-		resp, _, err := sender.SendReq(bo, req, batch.region.VerID(), client.ReadTimeoutMedium)
+		resp, _, err := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutMedium)
 		if batch.isPrimary && sender.GetRPCError() != nil {
 			c.setUndeterminedErr(errors.WithStack(sender.GetRPCError()))
 		}
@@ -349,7 +402,7 @@ func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back
 		},
 	})
 	sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient())
-	resp, _, err1 := sender.SendReq(bo, req, batch.region.VerID(), client.ReadTimeoutShort)
+	resp, _, err1 := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutShort)
 	if err1 != nil {
 		return nil, err1
 	}