WIP

Author: asmyasnikov

internal/balancer/balancer.go

@@ -42,7 +42,7 @@ type Balancer struct {
 	discoveryRepeater repeater.Repeater
 	localDCDetector   func(ctx context.Context, endpoints []endpoint.Endpoint) (string, error)
 
-	connectionsState atomic.Pointer[connectionsState]
+	connectionsState atomic.Pointer[state]
 
 	mu                         xsync.RWMutex
 	onApplyDiscoveredEndpoints []func(ctx context.Context, endpoints []endpoint.Info)
@@ -152,7 +152,7 @@ func (b *Balancer) applyDiscoveredEndpoints(ctx context.Context, newest []endpoi
 	}
 
 	info := balancerConfig.Info{SelfLocation: localDC}
-	state := newConnectionsState(connections, b.config.Filter, info, b.config.AllowFallback)
+	state := newState(connections, b.config.Filter, info, b.config.AllowFallback)
 
 	endpointsInfo := make([]endpoint.Info, len(newest))
 	for i, e := range newest {
@@ -319,7 +319,7 @@ func (b *Balancer) wrapCall(ctx context.Context, f func(ctx context.Context, cc
 	return nil
 }
 
-func (b *Balancer) connections() *connectionsState {
+func (b *Balancer) connections() *state {
 	return b.connectionsState.Load()
 }
 
@@ -351,7 +351,7 @@ func (b *Balancer) getConn(ctx context.Context) (c conn.Conn, err error) {
 		}
 	}()
 
-	c, failedCount = state.GetConnection(ctx)
+	c, failedCount = state.Next(ctx)
 	if c == nil {
 		return nil, xerrors.WithStackTrace(
 			fmt.Errorf("%w: cannot get connection from Balancer after %d attempts", ErrNoEndpoints, failedCount),