close dropped connections and ping added connections on balancer.Update()

Author: qrort

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Ping new connections on discovery attempt, close dropped ones, so `ydb_go_sdk_ydb_driver_conns` metric is correct
+
 ## v3.108.1
 * Supported `json.Marshaller` query parameter in `database/sql` driver
 

internal/balancer/balancer.go

@@ -186,10 +186,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 }),
@@ -202,6 +204,7 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 	for _, c := range connections {
 		b.pool.Allow(ctx, c)
 		c.Endpoint().Touch()
+		_ = c.Ping(ctx)
 	}
 
 	info := balancerConfig.Info{SelfLocation: localDC}
@@ -213,6 +216,13 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 	}
 
 	b.connectionsState.Store(state)
+
+	for _, e := range dropped {
+		c := b.pool.GetIfPresent(e)
+		if c != nil {
+			_ = c.Close(ctx)
+		}
+	}
 }
 
 func (b *Balancer) Close(ctx context.Context) (err error) {

internal/conn/conn.go

@@ -13,6 +13,7 @@ import (
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/stats"
 
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/closer"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/endpoint"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/meta"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/operation"
@@ -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()

internal/mock/conn.go

@@ -53,6 +53,10 @@ func (c *Conn) Park(ctx context.Context) (err error) {
 	panic("not implemented in mock")
 }
 
+func (c *Conn) Close(_ context.Context) error {
+	return nil
+}
+
 func (c *Conn) Ping(ctx context.Context) error {
 	return c.PingErr
 }