Apply suggestions from code review

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
Signed-off-by: Yusuke Kato <[email protected]>
Signed-off-by: kpango <[email protected]>

Author: kpango

internal/net/grpc/client.go

@@ -249,7 +249,7 @@ func (g *gRPCClient) StartConnectionMonitor(ctx context.Context) (<-chan error,
 				return ctx.Err()
 			case <-prTick.C:
 				if g.enablePoolRebalance {
-					err = g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+					err = g.rangeConns(ctx, true, func(addr string, p pool.Conn) bool {
 						// if addr or pool is nil or empty the registration of conns is invalid let's disconnect them
 						if addr == "" || p == nil {
 							disconnectTargets = append(disconnectTargets, addr)
@@ -286,7 +286,7 @@ func (g *gRPCClient) StartConnectionMonitor(ctx context.Context) (<-chan error,
 					})
 				}
 			case <-hcTick.C:
-				err = g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+				err = g.rangeConns(ctx, true, func(addr string, p pool.Conn) bool {
 					// if addr or pool is nil or empty the registration of conns is invalid let's disconnect them
 					if addr == "" || p == nil {
 						disconnectTargets = append(disconnectTargets, addr)
@@ -415,7 +415,7 @@ func (g *gRPCClient) Range(
 	if g.conns.Len() == 0 {
 		return errors.ErrGRPCClientConnNotFound("*")
 	}
-	err = g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+	err = g.rangeConns(ctx, false, func(addr string, p pool.Conn) bool {
 		ssctx, sspan := trace.StartSpan(sctx, apiName+"/Client.Range/"+addr)
 		defer func() {
 			if sspan != nil {
@@ -478,7 +478,7 @@ func (g *gRPCClient) RangeConcurrent(
 	if g.conns.Len() == 0 {
 		return errors.ErrGRPCClientConnNotFound("*")
 	}
-	err = g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+	err = g.rangeConns(ctx, false, func(addr string, p pool.Conn) bool {
 		eg.Go(safety.RecoverFunc(func() (err error) {
 			ssctx, sspan := trace.StartSpan(egctx, apiName+"/Client.RangeConcurrent/"+addr)
 			defer func() {
@@ -701,7 +701,7 @@ func (g *gRPCClient) RoundRobin(
 	}
 
 	do := func() (data any, err error) {
-		cerr := g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+		cerr := g.rangeConns(ctx, false, func(addr string, p pool.Conn) bool {
 			select {
 			case <-ctx.Done():
 				err = ctx.Err()
@@ -1087,7 +1087,7 @@ func (g *gRPCClient) Disconnect(ctx context.Context, addr string) error {
 
 func (g *gRPCClient) ConnectedAddrs(ctx context.Context) (addrs []string) {
 	addrs = make([]string, 0, g.conns.Len())
-	err := g.rangeConns(ctx, func(addr string, p pool.Conn) bool {
+	err := g.rangeConns(ctx, false, func(addr string, p pool.Conn) bool {
 		if p != nil && p.IsHealthy(ctx) {
 			addrs = append(addrs, addr)
 		}
@@ -1118,14 +1118,20 @@ func (g *gRPCClient) Close(ctx context.Context) (err error) {
 	return err
 }
 
-func (g *gRPCClient) rangeConns(ctx context.Context, fn func(addr string, p pool.Conn) bool) error {
+func (g *gRPCClient) rangeConns(ctx context.Context, force bool, fn func(addr string, p pool.Conn) bool) error {
 	var cnt int
 	g.conns.Range(func(addr string, p pool.Conn) bool {
+		if force {
+			cnt++
+			return fn(addr, p)
+		}
 		if p == nil || !p.IsHealthy(ctx) {
 			pc, err := p.Connect(ctx)
 			if pc == nil || err != nil || !pc.IsHealthy(ctx) {
 				if pc != nil {
-					pc.Disconnect(ctx)
+					if derr := pc.Disconnect(ctx); derr != nil {
+						log.Debugf("Failed to disconnect unhealthy connection for %s: %v", addr, derr)
+					}
 				}
 				log.Debugf("Unhealthy connection detected for %s during gRPC Connection Range over Loop:\t%s", addr, p.String())
 				return true

internal/net/grpc/client_test.go

@@ -3107,7 +3107,7 @@ package grpc
 // 				stopMonitor:            test.fields.stopMonitor,
 // 			}
 //
-// 			gotAddrs := g.ConnectedAddrs(context.Background)
+			gotAddrs := g.ConnectedAddrs(context.Background())
 // 			if err := checkFunc(test.want, gotAddrs); err != nil {
 // 				tt.Errorf("error = %v", err)
 // 			}