close dropped connection on balancer.Update()

Author: qrort

internal/balancer/balancer.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 	"sync/atomic"
 
 	"github.com/ydb-platform/ydb-go-genproto/Ydb_Discovery_V1"
@@ -127,10 +128,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 +148,14 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 		c.Endpoint().Touch()
 	}
 
+	droppedConnections := endpointsToConnections(b.pool, dropped)
+	for _, c := range droppedConnections {
+		sync.OnceFunc(func() {
+			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