close dropped connection on balancer.Update()

Author: qrort

internal/balancer/balancer.go

@@ -127,10 +127,12 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 		)
 		previous = b.connections().All()
 	)
+
+	_, added, dropped := xslices.Diff(previous, newest, func(lhs, rhs endpoint.Endpoint) int {
+		return strings.Compare(lhs.Address(), rhs.Address())
+	})
+
 	defer func() {
-		_, added, dropped := xslices.Diff(previous, newest, func(lhs, rhs endpoint.Endpoint) int {
-			return strings.Compare(lhs.Address(), rhs.Address())
-		})
 		onDone(
 			xslices.Transform(newest, func(t endpoint.Endpoint) trace.EndpointInfo { return t }),
 			xslices.Transform(added, func(t endpoint.Endpoint) trace.EndpointInfo { return t }),
@@ -145,6 +147,14 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 		c.Endpoint().Touch()
 	}
 
+	for _, e := range dropped {
+		c := b.pool.GetIfPresent(e)
+		if c != nil {
+			b.pool.Ban(ctx, c, errEndpointNotDiscovered)
+			_ = c.Close(ctx)
+		}
+	}
+
 	info := balancerConfig.Info{SelfLocation: localDC}
 	state := newConnectionsState(connections, b.balancerConfig.Filter, info, b.balancerConfig.AllowFallback)
 

internal/balancer/errors.go

@@ -0,0 +1,9 @@
+package balancer
+
+import (
+	"errors"
+)
+
+var (
+	errEndpointNotDiscovered = errors.New("endpoint is no longer discovered")
+)

internal/conn/conn.go

@@ -3,6 +3,7 @@ package conn
 import (
 	"context"
 	"fmt"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -36,6 +37,7 @@ var (
 
 type Conn interface {
 	grpc.ClientConnInterface
+	closer.Closer
 
 	Endpoint() endpoint.Endpoint
 

internal/conn/pool.go

@@ -40,6 +40,20 @@ func (p *Pool) GrpcDialOptions() []grpc.DialOption {
 	return p.dialOptions
 }
 
+func (p *Pool) GetIfPresent(endpoint endpoint.Endpoint) Conn {
+	var (
+		address = endpoint.Address()
+		cc      *conn
+		has     bool
+	)
+
+	if cc, has = p.conns.Get(address); has {
+		return cc
+	}
+
+	return nil
+}
+
 func (p *Pool) Get(endpoint endpoint.Endpoint) Conn {
 	var (
 		address = endpoint.Address()