kv: fix some goroutine leak when fail halfway in NewKVStore (#1765)

Signed-off-by: lance6716 <lance6716@gmail.com>

.github/workflows/test.yml

@@ -18,7 +18,7 @@ jobs:
         go-version: 1.23.2
 
     - name: Test
-      run: go test ./...
+      run: go test --tags=intest ./...
 
   race-test:
     runs-on: ubuntu-latest
@@ -31,7 +31,7 @@ jobs:
         go-version: 1.23.2
 
     - name: Test with race
-      run: go test -race ./...
+      run: go test -race --tags=intest ./...
 
   golangci:
     runs-on: ubuntu-latest

tikv/kv.go

@@ -66,6 +66,7 @@ import (
 	"github.com/tikv/client-go/v2/txnkv/txnlock"
 	"github.com/tikv/client-go/v2/txnkv/txnsnapshot"
 	"github.com/tikv/client-go/v2/util"
+	"github.com/tikv/client-go/v2/util/intest"
 	pd "github.com/tikv/pd/client"
 	"github.com/tikv/pd/client/clients/gc"
 	pdhttp "github.com/tikv/pd/client/http"
@@ -316,13 +317,24 @@ func requestHealthFeedbackFromKVClient(ctx context.Context, addr string, tikvCli
 }
 
 // NewKVStore creates a new TiKV store instance.
-func NewKVStore(uuid string, pdClient pd.Client, spkv SafePointKV, tikvclient Client, opt ...Option) (*KVStore, error) {
+func NewKVStore(
+	uuid string,
+	pdClient pd.Client,
+	spkv SafePointKV,
+	tikvclient Client,
+	opt ...Option,
+) (_ *KVStore, retErr error) {
 	o, err := oracles.NewPdOracle(pdClient, &oracles.PDOracleOptions{
 		UpdateInterval: defaultOracleUpdateInterval,
 	})
 	if err != nil {
 		return nil, err
 	}
+	defer func() {
+		if retErr != nil {
+			o.Close()
+		}
+	}()
 	ctx, cancel := context.WithCancel(context.Background())
 	var opts []locate.RegionCacheOpt
 	if config.NextGen {
@@ -349,8 +361,18 @@ func NewKVStore(uuid string, pdClient pd.Client, spkv SafePointKV, tikvclient Cl
 		cancel:                       cancel,
 		gP:                           NewSpool(128, 10*time.Second),
 	}
+	defer func() {
+		if retErr != nil {
+			regionCache.Close()
+		}
+	}()
 
 	txnSafePoint, err := store.loadTxnSafePoint(context.Background())
+	if intest.InTest {
+		if uuid == "TestErrorHalfwayInNewKVStore" {
+			err = errors.New("mock error for TestErrorHalfwayInNewKVStore")
+		}
+	}
 	if err != nil {
 		return nil, errors.WithStack(err)
 	}

tikv/kv_test.go

@@ -24,6 +24,7 @@ import (
 
 	"github.com/pingcap/kvproto/pkg/kvrpcpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
+	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
 	"github.com/tikv/client-go/v2/internal/mockstore/mocktikv"
 	"github.com/tikv/client-go/v2/oracle"
@@ -273,3 +274,9 @@ func (s *testKVSuite) TestMinSafeTsFromMixed2() {
 	s.Require().Equal(mockClient.tikvSafeTs, s.store.GetMinSafeTS("z1"))
 	s.Require().Equal(uint64(10), s.store.GetMinSafeTS("z2"))
 }
+
+func (s *testKVSuite) TestErrorHalfwayInNewKVStore() {
+	// this is a leak test, TestMain will check goroutine leak
+	_, err := NewKVStore("TestErrorHalfwayInNewKVStore", s.store.pdClient, NewMockSafePointKV(), &mocktikv.RPCClient{})
+	require.Error(s.T(), err)
+}