Simplify the code

Signed-off-by: JmPotato <ghzpotato@gmail.com>

internal/client/client_interceptor.go

@@ -26,8 +26,8 @@ import (
 )
 
 func init() {
-	resourceControlSwitch.Store(false)
-	resourceControlInterceptor = nil
+	ResourceControlSwitch.Store(false)
+	ResourceControlInterceptor = nil
 }
 
 var _ Client = interceptedClient{}
@@ -62,33 +62,20 @@ func (r interceptedClient) SendRequest(ctx context.Context, addr string, req *ti
 }
 
 var (
-	resourceControlSwitch      atomic.Value
-	resourceControlInterceptor resourceControlClient.ResourceGroupKVInterceptor
+	// 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
 )
 
-// EnableResourceControl enables the resource control.
-func EnableResourceControl() {
-	resourceControlSwitch.Store(true)
-}
-
-// DisableResourceControl disables the resource control.
-func DisableResourceControl() {
-	resourceControlSwitch.Store(false)
-}
-
-// SetResourceControlInterceptor sets the interceptor for resource control.
-func SetResourceControlInterceptor(interceptor resourceControlClient.ResourceGroupKVInterceptor) {
-	resourceControlInterceptor = interceptor
-}
-
 // buildResourceControlInterceptor builds a resource control interceptor with
 // the given resource group name.
 func buildResourceControlInterceptor(
 	ctx context.Context,
 	req *tikvrpc.Request,
 	resourceGroupName string,
 ) interceptor.RPCInterceptor {
-	if !resourceControlSwitch.Load().(bool) {
+	if !ResourceControlSwitch.Load().(bool) {
 		return nil
 	}
 	// When the group name is empty or "default", we don't need to
@@ -97,22 +84,22 @@ func buildResourceControlInterceptor(
 		return nil
 	}
 	// No resource group interceptor is set.
-	if resourceControlInterceptor == nil {
+	if ResourceControlInterceptor == nil {
 		return nil
 	}
 	// Make the request info.
 	reqInfo := resourcecontrol.MakeRequestInfo(req)
 	// Build the interceptor.
 	return func(next interceptor.RPCInterceptorFunc) interceptor.RPCInterceptorFunc {
 		return func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
-			err := resourceControlInterceptor.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 := resourcecontrol.MakeResponseInfo(resp)
-				resourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
+				ResourceControlInterceptor.OnResponse(ctx, resourceGroupName, reqInfo, respInfo)
 			}
 			return resp, err
 		}

tikv/kv.go

@@ -585,22 +585,22 @@ func (s *KVStore) updateSafeTS(ctx context.Context) {
 
 // EnableResourceControl enables the resource control.
 func EnableResourceControl() {
-	client.EnableResourceControl()
+	client.ResourceControlSwitch.Store(true)
 }
 
 // DisableResourceControl disables the resource control.
 func DisableResourceControl() {
-	client.DisableResourceControl()
+	client.ResourceControlSwitch.Store(false)
 }
 
 // SetResourceControlInterceptor sets the interceptor for resource control.
 func SetResourceControlInterceptor(interceptor resourceControlClient.ResourceGroupKVInterceptor) {
-	client.SetResourceControlInterceptor(interceptor)
+	client.ResourceControlInterceptor = interceptor
 }
 
 // UnsetResourceControlInterceptor un-sets the interceptor for resource control.
 func UnsetResourceControlInterceptor() {
-	client.SetResourceControlInterceptor(nil)
+	client.ResourceControlInterceptor = nil
 }
 
 // Variables defines the variables used by TiKV storage.