tikvclient: refine region-cache

ddl/table_split_test.go

@@ -79,7 +79,7 @@ func checkRegionStartWithTableID(c *C, id int64, store kvStore) {
 		c.Assert(err, IsNil)
 
 		// Region cache may be out of date, so we need to drop this expired region and load it again.
-		cache.DropRegion(loc.Region)
+		cache.InvalidateCachedRegion(loc.Region)
 		if bytes.Equal(loc.StartKey, []byte(regionStartKey)) {
 			return
 		}

store/tikv/coprocessor_test.go

@@ -34,6 +34,7 @@ func (s *testCoprocessorSuite) TestBuildTasks(c *C) {
 	_, regionIDs, _ := mocktikv.BootstrapWithMultiRegions(cluster, []byte("g"), []byte("n"), []byte("t"))
 	pdCli := &codecPDClient{mocktikv.NewPDClient(cluster)}
 	cache := NewRegionCache(pdCli)
+	defer cache.Close()
 
 	bo := NewBackoffer(context.Background(), 3000)
 
@@ -96,6 +97,7 @@ func (s *testCoprocessorSuite) TestSplitRegionRanges(c *C) {
 	mocktikv.BootstrapWithMultiRegions(cluster, []byte("g"), []byte("n"), []byte("t"))
 	pdCli := &codecPDClient{mocktikv.NewPDClient(cluster)}
 	cache := NewRegionCache(pdCli)
+	defer cache.Close()
 
 	bo := NewBackoffer(context.Background(), 3000)
 
@@ -148,6 +150,7 @@ func (s *testCoprocessorSuite) TestRebuild(c *C) {
 	storeID, regionIDs, peerIDs := mocktikv.BootstrapWithMultiRegions(cluster, []byte("m"))
 	pdCli := &codecPDClient{mocktikv.NewPDClient(cluster)}
 	cache := NewRegionCache(pdCli)
+	defer cache.Close()
 	bo := NewBackoffer(context.Background(), 3000)
 
 	tasks, err := buildCopTasks(bo, cache, buildCopRanges("a", "z"), false, false)
@@ -161,7 +164,7 @@ func (s *testCoprocessorSuite) TestRebuild(c *C) {
 	regionIDs = append(regionIDs, cluster.AllocID())
 	peerIDs = append(peerIDs, cluster.AllocID())
 	cluster.Split(regionIDs[1], regionIDs[2], []byte("q"), []uint64{peerIDs[2]}, storeID)
-	cache.DropRegion(tasks[1].region)
+	cache.InvalidateCachedRegion(tasks[1].region)
 
 	tasks, err = buildCopTasks(bo, cache, buildCopRanges("a", "z"), true, false)
 	c.Assert(err, IsNil)

store/tikv/kv.go

@@ -296,6 +296,7 @@ func (s *tikvStore) Close() error {
 	if s.txnLatches != nil {
 		s.txnLatches.Close()
 	}
+	s.regionCache.Close()
 	return nil
 }
 

store/tikv/rawkv.go

@@ -82,7 +82,15 @@ func NewRawKVClient(pdAddrs []string, security config.Security) (*RawKVClient, e
 
 // Close closes the client.
 func (c *RawKVClient) Close() error {
-	c.pdClient.Close()
+	if c.pdClient != nil {
+		c.pdClient.Close()
+	}
+	if c.regionCache != nil {
+		c.regionCache.Close()
+	}
+	if c.rpcClient == nil {
+		return nil
+	}
 	return c.rpcClient.Close()
 }