tikv
/
client-go
mirror of https://github.com/tikv/client-go.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Address the comments 

Signed-off-by: JmPotato <ghzpotato@gmail.com>
This commit is contained in:
JmPotato 2023-01-18 20:15:50 +08:00
parent 6ee294fedf
commit 35b648cf98
4 changed files with 33 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
go.mod
View File

@ -2,7 +2,7 @@ module github.com/tikv/client-go/v2
go 1.18
replace github.com/tikv/pd => github.com/CabinfeverB/pd v1.1.0-beta.0.20230117175918-55fedf50a131
replace github.com/tikv/pd => github.com/CabinfeverB/pd v1.1.0-beta.0.20230118092256-ed807b93fceb
require (
github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548

4
go.sum
View File

@ -2,8 +2,8 @@ cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMT
cloud.google.com/go v0.34.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/CabinfeverB/pd v1.1.0-beta.0.20230117175918-55fedf50a131 h1:agjNzxREwT+JqF/UocMsQa7JHhoOi167t9lcwRjtFq8=
github.com/CabinfeverB/pd v1.1.0-beta.0.20230117175918-55fedf50a131/go.mod h1:ZJUZ8xZBjCWZAQixiX0y2EE5+TFiNJ9lnP0bfB3gOu8=
github.com/CabinfeverB/pd v1.1.0-beta.0.20230118092256-ed807b93fceb h1:/SVKkznEk3eMSXtJYrle0hkkyI0GQ/b0hB6eIwQCeiM=
github.com/CabinfeverB/pd v1.1.0-beta.0.20230118092256-ed807b93fceb/go.mod h1:ZJUZ8xZBjCWZAQixiX0y2EE5+TFiNJ9lnP0bfB3gOu8=
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=
github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=

26
internal/client/client_interceptor.go
View File

@ -28,35 +28,29 @@ var _ Client = interceptedClient{}
type interceptedClient struct {
Client
resourceControlIt client.ResourceGroupKVInterceptor
}
// NewInterceptedClient creates a Client which can execute interceptor.
func NewInterceptedClient(client Client) Client {
return interceptedClient{client, nil}
}
// SetupResourceGroupKVInterceptor sets the interceptor for resource control.
func (r interceptedClient) SetupResourceGroupKVInterceptor(resourceControlIt client.ResourceGroupKVInterceptor) {
r.resourceControlIt = resourceControlIt
return interceptedClient{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 r.resourceControlIt != nil && len(resourceGroupName) > 0 {
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 := r.resourceControlIt.OnRequestWait(ctx, resourceGroupName, reqInfo)
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)
r.resourceControlIt.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
}
return resp, err
}
@ -81,3 +75,15 @@ func (r interceptedClient) SendRequest(ctx context.Context, addr string, req *ti
}
return r.Client.SendRequest(ctx, addr, req, timeout)
}
var resourceControlInterceptor client.ResourceGroupKVInterceptor
// SetupResourceControlInterceptor sets the interceptor for resource control.
func SetupResourceControlInterceptor(interceptor client.ResourceGroupKVInterceptor) {
resourceControlInterceptor = interceptor
}
// UnsetResourceControlInterceptor un-sets the interceptor for resource control.
func UnsetResourceControlInterceptor() {
resourceControlInterceptor = nil
}

14
internal/resource_control/resource_control.go
View File

@ -1,3 +1,17 @@
// 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 (