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

4
tikv/kv.go
View File

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