From a30ce371a330a206950fbf4f4119c0c877e2b319 Mon Sep 17 00:00:00 2001 From: husharp Date: Tue, 11 Jul 2023 17:09:49 +0800 Subject: [PATCH] address comment Signed-off-by: husharp --- internal/client/client_interceptor.go | 2 -- internal/resourcecontrol/resource_control.go | 9 +++------ util/request_source.go | 2 +- util/request_source_test.go | 14 +++++++------- 4 files changed, 11 insertions(+), 16 deletions(-) diff --git a/internal/client/client_interceptor.go b/internal/client/client_interceptor.go index 4e6ca99e1..9e55a1038 100644 --- a/internal/client/client_interceptor.go +++ b/internal/client/client_interceptor.go @@ -107,8 +107,6 @@ func buildResourceControlInterceptor( // Build the interceptor. interceptFn := func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc { return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) { - // Background requests don't need to update the runtime stats because it will record in TiKV. - reqInfo.SetRequestSource(util.RequestSourceFromCtx(ctx)) consumption, penalty, err := resourceControlInterceptor.OnRequestWait(ctx, resourceGroupName, reqInfo) if err != nil { return nil, err diff --git a/internal/resourcecontrol/resource_control.go b/internal/resourcecontrol/resource_control.go index 263143a58..635d8d2d6 100644 --- a/internal/resourcecontrol/resource_control.go +++ b/internal/resourcecontrol/resource_control.go @@ -40,7 +40,7 @@ type RequestInfo struct { // MakeRequestInfo extracts the relevant information from a BatchRequest. func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo { if !req.IsTxnWriteRequest() && !req.IsRawWriteRequest() { - return &RequestInfo{writeBytes: -1} + return &RequestInfo{writeBytes: -1, requestSource: req.RequestSource} } var writeBytes int64 @@ -58,7 +58,8 @@ func MakeRequestInfo(req *tikvrpc.Request) *RequestInfo { writeBytes += int64(len(k)) } } - return &RequestInfo{writeBytes: writeBytes, storeID: req.Context.Peer.StoreId, replicaNumber: req.ReplicaNumber} + return &RequestInfo{writeBytes: writeBytes, storeID: req.Context.Peer.StoreId, + replicaNumber: req.ReplicaNumber, requestSource: req.RequestSource} } // IsWrite returns whether the request is a write request. @@ -84,10 +85,6 @@ func (req *RequestInfo) RequestSource() string { return req.requestSource } -func (req *RequestInfo) SetRequestSource(requestSource string) { - req.requestSource = requestSource -} - // 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 diff --git a/util/request_source.go b/util/request_source.go index 5aca0d726..ed0c9bede 100644 --- a/util/request_source.go +++ b/util/request_source.go @@ -95,7 +95,7 @@ func WithInternalSourceType(ctx context.Context, source string) context.Context }) } -// BuildRequestSource builds a request source from internal, source and explicitSource. +// BuildRequestSource builds a request_source from internal, source and explicitSource. func BuildRequestSource(internal bool, source, explicitSource string) string { requestSource := RequestSource{ RequestSourceInternal: internal, diff --git a/util/request_source_test.go b/util/request_source_test.go index a1c60e9dd..c7548bddb 100644 --- a/util/request_source_test.go +++ b/util/request_source_test.go @@ -69,13 +69,8 @@ func TestGetRequestSource(t *testing.T) { func TestBuildRequestSource(t *testing.T) { // Test internal request - expected := "unknown_default" - actual := BuildRequestSource(true, "", "") - assert.Equal(t, expected, actual) - - // Test internal request - expected = "internal_test_lightning" - actual = BuildRequestSource(true, "test", "lightning") + expected := "internal_test_lightning" + actual := BuildRequestSource(true, "test", "lightning") assert.Equal(t, expected, actual) // Test external request @@ -92,4 +87,9 @@ func TestBuildRequestSource(t *testing.T) { expected = "external_unknown_lightning" actual = BuildRequestSource(false, "", "lightning") assert.Equal(t, expected, actual) + + // Test RequestSourceType && ExplicitRequestSourceType both empty + expected = "unknown_default" + actual = BuildRequestSource(true, "", "") + assert.Equal(t, expected, actual) }