From 864a5fc0c482ca6a3382c6655ad0fb705743a01b Mon Sep 17 00:00:00 2001
From: glorv <glorvs@163.com>
Date: Wed, 21 Jun 2023 11:10:26 +0800
Subject: [PATCH] make ResourceControlInterceptor atomic (#848)

Signed-off-by: glorv <glorvs@163.com>
---
 internal/client/client_interceptor.go | 13 ++++++++-----
 tikv/kv.go                            |  4 ++--
 2 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/internal/client/client_interceptor.go b/internal/client/client_interceptor.go
index e4b31b51..9e55a103 100644
--- a/internal/client/client_interceptor.go
+++ b/internal/client/client_interceptor.go
@@ -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
 				}
diff --git a/tikv/kv.go b/tikv/kv.go
index 5d3606a3..3e975361 100644
--- a/tikv/kv.go
+++ b/tikv/kv.go
@@ -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.