Repository: tikv/client-go
Pull request #268 (2021-08-14 20:36:30)

client: cleanup RPCClient options
Signed-off-by: disksing <i@disksing.com>

use optional security argument so we don't have to pass empty security config every time.


Patch:
diff --git a/internal/client/client.go b/internal/client/client.go
index e8cf1eb..bd619fb 100644
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -240,6 +240,16 @@ func (a *connArray) Close() {
 	close(a.done)
 }
 
+// Opt is the option for the client.
+type Opt func(*RPCClient)
+
+// WithSecurity is used to set the security config.
+func WithSecurity(security config.Security) Opt {
+	return func(c *RPCClient) {
+		c.security = security
+	}
+}
+
 // RPCClient is RPC client struct.
 // TODO: Add flow control between RPC clients in TiDB ond RPC servers in TiKV.
 // Since we use shared client connection to communicate to the same TiKV, it's possible
@@ -261,10 +271,9 @@ type RPCClient struct {
 }
 
 // NewRPCClient creates a client that manages connections and rpc calls with tikv-servers.
-func NewRPCClient(security config.Security, opts ...func(c *RPCClient)) *RPCClient {
+func NewRPCClient(opts ...Opt) *RPCClient {
 	cli := &RPCClient{
 		conns:       make(map[string]*connArray),
-		security:    security,
 		dialTimeout: dialTimeout,
 	}
 	for _, opt := range opts {
diff --git a/rawkv/rawkv.go b/rawkv/rawkv.go
index 8bb86fe..4bf4f84 100644
--- a/rawkv/rawkv.go
+++ b/rawkv/rawkv.go
@@ -96,7 +96,7 @@ func NewClient(ctx context.Context, pdAddrs []string, security config.Security,
 		clusterID:   pdCli.GetClusterID(ctx),
 		regionCache: locate.NewRegionCache(pdCli),
 		pdClient:    pdCli,
-		rpcClient:   client.NewRPCClient(security),
+		rpcClient:   client.NewRPCClient(client.WithSecurity(security)),
 	}, nil
 }
 
diff --git a/tikv/client.go b/tikv/client.go
index e1ac9e6..5a724ac 100644
--- a/tikv/client.go
+++ b/tikv/client.go
@@ -43,6 +43,14 @@ import (
 // It should not be used after calling Close().
 type Client = client.Client
 
+// ClientOpt defines the option to create RPC client.
+type ClientOpt = client.Opt
+
+// WithSecurity is used to set security config.
+func WithSecurity(security config.Security) ClientOpt {
+	return client.WithSecurity(security)
+}
+
 // Timeout durations.
 const (
 	ReadTimeoutMedium = client.ReadTimeoutMedium
@@ -50,6 +58,6 @@ const (
 )
 
 // NewRPCClient creates a client that manages connections and rpc calls with tikv-servers.
-func NewRPCClient(security config.Security, opts ...func(c *client.RPCClient)) *client.RPCClient {
-	return client.NewRPCClient(security, opts...)
+func NewRPCClient(opts ...ClientOpt) *client.RPCClient {
+	return client.NewRPCClient(opts...)
 }
diff --git a/tikv/kv.go b/tikv/kv.go
index 1201449..4a4d80b 100644
--- a/tikv/kv.go
+++ b/tikv/kv.go
@@ -211,7 +211,7 @@ func NewTxnClient(pdAddrs []string) (*KVStore, error) {
 		return nil, errors.Trace(err)
 	}
 
-	s, err := NewKVStore(uuid, pdClient, spkv, NewRPCClient(cfg.Security))
+	s, err := NewKVStore(uuid, pdClient, spkv, NewRPCClient(WithSecurity(cfg.Security)))
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -609,7 +609,7 @@ func NewLockResolver(etcdAddrs []string, security config.Security, opts ...pd.Cl
 		return nil, errors.Trace(err)
 	}
 
-	s, err := NewKVStore(uuid, locate.NewCodeCPDClient(pdCli), spkv, client.NewRPCClient(security))
+	s, err := NewKVStore(uuid, locate.NewCodeCPDClient(pdCli), spkv, client.NewRPCClient(WithSecurity(security)))
 	if err != nil {
 		return nil, errors.Trace(err)
 	}


Test patch:
diff --git a/integration_tests/util_test.go b/integration_tests/util_test.go
index 19101d2..48bb04c 100644
--- a/integration_tests/util_test.go
+++ b/integration_tests/util_test.go
@@ -74,7 +74,7 @@ func NewTestStore(t *testing.T) *tikv.KVStore {
 		require.Nil(t, err)
 		spKV, err := tikv.NewEtcdSafePointKV(addrs, tlsConfig)
 		require.Nil(t, err)
-		store, err := tikv.NewKVStore("test-store", &tikv.CodecPDClient{Client: pdClient}, spKV, tikv.NewRPCClient(securityConfig))
+		store, err := tikv.NewKVStore("test-store", &tikv.CodecPDClient{Client: pdClient}, spKV, tikv.NewRPCClient())
 		require.Nil(t, err)
 		err = clearStorage(store)
 		require.Nil(t, err)
diff --git a/internal/client/client_fail_test.go b/internal/client/client_fail_test.go
index ee4c171..d664026 100644
--- a/internal/client/client_fail_test.go
+++ b/internal/client/client_fail_test.go
@@ -59,9 +59,8 @@ func TestPanicInRecvLoop(t *testing.T) {
 	defer server.Stop()
 
 	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
-	rpcClient := NewRPCClient(config.Security{}, func(c *RPCClient) {
-		c.dialTimeout = time.Second / 3
-	})
+	rpcClient := NewRPCClient()
+	rpcClient.dialTimeout = time.Second / 3
 
 	// Start batchRecvLoop, and it should panic in `failPendingRequests`.
 	_, err := rpcClient.getConnArray(addr, true, func(cfg *config.TiKVClient) { cfg.GrpcConnectionCount = 1 })
@@ -94,7 +93,7 @@ func TestRecvErrorInMultipleRecvLoops(t *testing.T) {
 		conf.TiKVClient.MaxBatchSize = 128
 		conf.TiKVClient.GrpcConnectionCount = 1
 	})()
-	rpcClient := NewRPCClient(config.Security{})
+	rpcClient := NewRPCClient()
 	defer rpcClient.closeConns()
 
 	// Create 4 BatchCommands streams.
diff --git a/internal/client/client_test.go b/internal/client/client_test.go
index 9a4cc44..197c66f 100644
--- a/internal/client/client_test.go
+++ b/internal/client/client_test.go
@@ -59,7 +59,7 @@ func TestConn(t *testing.T) {
 		conf.TiKVClient.MaxBatchSize = 0
 	})()
 
-	client := NewRPCClient(config.Security{})
+	client := NewRPCClient()
 
 	addr := "127.0.0.1:6379"
 	conn1, err := client.getConnArray(addr, true)
@@ -92,7 +92,7 @@ func TestSendWhenReconnect(t *testing.T) {
 	server, port := startMockTikvService()
 	require.True(t, port > 0)
 
-	rpcClient := NewRPCClient(config.Security{})
+	rpcClient := NewRPCClient()
 	addr := fmt.Sprintf("%s:%d", "127.0.0.1", port)
 	conn, err := rpcClient.getConnArray(addr, true)
 	assert.Nil(t, err)
@@ -219,7 +219,7 @@ func TestForwardMetadataByUnaryCall(t *testing.T) {
 		conf.TiKVClient.MaxBatchSize = 0
 		conf.TiKVClient.GrpcConnectionCount = 1
 	})()
-	rpcClient := NewRPCClient(config.Security{})
+	rpcClient := NewRPCClient()
 	defer rpcClient.closeConns()
 
 	var checkCnt uint64
@@ -288,7 +288,7 @@ func TestForwardMetadataByBatchCommands(t *testing.T) {
 		conf.TiKVClient.MaxBatchSize = 128
 		conf.TiKVClient.GrpcConnectionCount = 1
 	})()
-	rpcClient := NewRPCClient(config.Security{})
+	rpcClient := NewRPCClient()
 	defer rpcClient.closeConns()
 
 	var checkCnt uint64
diff --git a/internal/locate/region_request_test.go b/internal/locate/region_request_test.go
index 642f850..2791e27 100644
--- a/internal/locate/region_request_test.go
+++ b/internal/locate/region_request_test.go
@@ -51,7 +51,6 @@ import (
 	"github.com/pingcap/kvproto/pkg/mpp"
 	"github.com/pingcap/kvproto/pkg/tikvpb"
 	"github.com/stretchr/testify/suite"
-	"github.com/tikv/client-go/v2/config"
 	"github.com/tikv/client-go/v2/internal/client"
 	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
 	"github.com/tikv/client-go/v2/internal/retry"
@@ -469,7 +468,7 @@ func (s *testRegionRequestToSingleStoreSuite) TestNoReloadRegionForGrpcWhenCtxCa
 		wg.Done()
 	}()
 
-	cli := client.NewRPCClient(config.Security{})
+	cli := client.NewRPCClient()
 	sender := NewRegionRequestSender(s.cache, cli)
 	req := tikvrpc.NewRequest(tikvrpc.CmdRawPut, &kvrpcpb.RawPutRequest{
 		Key:   []byte("key"),
@@ -486,7 +485,7 @@ func (s *testRegionRequestToSingleStoreSuite) TestNoReloadRegionForGrpcWhenCtxCa
 
 	// Just for covering error code = codes.Canceled.
 	client1 := &cancelContextClient{
-		Client:       client.NewRPCClient(config.Security{}),
+		Client:       client.NewRPCClient(),
 		redirectAddr: addr,
 	}
 	sender = NewRegionRequestSender(s.cache, client1)