From 0172e2693297ecfd18c73504ebdcaa6f1051b9f3 Mon Sep 17 00:00:00 2001
From: lance6716 <lance6716@gmail.com>
Date: Sun, 28 Sep 2025 18:50:07 +0800
Subject: [PATCH] kv: fix some goroutine leak when fail halfway in NewKVStore
 (#1765)

Signed-off-by: lance6716 <lance6716@gmail.com>
---
 .github/workflows/test.yml |  4 ++--
 tikv/kv.go                 | 24 +++++++++++++++++++++++-
 tikv/kv_test.go            |  7 +++++++
 3 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index a4d16b2d..f0860098 100644
--- a/.github/workflows/test.yml
+++ b/.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
diff --git a/tikv/kv.go b/tikv/kv.go
index f5e2c93e..f3e3b28a 100644
--- a/tikv/kv.go
+++ b/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)
 	}
diff --git a/tikv/kv_test.go b/tikv/kv_test.go
index 1a9a4618..8e609c49 100644
--- a/tikv/kv_test.go
+++ b/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)
+}