fix rpc client panic cause by concurrent close (#1359) (#1397)

close issue #1357 --------- Signed-off-by: crazycs520 <[email protected]>
tikv · Aug 21, 2024 · d28ae4a · d28ae4a
1 parent 48c3e31
commit d28ae4a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/internal/client/client.go b/internal/client/client.go
@@ -650,6 +650,10 @@ func (c *RPCClient) Close() error {
 // CloseAddr closes gRPC connections to the address.
 func (c *RPCClient) CloseAddr(addr string) error {
 	c.Lock()
+	if c.isClosed {
+		c.Unlock()
+		return nil
+	}
 	conn, ok := c.conns[addr]
 	if ok {
 		delete(c.conns, addr)

diff --git a/internal/client/client_test.go b/internal/client/client_test.go
@@ -715,3 +715,23 @@ func TestBatchClientRecoverAfterServerRestart(t *testing.T) {
 		require.NoError(t, err)
 	}
 }
+
+func TestConcurrentCloseConnPanic(t *testing.T) {
+	client := NewRPCClient()
+	addr := "127.0.0.1:6379"
+	_, err := client.getConnArray(addr, true)
+	assert.Nil(t, err)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		err := client.Close()
+		assert.Nil(t, err)
+	}()
+	go func() {
+		defer wg.Done()
+		err := client.CloseAddr(addr)
+		assert.Nil(t, err)
+	}()
+	wg.Wait()
+}