From cfee5d73dda3c5410a6d49f3e77074ed9d170375 Mon Sep 17 00:00:00 2001
From: lance6716 <lance6716@gmail.com>
Date: Tue, 23 Jan 2024 16:05:55 +0800
Subject: [PATCH] fix goroutine leak

Signed-off-by: lance6716 <lance6716@gmail.com>
---
 integration_tests/pd_api_test.go |  2 +-
 tikv/kv.go                       | 18 +++++++++++++-----
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/integration_tests/pd_api_test.go b/integration_tests/pd_api_test.go
index 1a19ba52..2da4c4bf 100644
--- a/integration_tests/pd_api_test.go
+++ b/integration_tests/pd_api_test.go
@@ -205,7 +205,7 @@ func (s *apiTestSuite) TestInitClusterMinResolvedTSZero() {
 
 func (s *apiTestSuite) TearDownTest() {
 	if s.store != nil {
-		s.Require().Nil(s.store.Close())
+		s.Require().NoError(s.store.Close())
 	}
 	s.Require().NoError(failpoint.Disable("tikvclient/mockFastSafeTSUpdater"))
 }
diff --git a/tikv/kv.go b/tikv/kv.go
index 2d21ab9e..47006f9e 100644
--- a/tikv/kv.go
+++ b/tikv/kv.go
@@ -198,11 +198,19 @@ func WithPDHTTPClient(
 	opts ...pdhttp.ClientOption,
 ) Option {
 	return func(o *KVStore) {
-		o.pdHttpClient = pdhttp.NewClient(
-			source,
-			pdAddrs,
-			opts...,
-		)
+		if cli := o.GetPDClient(); cli != nil {
+			o.pdHttpClient = pdhttp.NewClientWithServiceDiscovery(
+				source,
+				o.pdClient.GetServiceDiscovery(),
+				opts...,
+			)
+		} else {
+			o.pdHttpClient = pdhttp.NewClient(
+				source,
+				pdAddrs,
+				opts...,
+			)
+		}
 	}
 }