Add client_interceptor

Signed-off-by: mornyx <mornyx.z@gmail.com>

internal/client/client_collapse.go

@@ -42,7 +42,6 @@ import (
 
 	"github.com/pkg/errors"
 	"github.com/tikv/client-go/v2/tikvrpc"
-	"github.com/tikv/client-go/v2/tikvrpc/interceptor"
 	"golang.org/x/sync/singleflight"
 )
 
@@ -66,15 +65,6 @@ func (r reqCollapse) Close() error {
 }
 
 func (r reqCollapse) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
-	if it := interceptor.GetRPCInterceptorFromCtx(ctx); it != nil {
-		return it(func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
-			return r.sendRequest(ctx, target, req, timeout)
-		})(addr, req)
-	}
-	return r.sendRequest(ctx, addr, req, timeout)
-}
-
-func (r reqCollapse) sendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
 	if r.Client == nil {
 		panic("client should not be nil")
 	}

internal/client/client_interceptor.go

@@ -0,0 +1,43 @@
+// Copyright 2021 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package client
+
+import (
+	"context"
+	"time"
+
+	"github.com/tikv/client-go/v2/tikvrpc"
+	"github.com/tikv/client-go/v2/tikvrpc/interceptor"
+)
+
+var _ Client = interceptedClient{}
+
+type interceptedClient struct {
+	Client
+}
+
+// NewInterceptedClient creates a Client which can execute interceptor.
+func NewInterceptedClient(client Client) Client {
+	return interceptedClient{client}
+}
+
+func (r interceptedClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
+	if it := interceptor.GetRPCInterceptorFromCtx(ctx); it != nil {
+		return it(func(target string, req *tikvrpc.Request) (*tikvrpc.Response, error) {
+			return r.Client.SendRequest(ctx, target, req, timeout)
+		})(addr, req)
+	}
+	return r.Client.SendRequest(ctx, addr, req, timeout)
+}

tikv/kv.go

@@ -180,7 +180,7 @@ func NewKVStore(uuid string, pdClient pd.Client, spkv SafePointKV, tikvclient Cl
 		ctx:             ctx,
 		cancel:          cancel,
 	}
-	store.clientMu.client = client.NewReqCollapse(tikvclient)
+	store.clientMu.client = client.NewReqCollapse(client.NewInterceptedClient(tikvclient))
 	store.lockResolver = txnlock.NewLockResolver(store)
 
 	store.wg.Add(2)

tikvrpc/interceptor/interceptor.go

@@ -59,9 +59,9 @@ import (
 // ```
 //
 // NOTE: Interceptor calls may not correspond one-to-one with the underlying gRPC requests.
-// This is because there may be some exceptions, such as: request collapsed, request batched,
-// no valid connection etc. If you have questions about the execution location of RPCInterceptor,
-// please refer to: internal/client/client_collapse.go#SendRequest.
+// This is because there may be some exceptions, such as: request batched, no
+// valid connection etc. If you have questions about the execution location of
+// RPCInterceptor, please refer to: internal/client/client_collapse.go#SendRequest.
 type RPCInterceptor func(next RPCInterceptorFunc) RPCInterceptorFunc
 
 // RPCInterceptorFunc is a callable function used to initiate a request to TiKV.