From f73ec0e675f4478674837f3776328d85202d1f22 Mon Sep 17 00:00:00 2001
From: Lei Zhao <zlwgx1023@gmail.com>
Date: Wed, 6 Apr 2022 17:12:03 +0800
Subject: [PATCH] client: fix the race between connArray.Close() and
 connArray.Get() (#465)

Signed-off-by: youjiali1995 <zlwgx1023@gmail.com>
---
 internal/client/client.go       | 11 ++++-------
 internal/client/client_test.go  | 13 +++++++++++++
 internal/locate/region_cache.go |  2 +-
 3 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/internal/client/client.go b/internal/client/client.go
index 95541823..489db48a 100644
--- a/internal/client/client.go
+++ b/internal/client/client.go
@@ -137,7 +137,7 @@ func newConnArray(maxSize uint, addr string, security config.Security, idleNotif
 func (a *connArray) Init(addr string, security config.Security, idleNotify *uint32, enableBatch bool) error {
 	a.target = addr
 
-	opt := grpc.WithInsecure() //nolint
+	opt := grpc.WithInsecure()
 	if len(security.ClusterSSLCA) != 0 {
 		tlsConfig, err := security.ToTLSConfig()
 		if err != nil {
@@ -237,12 +237,9 @@ func (a *connArray) Close() {
 		a.batchConn.Close()
 	}
 
-	for i, c := range a.v {
-		if c != nil {
-			err := c.Close()
-			tikverr.Log(err)
-			a.v[i] = nil
-		}
+	for _, c := range a.v {
+		err := c.Close()
+		tikverr.Log(err)
 	}
 
 	close(a.done)
diff --git a/internal/client/client_test.go b/internal/client/client_test.go
index 8ca7a557..002acd2e 100644
--- a/internal/client/client_test.go
+++ b/internal/client/client_test.go
@@ -51,6 +51,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"github.com/tikv/client-go/v2/config"
 	"github.com/tikv/client-go/v2/tikvrpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -82,6 +83,18 @@ func TestConn(t *testing.T) {
 	assert.Nil(t, conn4)
 }
 
+func TestGetConnAfterClose(t *testing.T) {
+	client := NewRPCClient()
+
+	addr := "127.0.0.1:6379"
+	connArray, err := client.getConnArray(addr, true)
+	assert.Nil(t, err)
+	assert.Nil(t, client.CloseAddr(addr))
+	conn := connArray.Get()
+	state := conn.GetState()
+	assert.True(t, state == connectivity.Shutdown)
+}
+
 func TestCancelTimeoutRetErr(t *testing.T) {
 	req := new(tikvpb.BatchCommandsRequest_Request)
 	a := newBatchConn(1, 1, nil)
diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go
index d71c814e..f2e25ff3 100644
--- a/internal/locate/region_cache.go
+++ b/internal/locate/region_cache.go
@@ -2347,7 +2347,7 @@ func createKVHealthClient(ctx context.Context, addr string) (*grpc.ClientConn, h
 
 	cfg := config.GetGlobalConfig()
 
-	opt := grpc.WithInsecure() //nolint
+	opt := grpc.WithInsecure()
 	if len(cfg.Security.ClusterSSLCA) != 0 {
 		tlsConfig, err := cfg.Security.ToTLSConfig()
 		if err != nil {