Repository: tikv/client-go
Pull request #805 (2023-05-17 10:19:13)

Count only one tiflash replica in disaggregated mode.
Ref: pingcap/tiflash#7491

In disaggregated-mode(tiflash write-node), we only count 1 replica for tiflash to calculate 'WRU', no matter how many tiflash write-nodes there are.


Patch:
diff --git a/go.mod b/go.mod
index 8ad4ad5..f8ecbf5 100644
--- a/go.mod
+++ b/go.mod
@@ -14,29 +14,29 @@ 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-20230511011722-6e0e8a7deaa1
+	github.com/pingcap/kvproto v0.0.0-20230518070549-2a7e63366ffa
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pkg/errors v0.9.1
-	github.com/prometheus/client_golang v1.14.0
-	github.com/prometheus/client_model v0.3.0
+	github.com/prometheus/client_golang v1.15.1
+	github.com/prometheus/client_model v0.4.0
 	github.com/sony/gobreaker v0.5.0
 	github.com/stretchr/testify v1.8.2
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a
 	github.com/tikv/pd v0.0.0-20230420042919-0e0313adef1f
-	github.com/tikv/pd/client v0.0.0-20230420042919-0e0313adef1f
+	github.com/tikv/pd/client v0.0.0-20230518085737-f4241b0c82b0
 	github.com/twmb/murmur3 v1.1.3
 	go.etcd.io/etcd/api/v3 v3.5.2
 	go.etcd.io/etcd/client/v3 v3.5.2
-	go.uber.org/atomic v1.10.0
+	go.uber.org/atomic v1.11.0
 	go.uber.org/goleak v1.1.12
 	go.uber.org/zap v1.24.0
 	golang.org/x/sync v0.2.0
