From a7eab04e69b143d42f00dcb0d10c1981b9a1d127 Mon Sep 17 00:00:00 2001
From: MyonKeminta <MyonKeminta@users.noreply.github.com>
Date: Thu, 5 Nov 2020 16:56:05 +0800
Subject: [PATCH] cherry pick #20808 to release-4.0

Signed-off-by: MyonKeminta <MyonKeminta@users.noreply.github.com>
---
 store/tikv/kv.go        |  4 ++++
 store/tikv/safepoint.go | 11 +++++++++++
 2 files changed, 15 insertions(+)

diff --git a/store/tikv/kv.go b/store/tikv/kv.go
index 3be7398c545dc..9095e0d449bdb 100644
--- a/store/tikv/kv.go
+++ b/store/tikv/kv.go
@@ -351,6 +351,10 @@ func (s *tikvStore) Close() error {
 		s.txnLatches.Close()
 	}
 	s.regionCache.Close()
+
+	if err := s.kv.Close(); err != nil {
+		return errors.Trace(err)
+	}
 	return nil
 }
 
diff --git a/store/tikv/safepoint.go b/store/tikv/safepoint.go
index c6b119fd2af56..0e3d8de44495b 100644
--- a/store/tikv/safepoint.go
+++ b/store/tikv/safepoint.go
@@ -46,6 +46,7 @@ type SafePointKV interface {
 	Put(k string, v string) error
 	Get(k string) (string, error)
 	GetWithPrefix(k string) ([]*mvccpb.KeyValue, error)
+	Close() error
 }
 
 // MockSafePointKV implements SafePointKV at mock test
@@ -90,6 +91,11 @@ func (w *MockSafePointKV) GetWithPrefix(prefix string) ([]*mvccpb.KeyValue, erro
 	return kvs, nil
 }
 
+// Close implements the Close method for SafePointKV
+func (w *MockSafePointKV) Close() error {
+	return nil
+}
+
 // EtcdSafePointKV implements SafePointKV at runtime
 type EtcdSafePointKV struct {
 	cli *clientv3.Client
@@ -140,6 +146,11 @@ func (w *EtcdSafePointKV) GetWithPrefix(k string) ([]*mvccpb.KeyValue, error) {
 	return resp.Kvs, nil
 }
 
+// Close implements the Close for SafePointKV
+func (w *EtcdSafePointKV) Close() error {
+	return errors.Trace(w.cli.Close())
+}
+
 func saveSafePoint(kv SafePointKV, t uint64) error {
 	s := strconv.FormatUint(t, 10)
 	err := kv.Put(GcSavedSafePoint, s)