kvclient: fix the connection array race (#39706)

ref #33773
pingcap · Dec 7, 2022 · 825f4c8 · 825f4c8
1 parent 5d411f5
commit 825f4c8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/store/tikv/client/client.go b/store/tikv/client/client.go
@@ -208,12 +208,9 @@ func (a *connArray) Close() {
 		a.batchConn.Close()
 	}
 
-	for i, c := range a.v {
-		if c != nil {
-			err := c.Close()
-			terror.Log(errors.Trace(err))
-			a.v[i] = nil
-		}
+	for _, c := range a.v {
+		err := c.Close()
+		terror.Log(errors.Trace(err))
 	}
 
 	close(a.done)

diff --git a/store/tikv/client/client_test.go b/store/tikv/client/client_test.go
@@ -30,6 +30,7 @@ import (
 	"github.com/pingcap/kvproto/pkg/tikvpb"
 	"github.com/pingcap/tidb/store/tikv/config"
 	"github.com/pingcap/tidb/store/tikv/tikvrpc"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -467,3 +468,15 @@ func (s *testClientSuite) TestBatchCommandsBuilder(c *C) {
 	c.Assert(len(builder.forwardingReqs), Equals, 0)
 	c.Assert(builder.idAlloc, Not(Equals), 0)
 }
+
+func (s *testClientSuite) TestGetConnAfterClose(c *C) {
+	client := NewRPCClient(config.Security{})
+
+	addr := "127.0.0.1:6379"
+	connArray, err := client.getConnArray(addr, true)
+	c.Assert(err, IsNil)
+	c.Assert(client.Close(), IsNil)
+	conn := connArray.Get()
+	state := conn.GetState()
+	c.Assert(state, Equals, connectivity.Shutdown)
+}