store: fix routine/client leak when get a cached store (#42779) (#46143)

close #42778

DEPS.bzl

@@ -3703,8 +3703,8 @@ def go_deps():
         name = "com_github_tikv_client_go_v2",
         build_file_proto_mode = "disable_global",
         importpath = "github.com/tikv/client-go/v2",
-        sum = "h1:KL+gt28r/52vuHlVk9Q3klULt15LqkJeT01lP59tWGo=",
-        version = "v2.0.4-0.20230817125157-1bf6400f08f5",
+        sum = "h1:xL9X0pr5wZn8onnW6i98sf7KjaZmk0OIQnDhW1z+hRQ=",
+        version = "v2.0.4-0.20230817162610-e5fe1779769d",
     )
     go_repository(
         name = "com_github_tikv_pd_client",

go.mod

@@ -90,7 +90,7 @@ require (
 	github.com/stretchr/testify v1.8.0
 	github.com/tdakkota/asciicheck v0.1.1
 	github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2
-	github.com/tikv/client-go/v2 v2.0.4-0.20230817125157-1bf6400f08f5
+	github.com/tikv/client-go/v2 v2.0.4-0.20230817162610-e5fe1779769d
 	github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07
 	github.com/timakin/bodyclose v0.0.0-20210704033933-f49887972144
 	github.com/twmb/murmur3 v1.1.3

go.sum

@@ -935,8 +935,8 @@ github.com/tenntenn/text/transform v0.0.0-20200319021203-7eef512accb3 h1:f+jULpR
 github.com/tenntenn/text/transform v0.0.0-20200319021203-7eef512accb3/go.mod h1:ON8b8w4BN/kE1EOhwT0o+d62W65a6aPw1nouo9LMgyY=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2 h1:mbAskLJ0oJfDRtkanvQPiooDH8HvJ2FBh+iKT/OmiQQ=
 github.com/tiancaiamao/appdash v0.0.0-20181126055449-889f96f722a2/go.mod h1:2PfKggNGDuadAa0LElHrByyrz4JPZ9fFx6Gs7nx7ZZU=
-github.com/tikv/client-go/v2 v2.0.4-0.20230817125157-1bf6400f08f5 h1:KL+gt28r/52vuHlVk9Q3klULt15LqkJeT01lP59tWGo=
-github.com/tikv/client-go/v2 v2.0.4-0.20230817125157-1bf6400f08f5/go.mod h1:mmVCLP2OqWvQJPOIevQPZvGphzh/oq9vv8J5LDfpadQ=
+github.com/tikv/client-go/v2 v2.0.4-0.20230817162610-e5fe1779769d h1:xL9X0pr5wZn8onnW6i98sf7KjaZmk0OIQnDhW1z+hRQ=
+github.com/tikv/client-go/v2 v2.0.4-0.20230817162610-e5fe1779769d/go.mod h1:mmVCLP2OqWvQJPOIevQPZvGphzh/oq9vv8J5LDfpadQ=
 github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07 h1:ckPpxKcl75mO2N6a4cJXiZH43hvcHPpqc9dh1TmH1nc=
 github.com/tikv/pd/client v0.0.0-20221031025758-80f0d8ca4d07/go.mod h1:CipBxPfxPUME+BImx9MUYXCnAVLS3VJUr3mnSJwh40A=
 github.com/timakin/bodyclose v0.0.0-20210704033933-f49887972144 h1:kl4KhGNsJIbDHS9/4U9yQo1UcPQM0kOMJHn29EoH/Ro=

store/driver/tikv_driver.go

@@ -113,7 +113,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...)
@@ -122,7 +122,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,
@@ -135,15 +156,16 @@ func (d TiKVDriver) OpenWithOptions(path string, options ...Option) (kv.Storage,
 		),
 		pd.WithCustomTimeoutOption(time.Duration(d.pdConfig.PDServerTimeout)*time.Second),
 		pd.WithForwardingOption(config.GetGlobalConfig().EnableForwarding))
-	pdCli = util.InterceptedPDClient{Client: pdCli}
 
 	if err != nil {
 		return nil, errors.Trace(err)
 	}
+	pdCli = util.InterceptedPDClient{Client: pdCli}
 
 	// 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
 	}
 
@@ -152,13 +174,13 @@ 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)
 	}
 
 	pdClient := tikv.CodecPDClient{Client: pdCli}
-	s, err := tikv.NewKVStore(uuid, &pdClient, spkv, tikv.NewRPCClient(tikv.WithSecurity(d.security)), tikv.WithPDHTTPClient(tlsConfig, etcdAddrs))
+	s, err = tikv.NewKVStore(uuid, &pdClient, spkv, tikv.NewRPCClient(tikv.WithSecurity(d.security)), tikv.WithPDHTTPClient(tlsConfig, etcdAddrs))
 	if err != nil {
 		return nil, errors.Trace(err)
 	}

tests/realtikvtest/testkit.go

@@ -40,8 +40,12 @@ import (
 	"go.uber.org/goleak"
 )
 
-// WithRealTiKV is a flag identify whether tests run with real TiKV
-var WithRealTiKV = flag.Bool("with-real-tikv", false, "whether tests run with real TiKV")
+var (
+	// WithRealTiKV is a flag identify whether tests run with real TiKV
+	WithRealTiKV = flag.Bool("with-real-tikv", false, "whether tests run with real TiKV")
+	// TiKVPath is the path of the TiKV Storage.
+	TiKVPath = flag.String("tikv-path", "tikv://127.0.0.1:2379?disableGC=true", "TiKV addr")
+)
 
 // RunTestMain run common setups for all real tikv tests.
 func RunTestMain(m *testing.M) {

tests/realtikvtest/txntest/BUILD.bazel

@@ -16,10 +16,12 @@ go_test(
         "//kv",
         "//parser",
         "//session/txninfo",
+        "//store/driver",
         "//testkit",
         "//tests/realtikvtest",
         "//util",
         "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
+        "@io_opencensus_go//stats/view",
     ],
 )

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.