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

make ResourceControlInterceptor atomic (#848) 

Signed-off-by: glorv <glorvs@163.com>
This commit is contained in:
glorv 2023-06-21 11:10:26 +08:00 committed by GitHub
parent 18b17d4b72
commit 864a5fc0c4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
internal/client/client_interceptor.go
View File

@ -29,7 +29,6 @@ import (
func init() {
ResourceControlSwitch.Store(false)
ResourceControlInterceptor = nil
}
var _ Client = interceptedClient{}
@ -77,7 +76,7 @@ var (
// ResourceControlSwitch is used to control whether to enable the resource control.
ResourceControlSwitch atomic.Value
// ResourceControlInterceptor is used to build the resource control interceptor.
ResourceControlInterceptor resourceControlClient.ResourceGroupKVInterceptor
ResourceControlInterceptor atomic.Pointer[resourceControlClient.ResourceGroupKVInterceptor]
)
// buildResourceControlInterceptor builds a resource control interceptor with
@ -95,16 +94,20 @@ func buildResourceControlInterceptor(
if len(resourceGroupName) == 0 {
return nil
}
rcInterceptor := ResourceControlInterceptor.Load()
// No resource group interceptor is set.
if ResourceControlInterceptor == nil {
if rcInterceptor == nil {
return nil
}
resourceControlInterceptor := *rcInterceptor
// Make the request info.
reqInfo := resourcecontrol.MakeRequestInfo(req)
// Build the interceptor.
interceptFn := func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc {
return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
consumption, penalty, err := ResourceControlInterceptor.OnRequestWait(ctx, resourceGroupName, reqInfo)
consumption, penalty, err := resourceControlInterceptor.OnRequestWait(ctx, resourceGroupName, reqInfo)
if err != nil {
return nil, err
}
@ -113,7 +116,7 @@ func buildResourceControlInterceptor(
resp, err := next(target, req)
if resp != nil {
respInfo := resourcecontrol.MakeResponseInfo(resp)
consumption, err = ResourceControlInterceptor.OnResponse(resourceGroupName, reqInfo, respInfo)
consumption, err = resourceControlInterceptor.OnResponse(resourceGroupName, reqInfo, respInfo)
if err != nil {
return nil, err
}

4
tikv/kv.go
View File

@ -705,12 +705,12 @@ func DisableResourceControl() {
// SetResourceControlInterceptor sets the interceptor for resource control.
func SetResourceControlInterceptor(interceptor resourceControlClient.ResourceGroupKVInterceptor) {
client.ResourceControlInterceptor = interceptor
client.ResourceControlInterceptor.Store(&interceptor)
}
// UnsetResourceControlInterceptor un-sets the interceptor for resource control.
func UnsetResourceControlInterceptor() {
client.ResourceControlInterceptor = nil
client.ResourceControlInterceptor.Store(nil)
}
// Variables defines the variables used by TiKV storage.