Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ require (
modernc.org/mathutil v1.7.1
)

replace github.com/pingcap/kvproto => github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa

require (
github.com/beorn7/perks v1.0.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@ github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa h1:dryZb9j5lsPeqPHpwAsJXauQhLU9PiwLXF4bDh7TjLY=
github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=
Expand Down Expand Up @@ -81,8 +83,6 @@ github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXH
github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
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-20260622063236-b41e86365ce0 h1:MalBpLjhK/cS9ndCoSAg2C7aBzVojAqFVfzTKmuy0ho=
github.com/pingcap/kvproto v0.0.0-20260622063236-b41e86365ce0/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
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=
Expand Down
1 change: 1 addition & 0 deletions integration_tests/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -177,5 +177,6 @@ require (

replace (
github.com/go-ldap/ldap/v3 => github.com/YangKeao/ldap/v3 v3.4.5-0.20230421065457-369a3bab1117
github.com/pingcap/kvproto => github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa
github.com/tikv/client-go/v2 => ../
)
2 changes: 2 additions & 0 deletions integration_tests/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -857,6 +857,8 @@ github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1o
github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
github.com/YangKeao/ldap/v3 v3.4.5-0.20230421065457-369a3bab1117 h1:+OqGGFc2YHFd82aSHmjlILVt1t4JWJjrNIfV8cVEPow=
github.com/YangKeao/ldap/v3 v3.4.5-0.20230421065457-369a3bab1117/go.mod h1:bMGIq3AGbytbaMwf8wdv5Phdxz0FWHTIYMSzyrYgnQs=
github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa h1:dryZb9j5lsPeqPHpwAsJXauQhLU9PiwLXF4bDh7TjLY=
github.com/YuhaoZhang00/kvproto v0.0.0-20260717043146-2c5a1af9f8fa/go.mod h1:z6+aAHB7dBkA+LyinEX+48/ImRJ3jag0Hg0c7wkhEvE=
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d h1:licZJFw2RwpHMqeKTCYkitsPqHNxTmd4SNR5r94FGM8=
github.com/acarl005/stripansi v0.0.0-20180116102854-5a71ef0e047d/go.mod h1:asat636LX7Bqt5lYEZ27JNDcqxfjdBQuJ/MM4CN/Lzo=
github.com/ajstarks/deck v0.0.0-20200831202436-30c9fc6549a9/go.mod h1:JynElWSGnm/4RlzPXRlREEwqTHAN3T56Bv2ITsFT3gY=
Expand Down
32 changes: 23 additions & 9 deletions internal/resourcecontrol/resource_control.go
Original file line number Diff line number Diff line change
Expand Up @@ -191,9 +191,10 @@ func isCopRequest(req *tikvrpc.Request) bool {
// 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
kvCPU time.Duration
respSize uint64
readBytes uint64
remoteReadBytes uint64
kvCPU time.Duration
respSize uint64
}

// MakeResponseInfo extracts the relevant information from a BatchResponse.
Expand All @@ -203,9 +204,10 @@ func MakeResponseInfo(resp *tikvrpc.Response) *ResponseInfo {
}
// Parse the response to extract the info.
var (
readBytes uint64
detailsV2 *kvrpcpb.ExecDetailsV2
details *kvrpcpb.ExecDetails
readBytes uint64
remoteReadBytes uint64
detailsV2 *kvrpcpb.ExecDetailsV2
details *kvrpcpb.ExecDetails
)
switch r := resp.Resp.(type) {
case *coprocessor.Response:
Expand Down Expand Up @@ -242,6 +244,10 @@ func MakeResponseInfo(resp *tikvrpc.Response) *ResponseInfo {
// processed versions size is greater than the total versions size,
// we use the processed versions size as the read bytes.
readBytes = max(scanDetail.GetTotalVersionsSize(), scanDetail.GetProcessedVersionsSize())
remoteReadBytes = max(
scanDetail.GetRemoteTotalVersionsSize(),
scanDetail.GetRemoteProcessedVersionsSize(),
)
} else {
// NOTE: The original design intended to account for all MVCC read
// overhead, but TotalVersionsSize did not exist at the time, so
Expand All @@ -257,9 +263,10 @@ func MakeResponseInfo(resp *tikvrpc.Response) *ResponseInfo {
// Get the KV CPU time in milliseconds from the execution time details.
kvCPU := getKVCPU(detailsV2, details)
return &ResponseInfo{
readBytes: readBytes,
kvCPU: kvCPU,
respSize: uint64(resp.GetSize()),
readBytes: readBytes,
remoteReadBytes: remoteReadBytes,
kvCPU: kvCPU,
respSize: uint64(resp.GetSize()),
}
}

Expand All @@ -282,6 +289,13 @@ func (res *ResponseInfo) ReadBytes() uint64 {
return res.readBytes
}

// RemoteReadBytes returns the factual subset of ReadBytes processed by a
// remote coprocessor. It is zero outside NextGen so legacy RU pricing remains
// unchanged.
func (res *ResponseInfo) RemoteReadBytes() uint64 {
return res.remoteReadBytes
}

// KVCPU returns the KV CPU time of the response.
func (res *ResponseInfo) KVCPU() time.Duration {
return res.kvCPU
Expand Down
15 changes: 11 additions & 4 deletions internal/resourcecontrol/resource_control_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -107,17 +107,21 @@ func TestResponseInfoReadBytes(t *testing.T) {
Resp: &coprocessor.Response{
ExecDetailsV2: &kvrpcpb.ExecDetailsV2{
ScanDetailV2: &kvrpcpb.ScanDetailV2{
TotalVersionsSize: 100,
ProcessedVersionsSize: 80,
TotalVersionsSize: 100,
ProcessedVersionsSize: 80,
RemoteTotalVersionsSize: 40,
RemoteProcessedVersionsSize: 30,
},
},
},
}
info := MakeResponseInfo(resp)
if config.NextGen {
assert.Equal(t, uint64(100), info.ReadBytes())
assert.Equal(t, uint64(40), info.RemoteReadBytes())
} else {
assert.Equal(t, uint64(80), info.ReadBytes())
assert.Zero(t, info.RemoteReadBytes())
}

if config.NextGen {
Expand All @@ -126,13 +130,16 @@ func TestResponseInfoReadBytes(t *testing.T) {
Resp: &coprocessor.Response{
ExecDetailsV2: &kvrpcpb.ExecDetailsV2{
ScanDetailV2: &kvrpcpb.ScanDetailV2{
TotalVersionsSize: 80,
ProcessedVersionsSize: 100,
TotalVersionsSize: 80,
ProcessedVersionsSize: 100,
RemoteTotalVersionsSize: 60,
RemoteProcessedVersionsSize: 70,
},
},
},
}
infoCompat := MakeResponseInfo(respCompat)
assert.Equal(t, uint64(100), infoCompat.ReadBytes())
assert.Equal(t, uint64(70), infoCompat.RemoteReadBytes())
}
}
6 changes: 6 additions & 0 deletions tikvrpc/tikvrpc.go
Original file line number Diff line number Diff line change
Expand Up @@ -884,6 +884,12 @@ func AttachContext(req *Request, rpcCtx kvrpcpb.Context) bool {
if cmd == CmdCopStream {
cmd = CmdCop
}
if cmd == CmdCop {
// Only NextGen's RU-v1 controller can price the factual remote read
// subset separately. Legacy clients must keep the worker's pre-priced
// compatibility response.
ctx.ClientPricesRemoteReads = config.NextGen
}
if patchCmdCtx(req, cmd, ctx) {
return true
}
Expand Down
12 changes: 12 additions & 0 deletions tikvrpc/tikvrpc_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,18 @@ func TestDefaultRequestOrigin(t *testing.T) {
require.Equal(t, kvrpcpb.RequestOrigin_RequestOriginTiDB, req.GetRequestOrigin())
}

func TestAttachContextAdvertisesRemoteReadPricingForNextGenCop(t *testing.T) {
for _, cmd := range []CmdType{CmdCop, CmdCopStream} {
req := NewRequest(cmd, &coprocessor.Request{})
require.True(t, AttachContext(req, kvrpcpb.Context{}))
require.Equal(t, config.NextGen, req.Cop().GetContext().GetClientPricesRemoteReads())
}

req := NewRequest(CmdGet, &kvrpcpb.GetRequest{})
require.True(t, AttachContext(req, kvrpcpb.Context{}))
require.False(t, req.Get().GetContext().GetClientPricesRemoteReads())
}

func TestAttachContextSetsRequestContext(t *testing.T) {
rpcCtx := kvrpcpb.Context{
RegionId: 123,
Expand Down
Loading