Introduce the resource control client

config/config.go

@@ -44,6 +44,7 @@ import (
 	"github.com/tikv/client-go/v2/internal/logutil"
 	"github.com/tikv/client-go/v2/oracle"
 	"github.com/tikv/client-go/v2/util"
+	"github.com/tikv/pd/pkg/mcs/resource_manager/client"
 	"go.uber.org/zap"
 )
 
@@ -78,6 +79,7 @@ type Config struct {
 	TxnScope              string
 	EnableAsyncCommit     bool
 	Enable1PC             bool
+	ResourceControl       client.RequestUnitConfig
 }
 
 // DefaultConfig returns the default configuration.
@@ -95,6 +97,7 @@ func DefaultConfig() Config {
 		TxnScope:              "",
 		EnableAsyncCommit:     false,
 		Enable1PC:             false,
+		ResourceControl:       *client.DefaultRequestUnitConfig(),
 	}
 }
 

go.mod

@@ -2,6 +2,8 @@ module github.com/tikv/client-go/v2
 
 go 1.18
 
+replace github.com/tikv/pd => github.com/CabinfeverB/pd v1.1.0-beta.0.20230118184315-8efd2e043066
+
 require (
 	github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548
 	github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2
@@ -22,6 +24,7 @@ require (
 	github.com/stathat/consistent v1.0.0
 	github.com/stretchr/testify v1.8.1
 	github.com/tiancaiamao/gp v0.0.0-20221230034425-4025bc8a4d4a
+	github.com/tikv/pd v1.1.0-beta.0.20230118040950-082fc6a9bc2e
 	github.com/tikv/pd/client v0.0.0-20230118040950-082fc6a9bc2e
 	github.com/twmb/murmur3 v1.1.3
 	go.etcd.io/etcd/api/v3 v3.5.2
@@ -40,10 +43,11 @@ require (
 	github.com/coreos/go-semver v0.3.0 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
-	github.com/golang/snappy v0.0.3 // indirect
+	github.com/golang/snappy v0.0.4 // indirect
+	github.com/kr/text v0.2.0 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
+	github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e // indirect
 	github.com/onsi/ginkgo v1.16.5 // indirect
-	github.com/onsi/gomega v1.18.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/prometheus/common v0.26.0 // indirect
 	github.com/prometheus/procfs v0.6.0 // indirect
@@ -55,6 +59,7 @@ require (
 	golang.org/x/text v0.4.0 // indirect
 	google.golang.org/genproto v0.0.0-20221202195650-67e5cbc046fd // indirect
 	google.golang.org/protobuf v1.28.1 // indirect
+	gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	stathat.com/c/consistent v1.0.0 // indirect

go.sum

@@ -377,6 +377,8 @@ dmitri.shuralyov.com/gpu/mtl v0.0.0-20190408044501-666a987793e9/go.mod h1:H6x//7
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
 github.com/BurntSushi/xgb v0.0.0-20160522181843-27f122750802/go.mod h1:IVnqGOEym/WlBOVXweHU+Q+/VP0lqqI8lqeDx9IjBqo=
+github.com/CabinfeverB/pd v1.1.0-beta.0.20230118184315-8efd2e043066 h1:Dzy4RAD15MmCOgHChzal7H+z+5uvwpVmmAB3C65YiEw=
+github.com/CabinfeverB/pd v1.1.0-beta.0.20230118184315-8efd2e043066/go.mod h1:ZJUZ8xZBjCWZAQixiX0y2EE5+TFiNJ9lnP0bfB3gOu8=
 github.com/OneOfOne/xxhash v1.2.2/go.mod h1:HSdplMjZKSmBqAxg5vPj2TmRDmfkzw+cTzAElWljhcU=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
@@ -412,6 +414,7 @@ github.com/coreos/go-semver v0.3.0 h1:wkHLiw0WNATZnSG7epLsujiMCgPAc9xhjJ4tgnAxmf
 github.com/coreos/go-semver v0.3.0/go.mod h1:nnelYz7RCh+5ahJtPPxZlU+153eP4D4r3EedlOD2RNk=
 github.com/coreos/go-systemd/v22 v22.3.2 h1:D9/bQk5vlXQFZ6Kwuu6zaiXJ9oTPe68++AzAJc1DzSI=
 github.com/coreos/go-systemd/v22 v22.3.2/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
 github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548 h1:iwZdTE0PVqJCos1vaoKsclOGD3ADKpshg3SRtYBbwso=
 github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548/go.mod h1:e6NPNENfs9mPDVNRekM7lKScauxd5kXTr1Mfyig6TDM=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -481,8 +484,9 @@ github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaS
 github.com/golang/protobuf v1.5.1/go.mod h1:DopwsBzvsk0Fs44TXzsVbJyPhcCPeIwnvohx4u74HPM=
 github.com/golang/protobuf v1.5.2 h1:ROPKBNFfQgOUMifHyP+KYbvpjbdoFNs+aK7DXlji0Tw=
 github.com/golang/protobuf v1.5.2/go.mod h1:XVQd3VNwM+JqD3oG2Ue2ip4fOMUkwXdXDdiuN0vRsmY=
-github.com/golang/snappy v0.0.3 h1:fHPg5GQYlCeLIPB9BZqMVR5nR9A+IM5zcgeTdjMYmLA=
 github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
+github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.0.0/go.mod h1:lNA+9X1NB3Zf8V7Ke586lFgjr2dZNuvo3lPJSGZ5JPQ=
 github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU=
@@ -519,7 +523,6 @@ github.com/google/pprof v0.0.0-20201023163331-3e6fc7fc9c4c/go.mod h1:kpwsk12EmLe
 github.com/google/pprof v0.0.0-20201203190320-1bf35d6f28c2/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/pprof v0.0.0-20210122040257-d980be63207e/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/pprof v0.0.0-20210226084205-cbba55b83ad5/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
-github.com/google/pprof v0.0.0-20210407192527-94a9f03dee38/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/pprof v0.0.0-20210601050228-01bbb1931b22/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/pprof v0.0.0-20210609004039-a478d1d731e9/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
 github.com/google/pprof v0.0.0-20210720184732-4bb14d4b1be1/go.mod h1:kpwsk12EmLew5upagYY7GY0pfYCcupk39gWOCRROcvE=
@@ -564,11 +567,11 @@ github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+o
 github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
 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 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
 github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORNo=
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
-github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 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=
+github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/matttproud/golang_protobuf_extensions v1.0.1 h1:4hp9jkHxhMHkqkrB3Ix0jegS5sx/RkqARlsWZ6pIwiU=
 github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
@@ -577,20 +580,18 @@ github.com/modern-go/reflect2 v0.0.0-20180701023420-4b7aa43c6742/go.mod h1:bx2lN
 github.com/modern-go/reflect2 v1.0.1/go.mod h1:bx2lNnkwVCuqBIxFjflWJWanXIb3RllmbCylyMrvgv0=
 github.com/mwitkow/go-conntrack v0.0.0-20161129095857-cc309e4a2223/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
 github.com/mwitkow/go-conntrack v0.0.0-20190716064945-2f068394615f/go.mod h1:qRWi+5nqEBWmkhHvq77mSJWrCKwh8bxhgT7d/eI7P4U=
+github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e h1:fD57ERR4JtEqsWbfPhv4DMiApHyliiK5xCTNVSPiaAs=
+github.com/niemeyer/pretty v0.0.0-20200227124842-a10e7caefd8e/go.mod h1:zD1mROLANZcx1PVRCS0qkT7pwLkGfwJo4zjcN/Tysno=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
 github.com/nxadm/tail v1.4.8 h1:nPr65rt6Y5JFSKQO7qToXr7pePgD6Gwiw05lkbyAQTE=
 github.com/nxadm/tail v1.4.8/go.mod h1:+ncqLTQzXmGhMZNUePPaPqPvBxHAIsmXswZKocGu+AU=
 github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
-github.com/onsi/ginkgo v1.16.4/go.mod h1:dX+/inL/fNMqNlz0e9LfyB9TswhZpCVdJM/Z6Vvnwo0=
 github.com/onsi/ginkgo v1.16.5 h1:8xi0RTUf59SOSfEtZMvwTvXYMzG4gV23XVHOZiXNtnE=
 github.com/onsi/ginkgo v1.16.5/go.mod h1:+E8gABHa3K6zRBolWtd+ROzc/U5bkGt0FwiG042wbpU=
-github.com/onsi/ginkgo/v2 v2.0.0/go.mod h1:vw5CSIxN1JObi/U8gcbwft7ZxR2dgaR70JSE3/PpL4c=
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
-github.com/onsi/gomega v1.17.0/go.mod h1:HnhC7FXeEQY45zxNK3PPoIUhzk/80Xly9PcubAlGdZY=
-github.com/onsi/gomega v1.18.1 h1:M1GfJqGRrBrrGGsbxzV5dqM2U2ApXefZCQpkukxYRLE=
-github.com/onsi/gomega v1.18.1/go.mod h1:0q+aL8jAiMXy9hbwj2mr5GziHiwhAIQpFmmtT5hitRs=
+github.com/onsi/gomega v1.20.1 h1:PA/3qinGoukvymdIDV8pii6tiZgC8kbmJO6Z5+b002Q=
 github.com/opentracing/opentracing-go v1.1.0/go.mod h1:UkNAQd3GIcIGf0SeVgPpRdFStlNbqXla1AfSYxPUl2o=
 github.com/opentracing/opentracing-go v1.2.0 h1:uEJPy/1a5RIPAJ0Ov+OIO8OxWu77jEv+1B0VhjKrZUs=
 github.com/opentracing/opentracing-go v1.2.0/go.mod h1:GxEUsuufX4nBwe+T+Wl9TAgYrxe9dPLANfrWvHYVTgc=
@@ -785,7 +786,6 @@ golang.org/x/net v0.0.0-20210119194325-5f4716e94777/go.mod h1:m0MpNAwzfU5UDzcl9v
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
 golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc=
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
-golang.org/x/net v0.0.0-20210428140749-89ef3d95e781/go.mod h1:OJAsFXCWl8Ukc7SiCT/9KSuxbyM7479/AVlXFRxuMCk=
 golang.org/x/net v0.0.0-20210503060351-7fd8e65b6420/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
 golang.org/x/net v0.0.0-20220127200216-cd36cc0744dd/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
@@ -1230,8 +1230,9 @@ google.golang.org/protobuf v1.28.1/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqw
 gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f h1:BLraFXnmrev5lT+xlilqcH8XK9/i0At2xKjWk4p6zsU=
+gopkg.in/check.v1 v1.0.0-20200227125254-8fa46927fb4f/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/errgo.v2 v2.1.0/go.mod h1:hNsd1EY+bozCKY1Ytp96fpM3vjJbqLJn88ws8XvfDNI=
 gopkg.in/fsnotify.v1 v1.4.7/go.mod h1:Tz8NjZHkW78fSQdbUxIjBTcgA1z1m8ZHf0WmKUhAMys=
 gopkg.in/natefinch/lumberjack.v2 v2.0.0 h1:1Lc07Kr7qY4U2YPouBjpCLxpiyxIVoxqXgkXLknAOE8=

internal/client/client_interceptor.go

@@ -18,8 +18,10 @@ import (
 	"context"
 	"time"
 
+	"github.com/tikv/client-go/v2/internal/resource_control"
 	"github.com/tikv/client-go/v2/tikvrpc"
 	"github.com/tikv/client-go/v2/tikvrpc/interceptor"
+	"github.com/tikv/pd/pkg/mcs/resource_manager/client"
 )
 
 var _ Client = interceptedClient{}
@@ -34,10 +36,54 @@ func NewInterceptedClient(client Client) Client {
 }
 
 func (r interceptedClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
+	// Build the resource control interceptor if there is one and the resource group name is given.
+	var rcInterceptor interceptor.RPCInterceptor
+	resourceGroupName := req.GetResourceGroupName()
+	if resourceControlInterceptor != nil && len(resourceGroupName) > 0 {
+		rcInterceptor = func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc {
+			return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
+				reqInfo := resource_control.MakeRequestInfo(req)
+				err := resourceControlInterceptor.OnRequestWait(ctx, resourceGroupName, reqInfo)
+				if err != nil {
+					return nil, err
+				}
+				resp, err := next(target, req)
+				if resp != nil {
+					respInfo := resource_control.MakeResponseInfo(resp)
+					resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
+				}
+				return resp, err
+			}
+		}
+	}
+	// Chain the interceptors if there are multiple interceptors.
+	var finalInterceptor interceptor.RPCInterceptor
+	if rcInterceptor != nil {
+		finalInterceptor = rcInterceptor
+	}
 	if it := interceptor.GetRPCInterceptorFromCtx(ctx); it != nil {
-		return it(func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
+		if finalInterceptor != nil {
+			finalInterceptor = interceptor.ChainRPCInterceptors(finalInterceptor, it)
+		} else {
+			finalInterceptor = it
+		}
+	}
+	if finalInterceptor != nil {
+		return finalInterceptor(func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
 			return r.Client.SendRequest(ctx, target, req, timeout)
 		})(addr, req)
 	}
 	return r.Client.SendRequest(ctx, addr, req, timeout)
 }
+
+var resourceControlInterceptor client.ResourceGroupKVInterceptor
+
+// SetResourceControlInterceptor sets the interceptor for resource control.
+func SetResourceControlInterceptor(interceptor client.ResourceGroupKVInterceptor) {
+	resourceControlInterceptor = interceptor
+}
+
+// UnsetResourceControlInterceptor un-sets the interceptor for resource control.
+func UnsetResourceControlInterceptor() {
+	resourceControlInterceptor = nil
+}

internal/resource_control/resource_control.go

@@ -0,0 +1,135 @@
+// Copyright 2023 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package resource_control
+
+import (
+	"reflect"
+	"unsafe"
+
+	"github.com/pingcap/kvproto/pkg/coprocessor"
+	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/pingcap/log"
+	"github.com/tikv/client-go/v2/tikvrpc"
+	"go.uber.org/zap"
+)
+
+// RequestInfo contains information about a request that is able to calculate the RU cost
+// before the request is sent. Specifically, the write bytes RU cost of a write request
+// could be calculated by its key size to write.
+type RequestInfo struct {
+	// writeBytes is the actual write size if the request is a write request,
+	// or -1 if it's a read request.
+	writeBytes int64
+}
+
+// MakeRequestInfo extracts the relevant information from a BatchRequest.
+func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo {
+	if !req.IsTxnWriteRequest() && !req.IsRawWriteRequest() {
+		return &RequestInfo{writeBytes: -1}
+	}
+
+	var writeBytes int64
+	switch r := req.Req.(type) {
+	case *kvrpcpb.PrewriteRequest:
+		writeBytes += int64(r.TxnSize)
+	case *kvrpcpb.CommitRequest:
+		writeBytes += int64(unsafe.Sizeof(r.Keys))
+	}
+
+	return &RequestInfo{writeBytes: writeBytes}
+}
+
+// IsWrite returns whether the request is a write request.
+func (req *RequestInfo) IsWrite() bool {
+	return req.writeBytes > -1
+}
+
+// WriteBytes returns the actual write size of the request,
+// -1 will be returned if it's not a write request.
+func (req *RequestInfo) WriteBytes() uint64 {
+	return uint64(req.writeBytes)
+}
+
+// ResponseInfo contains information about a response that is able to calculate the RU cost
+// after the response is received. Specifically, the read bytes RU cost of a read request
+// could be calculated by its response size, and the KV CPU time RU cost of a request could
+// be calculated by its execution details info.
+type ResponseInfo struct {
+	readBytes uint64
+	kvCPUMs   uint64
+}
+
+// MakeResponseInfo extracts the relevant information from a BatchResponse.
+func MakeResponseInfo(resp *tikvrpc.Response) *ResponseInfo {
+	if resp.Resp == nil {
+		return &ResponseInfo{}
+	}
+	// Parse the response to extract the info.
+	var (
+		readBytes uint64
+		detailsV2 *kvrpcpb.ExecDetailsV2
+		details   *kvrpcpb.ExecDetails
+	)
+	switch r := resp.Resp.(type) {
+	case *coprocessor.Response:
+		detailsV2 = r.GetExecDetailsV2()
+		details = r.GetExecDetails()
+		readBytes = uint64(r.Data.Size())
+	case *tikvrpc.CopStreamResponse:
+		// Streaming request returns `io.EOF``, so the first `CopStreamResponse.Response`` may be nil.
+		if r.Response != nil {
+			detailsV2 = r.Response.GetExecDetailsV2()
+			details = r.Response.GetExecDetails()
+		}
+		readBytes = uint64(r.Data.Size())
+	case *kvrpcpb.GetResponse:
+		detailsV2 = r.GetExecDetailsV2()
+	case *kvrpcpb.BatchGetResponse:
+		detailsV2 = r.GetExecDetailsV2()
+	case *kvrpcpb.ScanResponse:
+		// TODO: using a more accurate size rather than using the whole response size as the read bytes.
+		readBytes = uint64(r.Size())
+	default:
+		log.Warn("[kv resource] unknown response type to collect the info", zap.Any("type", reflect.TypeOf(r)))
+		return &ResponseInfo{}
+	}
+	// Try to get read bytes from the `detailsV2`.
+	// TODO: clarify whether we should count the underlying storage engine read bytes or not.
+	if scanDetail := detailsV2.GetScanDetailV2(); scanDetail != nil {
+		readBytes = scanDetail.GetProcessedVersionsSize()
+	}
+	// Get the KV CPU time in milliseconds from the execution time details.
+	kvCPUMs := getKVCPUMs(detailsV2, details)
+	return &ResponseInfo{readBytes: readBytes, kvCPUMs: kvCPUMs}
+}
+
+// TODO: find out a more accurate way to get the actual KV CPU time.
+func getKVCPUMs(detailsV2 *kvrpcpb.ExecDetailsV2, details *kvrpcpb.ExecDetails) uint64 {
+	if timeDetail := detailsV2.GetTimeDetail(); timeDetail != nil {
+		return timeDetail.GetProcessWallTimeMs()
+	}
+	if timeDetail := details.GetTimeDetail(); timeDetail != nil {
+		return timeDetail.GetProcessWallTimeMs()
+	}
+	return 0
+}
+
+func (res *ResponseInfo) ReadBytes() uint64 {
+	return res.readBytes
+}
+
+func (res *ResponseInfo) KVCPUMs() uint64 {
+	return res.kvCPUMs
+}

tikvrpc/tikvrpc.go

@@ -1263,5 +1263,15 @@ func (req *Request) IsTxnWriteRequest() bool {
 	return false
 }
 
+// IsRawWriteRequest checks if the request is a raw write request.
+func (req *Request) IsRawWriteRequest() bool {
+	if req.Type == CmdRawPut ||
+		req.Type == CmdRawBatchPut ||
+		req.Type == CmdRawDelete {
+		return true
+	}
+	return false
+}
+
 // ResourceGroupTagger is used to fill the ResourceGroupTag in the kvrpcpb.Context.
 type ResourceGroupTagger func(req *Request)