address comment

Signed-off-by: Cabinfever_B <[email protected]>
CabinfeverB · Dec 5, 2023 · 0891cf1 · 0891cf1
1 parent 013fbce
commit 0891cf1
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/client/pd_service_discovery.go b/client/pd_service_discovery.go
@@ -210,7 +210,10 @@ func (c *pdServiceClient) GetClientConn() *grpc.ClientConn {
 
 // NeedRetry implements ServiceClient.
 func (c *pdServiceClient) NeedRetry(pdErr *pdpb.Error, err error) bool {
-	return !c.IsLeader()
+	if c.IsLeader() {
+		return false
+	}
+	return !(err == nil && pdErr == nil)
 }
 
 type errFn func(pdErr *pdpb.Error) bool

diff --git a/client/pd_service_discovery_test.go b/client/pd_service_discovery_test.go
@@ -149,6 +149,11 @@ func (suite *serviceClientTestSuite) SetupSuite() {
 }
 
 func (suite *serviceClientTestSuite) TearDownTest() {
+	suite.leaderServer.server.resetCount()
+	suite.followerServer.server.resetCount()
+}
+
+func (suite *serviceClientTestSuite) TearDownSuite() {
 	suite.leaderServer.grpcServer.GracefulStop()
 	suite.followerServer.grpcServer.GracefulStop()
 	suite.clean()