store: fix routine/client leak when get a cached store

store/driver/tikv_driver.go

@@ -145,7 +145,7 @@ func (d *TiKVDriver) setDefaultAndOptions(options ...Option) {
 
 // OpenWithOptions is used by other program that use tidb as a library, to avoid modifying GlobalConfig
 // unspecified options will be set to global config
-func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage, error) {
+func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (resStore kv.Storage, err error) {
 	mc.Lock()
 	defer mc.Unlock()
 	d.setDefaultAndOptions(options...)
@@ -154,7 +154,28 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		return nil, errors.Trace(err)
 	}
 
-	pdCli, err := pd.NewClient(etcdAddrs, pd.SecurityOption{
+	var (
+		pdCli pd.Client
+		spkv  *tikv.EtcdSafePointKV
+		s     *tikv.KVStore
+	)
+	defer func() {
+		if err != nil {
+			if s != nil {
+				// if store is created, it will close spkv and pdCli inside
+				_ = s.Close()
+				return
+			}
+			if spkv != nil {
+				_ = spkv.Close()
+			}
+			if pdCli != nil {
+				pdCli.Close()
+			}
+		}
+	}()
+
+	pdCli, err = pd.NewClient(etcdAddrs, pd.SecurityOption{
 		CAPath:   d.security.ClusterSSLCA,
 		CertPath: d.security.ClusterSSLCert,
 		KeyPath:  d.security.ClusterSSLKey,
@@ -176,6 +197,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 	// FIXME: uuid will be a very long and ugly string, simplify it.
 	uuid := fmt.Sprintf("tikv-%v", pdCli.GetClusterID(context.TODO()))
 	if store, ok := mc.cache[uuid]; ok {
+		pdCli.Close()
 		return store, nil
 	}
 
@@ -184,7 +206,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		return nil, errors.Trace(err)
 	}
 
-	spkv, err := tikv.NewEtcdSafePointKV(etcdAddrs, tlsConfig)
+	spkv, err = tikv.NewEtcdSafePointKV(etcdAddrs, tlsConfig)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
@@ -212,7 +234,7 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		tikv.WithCodec(codec),
 	)
 
-	s, err := tikv.NewKVStore(uuid, pdClient, spkv, rpcClient)
+	s, err = tikv.NewKVStore(uuid, pdClient, spkv, rpcClient)
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

tests/realtikvtest/txntest/isolation_test.go

@@ -17,12 +17,28 @@ package txntest
 import (
 	"testing"
 
+	"github.com/pingcap/tidb/store/driver"
 	"github.com/pingcap/tidb/testkit"
 	"github.com/pingcap/tidb/tests/realtikvtest"
 	"github.com/pingcap/tidb/util"
 	"github.com/stretchr/testify/require"
+	"go.opencensus.io/stats/view"
 )
 
+func TestGetCachedStore(t *testing.T) {
+	defer view.Stop()
+	var d driver.TiKVDriver
+	// when get the cached store, there should not have routine leak.
+	store1, err := d.Open(*realtikvtest.TiKVPath)
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, store1.Close())
+	}()
+	store2, err := d.Open(*realtikvtest.TiKVPath)
+	require.NoError(t, err)
+	require.Equal(t, store1, store2)
+}
+
 /*
 These test cases come from the paper <A Critique of ANSI SQL Isolation Levels>.
 The sign 'P0', 'P1'.... can be found in the paper. These cases will run under snapshot isolation.