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

Export the switch 

Signed-off-by: JmPotato <ghzpotato@gmail.com>
This commit is contained in:
JmPotato 2023-01-19 17:18:56 +08:00
parent 8553f26149
commit 2a19c65068
2 changed files with 24 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
internal/client/client_interceptor.go
View File

@ -22,7 +22,7 @@ import (
"github.com/tikv/client-go/v2/internal/resourcecontrol" "github.com/tikv/client-go/v2/internal/resourcecontrol"
"github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/tikvrpc"
"github.com/tikv/client-go/v2/tikvrpc/interceptor" "github.com/tikv/client-go/v2/tikvrpc/interceptor"
"github.com/tikv/pd/pkg/mcs/resource_manager/client" resourceControlClient "github.com/tikv/pd/pkg/mcs/resource_manager/client"
) )
func init() { func init() {
@ -63,7 +63,7 @@ func (r interceptedClient) SendRequest(ctx context.Context, addr string, req *ti
var ( var (
resourceControlSwitch atomic.Value resourceControlSwitch atomic.Value
resourceControlInterceptor client.ResourceGroupKVInterceptor resourceControlInterceptor resourceControlClient.ResourceGroupKVInterceptor
) )
// EnableResourceControl enables the resource control. // EnableResourceControl enables the resource control.
@ -77,15 +77,10 @@ func DisableResourceControl() {
} }
// SetResourceControlInterceptor sets the interceptor for resource control. // SetResourceControlInterceptor sets the interceptor for resource control.
func SetResourceControlInterceptor(interceptor client.ResourceGroupKVInterceptor) { func SetResourceControlInterceptor(interceptor resourceControlClient.ResourceGroupKVInterceptor) {
resourceControlInterceptor = interceptor resourceControlInterceptor = interceptor
} }
// UnsetResourceControlInterceptor un-sets the interceptor for resource control.
func UnsetResourceControlInterceptor() {
resourceControlInterceptor = nil
}
// buildResourceControlInterceptor builds a resource control interceptor with // buildResourceControlInterceptor builds a resource control interceptor with
// the given resource group name. // the given resource group name.
func buildResourceControlInterceptor( func buildResourceControlInterceptor(

21
tikv/kv.go
View File

@ -67,6 +67,7 @@ import (
"github.com/tikv/client-go/v2/txnkv/txnsnapshot" "github.com/tikv/client-go/v2/txnkv/txnsnapshot"
"github.com/tikv/client-go/v2/util" "github.com/tikv/client-go/v2/util"
pd "github.com/tikv/pd/client" pd "github.com/tikv/pd/client"
resourceControlClient "github.com/tikv/pd/pkg/mcs/resource_manager/client"
clientv3 "go.etcd.io/etcd/client/v3" clientv3 "go.etcd.io/etcd/client/v3"
atomicutil "go.uber.org/atomic" atomicutil "go.uber.org/atomic"
"go.uber.org/zap" "go.uber.org/zap"
@ -582,6 +583,26 @@ func (s *KVStore) updateSafeTS(ctx context.Context) {
wg.Wait() wg.Wait()
} }
// EnableResourceControl enables the resource control.
func (s *KVStore) EnableResourceControl() {
client.EnableResourceControl()
}
// DisableResourceControl disables the resource control.
func (s *KVStore) DisableResourceControl() {
client.DisableResourceControl()
}
// SetResourceControlInterceptor sets the interceptor for resource control.
func (s *KVStore) SetResourceControlInterceptor(interceptor resourceControlClient.ResourceGroupKVInterceptor) {
client.SetResourceControlInterceptor(interceptor)
}
// UnsetResourceControlInterceptor un-sets the interceptor for resource control.
func (s *KVStore) UnsetResourceControlInterceptor() {
client.SetResourceControlInterceptor(nil)
}
// Variables defines the variables used by TiKV storage. // Variables defines the variables used by TiKV storage.
type Variables = kv.Variables type Variables = kv.Variables