-	google.golang.org/grpc v1.54.0
+	google.golang.org/grpc v1.55.0
 )
 
 require (
 	github.com/BurntSushi/toml v1.2.1 // indirect
-	github.com/benbjohnson/clock v1.3.0 // indirect
+	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
 	github.com/cespare/xxhash/v2 v2.2.0 // indirect
 	github.com/coreos/go-semver v0.3.0 // indirect
@@ -48,14 +48,13 @@ require (
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/golang/snappy v0.0.4 // indirect
 	github.com/jonboulle/clockwork v0.4.0 // indirect
-	github.com/kr/pretty v0.2.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
 	github.com/onsi/ginkgo v1.16.5 // indirect
 	github.com/petermattis/goid v0.0.0-20211229010228-4d14c490ee36 // indirect
 	github.com/phf/go-queue v0.0.0-20170504031614-9abe38d0371d // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/common v0.39.0 // indirect
+	github.com/prometheus/common v0.43.0 // indirect
 	github.com/prometheus/procfs v0.9.0 // indirect
 	github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 // indirect
 	github.com/sasha-s/go-deadlock v0.2.0 // indirect
diff --git a/go.sum b/go.sum
index bcf2847..7140684 100644
--- a/go.sum
+++ b/go.sum
@@ -10,8 +10,8 @@ github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRF
 github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
 github.com/antihax/optional v1.0.0/go.mod h1:uupD/76wgC+ih3iEmQUL+0Ugr19nfwCT1kdvxnR2qWY=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
-github.com/benbjohnson/clock v1.3.0 h1:ip6w0uFQkncKQ979AypyG0ER7mqUSBdKLOgAle/AT8A=
-github.com/benbjohnson/clock v1.3.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
+github.com/benbjohnson/clock v1.3.5 h1:VvXlSJBzZpA/zum6Sj74hxwYI2DIxRWuNIoXAzHZz5o=
+github.com/benbjohnson/clock v1.3.5/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v0.0.0-20180321164747-3a771d992973/go.mod h1:Dwedo/Wpr24TaqPxmxbtue+5NUziq4I4S80YR8gNf3Q=
 github.com/beorn7/perks v1.0.0/go.mod h1:KWe93zE9D1o94FZ5RNwFwVgaQK1VOXiVxmqh+CedLV8=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
@@ -76,7 +76,6 @@ github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5y
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.3/go.mod h1:vzj43D7+SQXF/4pzW/hwtAqwc6iTitCiVSaWz5lYuqw=
-github.com/golang/protobuf v1.3.5/go.mod h1:6O5/vntMXwX2lRkT1hjjk0nAC1IDOTvTlVgjlRvqsdk=
 github.com/golang/protobuf v1.4.0-rc.1/go.mod h1:ceaxUfeHdC40wWswd/P6IGgMaK3YpKi5j83Wpe3EHw8=
 github.com/golang/protobuf v1.4.0-rc.1.0.20200221234624-67d41d38c208/go.mod h1:xKAWHe0F5eneWXFV3EuXVDTCmh+JuBKY0li0aMyXATA=
 github.com/golang/protobuf v1.4.0-rc.2/go.mod h1:LlEzMj4AhA7rCAGe4KMBDvJI+AwstrUpVNzEA03Pprs=
@@ -130,8 +129,7 @@ github.com/konsorten/go-windows-terminal-sequences v1.0.3 h1:CE8S1cTafDpPvMhIxNJ
 github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 github.com/kr/logfmt v0.0.0-20140226030751-b84e30acd515/go.mod h1:+0opPa2QZZtGFBFZlji/RkVcI2GknAs/DXo4wKdlNEc=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
-github.com/kr/pretty v0.2.1 h1:Fmg33tUaq4/8ym9TJN1x7sLJnHVwhP33CNkpYV/7rwI=
-github.com/kr/pretty v0.2.1/go.mod h1:ipq/a2n7PKx3OHsz4KJII5eveXtPO4qwEXGdVfWzfnI=
+github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
 github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
@@ -172,8 +170,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-20230511011722-6e0e8a7deaa1 h1:VXQ6Du/nKZ9IQnI9NWMzKbftWu8NV5pQkSLKIRzzGN4=
-github.com/pingcap/kvproto v0.0.0-20230511011722-6e0e8a7deaa1/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
+github.com/pingcap/kvproto v0.0.0-20230518070549-2a7e63366ffa h1:r7A9U1xjhTK26d6h+J1N5OR+d5Y2conQgi5qElNcKNY=
+github.com/pingcap/kvproto v0.0.0-20230518070549-2a7e63366ffa/go.mod h1:guCyM5N+o+ru0TsoZ1hi9lDjUMs2sIBjW3ARTEpVbnk=
 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.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
@@ -186,19 +184,19 @@ github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXP
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
 github.com/prometheus/client_golang v1.7.1/go.mod h1:PY5Wy2awLA44sXw4AOSfFBetzPP4j5+D6mVACh+pe2M=
 github.com/prometheus/client_golang v1.11.0/go.mod h1:Z6t4BnS23TR94PD6BsDNk8yVqroYurpAkEiz0P2BEV0=
-github.com/prometheus/client_golang v1.14.0 h1:nJdhIvne2eSX/XRAFV9PcvFFRbrjbcTUj0VP62TMhnw=
-github.com/prometheus/client_golang v1.14.0/go.mod h1:8vpkKitgIVNcqrRBWh1C4TIUQgYNtG/XQE4E/Zae36Y=
+github.com/prometheus/client_golang v1.15.1 h1:8tXpTmJbyH5lydzFPoxSIJ0J46jdh3tylbvM1xCv0LI=
+github.com/prometheus/client_golang v1.15.1/go.mod h1:e9yaBhRPU2pPNsZwE+JdQl0KEt1N9XgF6zxWmaC0xOk=
 github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
 github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
 github.com/prometheus/client_model v0.2.0/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
-github.com/prometheus/client_model v0.3.0 h1:UBgGFHqYdG/TPFD1B1ogZywDqEkwp3fBMvqdiQ7Xew4=
-github.com/prometheus/client_model v0.3.0/go.mod h1:LDGWKZIo7rky3hgvBe+caln+Dr3dPggB5dvjtD7w9+w=
+github.com/prometheus/client_model v0.4.0 h1:5lQXD3cAg1OXBf4Wq03gTrXHeaV0TQvGfUooCfx1yqY=
+github.com/prometheus/client_model v0.4.0/go.mod h1:oMQmHW1/JoDwqLtg57MGgP/Fb1CJEYF2imWWhWtMkYU=
 github.com/prometheus/common v0.4.1/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
 github.com/prometheus/common v0.10.0/go.mod h1:Tlit/dnDKsSWFlCLTWaA1cyBgKHSMdTB80sz/V91rCo=
 github.com/prometheus/common v0.26.0/go.mod h1:M7rCNAaPfAosfx8veZJCuw84e35h3Cfd9VFqTh1DIvc=
-github.com/prometheus/common v0.39.0 h1:oOyhkDq05hPZKItWVBkJ6g6AtGxi+fy7F4JvUV8uhsI=
-github.com/prometheus/common v0.39.0/go.mod h1:6XBZ7lYdLCbkAVhwRsWTZn+IN5AB9F/NXd5w0BbEX0Y=
+github.com/prometheus/common v0.43.0 h1:iq+BVjvYLei5f27wiuNiB1DN6DYQkp1c8Bx0Vykh5us=
+github.com/prometheus/common v0.43.0/go.mod h1:NCvr5cQIh3Y/gy73/RdVtC9r8xxrxwJnB+2lB3BxrFc=
 github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
 github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
 github.com/prometheus/procfs v0.1.3/go.mod h1:lV6e/gmhEcM9IjHGsFOCxxuZ+z1YqCvr4OA4YeYWdaU=
@@ -208,6 +206,7 @@ github.com/prometheus/procfs v0.9.0/go.mod h1:+pB4zwohETzFnmlpe6yd2lSc+0/46IYZRB
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0 h1:OdAsTTz6OkFY5QxjkYwrChwuRruF69c169dPK26NUlk=
 github.com/remyoudompheng/bigfft v0.0.0-20200410134404-eec4a21b6bb0/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo=
 github.com/rogpeppe/fastuuid v1.2.0/go.mod h1:jVj6XXZzXRy/MSR5jhDC/2q6DgLz+nrA6LYCDYWNEvQ=
+github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ=
 github.com/sasha-s/go-deadlock v0.2.0 h1:lMqc+fUb7RrFS3gQLtoQsJ7/6TV/pAIFvBsqX73DK8Y=
 github.com/sasha-s/go-deadlock v0.2.0/go.mod h1:StQn567HiB1fF2yJ44N9au7wOhrPS3iZqiDbRupzT10=
 github.com/sergi/go-diff v1.1.0/go.mod h1:STckp+ISIX8hZLjrqAeVduY0gWCT9IjLuqbuNXdaHfM=
@@ -237,8 +236,8 @@ github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a h1:J/YdBZ46WKpXsxsW
 github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a/go.mod h1:h4xBhSNtOeEosLJ4P7JyKXX7Cabg7AVkWCK5gV2vOrM=
 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-20230420042919-0e0313adef1f h1:Qa11LxEMz3tPlkcEJCnkc+phV00AaqXNNpqybx+QX04=
-github.com/tikv/pd/client v0.0.0-20230420042919-0e0313adef1f/go.mod h1:5vgcvO020ZCdMZkTrRdS/wFZQUab82BSfKE38T61ro0=
+github.com/tikv/pd/client v0.0.0-20230518085737-f4241b0c82b0 h1:3CCUDAeumhOwVYLlMd8tuUGTn4FsZ6jyH/2qjPbaAxk=
+github.com/tikv/pd/client v0.0.0-20230518085737-f4241b0c82b0/go.mod h1:mpFdCB3vLT/8wP0xEtDe5aSqEtjwiSqI1hyqXgqDV40=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20200427203606-3cfed13b9966 h1:j6JEOq5QWFker+d7mFQYOhjTZonQ7YkLTHm56dbn+yM=
 github.com/twmb/murmur3 v1.1.3 h1:D83U0XYKcHRYwYIpBKf3Pks91Z0Byda/9SJ8B6EMRcA=
 github.com/twmb/murmur3 v1.1.3/go.mod h1:Qq/R7NUyOfr65zD+6Q5IHKsJLwP7exErjN6lyyq3OSQ=
@@ -258,8 +257,8 @@ go.uber.org/atomic v1.4.0/go.mod h1:gD2HeocX3+yG+ygLZcrzQJaqmWj9AIm7n08wl/qW/PE=
 go.uber.org/atomic v1.6.0/go.mod h1:sABNBOSYdrvTF6hTgEIbc7YasKWGhgEQZyfxyTvoXHQ=
 go.uber.org/atomic v1.7.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
-go.uber.org/atomic v1.10.0 h1:9qC72Qh0+3MqyJbAn8YU5xVq1frD8bn3JtD2oXtafVQ=
-go.uber.org/atomic v1.10.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
+go.uber.org/atomic v1.11.0 h1:ZvwS0R+56ePWxUNi+Atn9dWONBPp/AUETXlHW0DxSjE=
+go.uber.org/atomic v1.11.0/go.mod h1:LUxbIzbOniOlMKjJjyPfpl4v+PKK2cNJn91OQbhoJI0=
 go.uber.org/goleak v1.1.10/go.mod h1:8a7PlsEVH3e/a/GLqe5IIrQx6GzcnRmZEufDUTk4A7A=
 go.uber.org/goleak v1.1.12 h1:gZAh5/EyT/HQwlpkCy6wTpqfH9H8Lz8zbm3dZh+OyzA=
 go.uber.org/goleak v1.1.12/go.mod h1:cwTWslyiVhfpKIDGSZEM2HlOvcqm+tG4zioyIeLoqMQ=
@@ -393,8 +392,8 @@ google.golang.org/grpc v1.25.1/go.mod h1:c3i+UQWmh7LiEpx4sFZnkU36qjEYZ0imhYfXVyQ
 google.golang.org/grpc v1.27.0/go.mod h1:qbnxyOmOxrQa7FizSgH+ReBfzJrCY1pSN7KXBS8abTk=
 google.golang.org/grpc v1.33.1/go.mod h1:fr5YgcSWrqhRRxogOsw7RzIpsmvOZ6IcH4kBYTpR3n0=
 google.golang.org/grpc v1.38.0/go.mod h1:NREThFqKR1f3iQ6oBuvc5LadQuXVGo9rkm5ZGrQdJfM=
-google.golang.org/grpc v1.54.0 h1:EhTqbhiYeixwWQtAEZAxmV9MGqcjEU2mFx52xCzNyag=
-google.golang.org/grpc v1.54.0/go.mod h1:PUSEXI6iWghWaB6lXM4knEgpJNu2qUcKfDtNci3EC2g=
+google.golang.org/grpc v1.55.0 h1:3Oj82/tFSCeUrRTg/5E/7d/W5A1tj6Ky1ABAuZuv5ag=
+google.golang.org/grpc v1.55.0/go.mod h1:iYEXKGkEBhg1PjZQvoYEVPTDkHo1/bjTnfwTeGONTY8=
 google.golang.org/protobuf v0.0.0-20200109180630-ec00e32a8dfd/go.mod h1:DFci5gLYBciE7Vtevhsrf46CRTquxDuWsQurQQe4oz8=
 google.golang.org/protobuf v0.0.0-20200221191635-4d8936d0db64/go.mod h1:kwYJMbMJ01Woi6D6+Kah6886xMZcty6N08ah7+eCXa0=
 google.golang.org/protobuf v0.0.0-20200228230310-ab0ca4ff8a60/go.mod h1:cfTl7dwQJ+fmap5saPgwCLgHXTUD7jkjRqWcaiX5VyM=
diff --git a/internal/locate/region_request.go b/internal/locate/region_request.go
index 52d81b3..942ab91 100644
--- a/internal/locate/region_request.go
+++ b/internal/locate/region_request.go
@@ -1313,6 +1313,32 @@ func fetchRespInfo(resp *tikvrpc.Response) string {
 	return extraInfo
 }
 
+func (s *RegionRequestSender) countReplicaNumber(peers []*metapb.Peer) int {
+	isTiFlashWriteNode := func(storeId uint64) bool {
+		store := s.regionCache.getStoreByStoreID(storeId)
+		engine, _ := store.GetLabelValue("engine")
+		engineRole, _ := store.GetLabelValue("engine_role")
+		return engine == "tiflash" && engineRole == "write"
+	}
+	// In disaggregated-mode(tiflash write-node), only count 1 replica for tiflash, no matter how many tiflash write-nodes there are.
+	replicaNumber := 0
+	hasTiFlashWriteNode := false
+	for _, peer := range peers {
+		role := peer.GetRole()
+		if role == metapb.PeerRole_Voter {
+			replicaNumber++
+		} else if role == metapb.PeerRole_Learner {
+			if !isTiFlashWriteNode(peer.StoreId) {
+				replicaNumber++
+			} else if !hasTiFlashWriteNode {
+				hasTiFlashWriteNode = true
+				replicaNumber++
+			}
+		}
+	}
+	return replicaNumber
+}
+
 func (s *RegionRequestSender) sendReqToRegion(
 	bo *retry.Backoffer, rpcCtx *RPCContext, req *tikvrpc.Request, timeout time.Duration,
 ) (resp *tikvrpc.Response, retry bool, err error) {
@@ -1348,13 +1374,7 @@ func (s *RegionRequestSender) sendReqToRegion(
 	// Count the replica number as the RU cost factor.
 	req.ReplicaNumber = 1
 	if rpcCtx.Meta != nil && len(rpcCtx.Meta.GetPeers()) > 0 {
-		req.ReplicaNumber = 0
-		for _, peer := range rpcCtx.Meta.GetPeers() {
-			role := peer.GetRole()
-			if role == metapb.PeerRole_Voter || role == metapb.PeerRole_Learner {
-				req.ReplicaNumber++
-			}
-		}
+		req.ReplicaNumber = int64(s.countReplicaNumber(rpcCtx.Meta.GetPeers()))
 	}
 
 	var sessionID uint64
diff --git a/rawkv/rawkv.go b/rawkv/rawkv.go
index cebd534..e1026e7 100644
--- a/rawkv/rawkv.go
+++ b/rawkv/rawkv.go
@@ -200,32 +200,33 @@ func NewClientWithOpts(ctx context.Context, pdAddrs []string, opts ...ClientOpt)
 		o(opt)
 	}
 
-	// Use an unwrapped PDClient to obtain keyspace meta.
-	pdCli, err := pd.NewClient(pdAddrs, pd.SecurityOption{
-		CAPath:   opt.security.ClusterSSLCA,
-		CertPath: opt.security.ClusterSSLCert,
-		KeyPath:  opt.security.ClusterSSLKey,
-	}, opt.pdOptions...)
-	if err != nil {
-		return nil, errors.WithStack(err)
-	}
-
-	// Build a CodecPDClient
-	var codecCli *tikv.CodecPDClient
+	var (
+		pdClient pd.Client
+		codecCli *tikv.CodecPDClient
+		err      error
+	)
 
 	switch opt.apiVersion {
 	case kvrpcpb.APIVersion_V1, kvrpcpb.APIVersion_V1TTL:
-		codecCli = locate.NewCodecPDClient(tikv.ModeRaw, pdCli)
+		pdClient, err = tikv.NewPDClientWithAPIContext(pdAddrs, pd.NewAPIContextV1())
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		codecCli = locate.NewCodecPDClient(tikv.ModeRaw, pdClient)
 	case kvrpcpb.APIVersion_V2:
-		codecCli, err = tikv.NewCodecPDClientWithKeyspace(tikv.ModeRaw, pdCli, opt.keyspace)
+		pdClient, err = tikv.NewPDClientWithAPIContext(pdAddrs, pd.NewAPIContextV2(opt.keyspace))
+		if err != nil {
+			return nil, errors.WithStack(err)
+		}
+		codecCli, err = tikv.NewCodecPDClientWithKeyspace(tikv.ModeRaw, pdClient, opt.keyspace)
 		if err != nil {
 			return nil, err
 		}
 	default:
-		return nil, errors.Errorf("unknown api version: %d", opt.apiVersion)
+		return nil, errors.Errorf("unknown API version: %d", opt.apiVersion)
 	}
 
-	pdCli = codecCli
+	pdClient = codecCli
 
 	rpcCli := client.NewRPCClient(
 		client.WithSecurity(opt.security),
@@ -235,9 +236,9 @@ func NewClientWithOpts(ctx context.Context, pdAddrs []string, opts ...ClientOpt)
 
 	return &Client{
 		apiVersion:  opt.apiVersion,
-		clusterID:   pdCli.GetClusterID(ctx),
-		regionCache: locate.NewRegionCache(pdCli),
-		pdClient:    pdCli,
+		clusterID:   pdClient.GetClusterID(ctx),
+		regionCache: locate.NewRegionCache(pdClient),
+		pdClient:    pdClient,
 		rpcClient:   rpcCli,
 	}, nil
 }
diff --git a/tikv/kv.go b/tikv/kv.go
index 3621a17..a9b4f7a 100644
--- a/tikv/kv.go
+++ b/tikv/kv.go
@@ -235,9 +235,15 @@ func NewKVStore(uuid string, pdClient pd.Client, spkv SafePointKV, tikvclient Cl
 
 // NewPDClient returns an unwrapped pd client.
 func NewPDClient(pdAddrs []string) (pd.Client, error) {
+	return NewPDClientWithAPIContext(pdAddrs, pd.NewAPIContextV1())
+}
+
+// NewPDClientWithAPIContext returns an unwrapped pd client with API context.
+func NewPDClientWithAPIContext(pdAddrs []string, apiContext pd.APIContext) (pd.Client, error) {
 	cfg := config.GetGlobalConfig()
 	// init pd-client
-	pdCli, err := pd.NewClient(
+	pdCli, err := pd.NewClientWithAPIContext(
+		context.Background(), apiContext,
 		pdAddrs, pd.SecurityOption{
 			CAPath:   cfg.Security.ClusterSSLCA,
 			CertPath: cfg.Security.ClusterSSLCert,
diff --git a/txnkv/client.go b/txnkv/client.go
index cfb40f5..30b7dd8 100644
--- a/txnkv/client.go
+++ b/txnkv/client.go
@@ -26,6 +26,7 @@ import (
 	"github.com/tikv/client-go/v2/tikv"
 	"github.com/tikv/client-go/v2/txnkv/transaction"
 	"github.com/tikv/client-go/v2/util"
+	pd "github.com/tikv/pd/client"
 )
 
 // Client is a txn client.
@@ -62,12 +63,24 @@ func NewClient(pdAddrs []string, opts ...ClientOpt) (*Client, error) {
 	for _, o := range opts {
 		o(opt)
 	}
-	// Use an unwrapped PDClient to obtain keyspace meta.
-	pdClient, err := tikv.NewPDClient(pdAddrs)
+
+	var (
+		pdClient pd.Client
+		apiContext pd.APIContext
+		err      error
+	)
+	switch opt.apiVersion {
+	case kvrpcpb.APIVersion_V1:
+		apiContext= pd.NewAPIContextV1()
+	case kvrpcpb.APIVersion_V2:
+		apiContext=pd.NewAPIContextV2(opt.keyspaceName)
+	default:
+		return nil, errors.Errorf("unknown API version: %d", opt.apiVersion)
+	}
+	pdClient, err = tikv.NewPDClientWithAPIContext(pdAddrs, apiContext)
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}
-
 	pdClient = util.InterceptedPDClient{Client: pdClient}
 
 	// Construct codec from options.
@@ -81,7 +94,7 @@ func NewClient(pdAddrs []string, opts ...ClientOpt) (*Client, error) {
 			return nil, err
 		}
 	default:
-		return nil, errors.Errorf("unknown api version: %d", opt.apiVersion)
+		return nil, errors.Errorf("unknown API version: %d", opt.apiVersion)
 	}
 
 	pdClient = codecCli


Test patch:
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index edd013c..66c5f0e 100644
--- a/internal/locate/region_request_test.go
+++ b/internal/locate/region_request_test.go
@@ -508,6 +508,10 @@ func (s *mockTikvGrpcServer) GetDisaggConfig(context.Context, *disaggregated.Get
 	return nil, errors.New("unreachable")
 }
 
+func (s *mockTikvGrpcServer) CancelDisaggTask(context.Context, *disaggregated.CancelDisaggTaskRequest) (*disaggregated.CancelDisaggTaskResponse, error) {
+	return nil, errors.New("unreachable")
+}
+
 func (s *testRegionRequestToSingleStoreSuite) TestNoReloadRegionForGrpcWhenCtxCanceled() {
 	// prepare a mock tikv grpc server
 	addr := "localhost:56341"
@@ -653,3 +657,43 @@ func (s *testRegionRequestToSingleStoreSuite) TestCloseConnectionOnStoreNotMatch
 	s.NotNil(regionErr)
 	s.Equal(target, client.closedAddr)
 }
+
+func (s *testRegionRequestToSingleStoreSuite) TestCountReplicaNumber() {
+	fmt.Println("TestCountReplicaNumber")
+	fmt.Println(s.cache.storeMu.stores)
+	tikvLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tikv"}}
+	tiflashLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}}
+	tiflashWNLabels := []*metapb.StoreLabel{{Key: "engine", Value: "tiflash"}, {Key: "engine_role", Value: "write"}}
+
+	s.cache.SetRegionCacheStore(1, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(2, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(3, "", "", tikvrpc.TiKV, 0, tikvLabels)
+	s.cache.SetRegionCacheStore(4, "", "", tikvrpc.TiFlash, 0, tiflashLabels)
+	s.cache.SetRegionCacheStore(5, "", "", tikvrpc.TiFlash, 0, tiflashLabels)
+	{
+		peers := []*metapb.Peer{{StoreId: 1, Role: metapb.PeerRole_Voter}}
+		s.Equal(1, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 2, Role: metapb.PeerRole_Voter})
+		s.Equal(2, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 3, Role: metapb.PeerRole_Voter})
+		s.Equal(3, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 4, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 5, Role: metapb.PeerRole_Learner})
+		s.Equal(5, s.regionRequestSender.countReplicaNumber(peers))
+	}
+	s.cache.SetRegionCacheStore(4, "", "", tikvrpc.TiFlash, 0, tiflashWNLabels)
+	s.cache.SetRegionCacheStore(5, "", "", tikvrpc.TiFlash, 0, tiflashWNLabels)
+	{
+		peers := []*metapb.Peer{{StoreId: 1, Role: metapb.PeerRole_Voter}}
+		s.Equal(1, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 2, Role: metapb.PeerRole_Voter})
+		s.Equal(2, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 3, Role: metapb.PeerRole_Voter})
+		s.Equal(3, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 4, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers))
+		peers = append(peers, &metapb.Peer{StoreId: 5, Role: metapb.PeerRole_Learner})
+		s.Equal(4, s.regionRequestSender.countReplicaNumber(peers))   // Only count 1 tiflash replica for tiflash write-nodes.
+	}
+}