xds: use locality from the connected address for load reporting

balancer/balancer.go

@@ -72,8 +72,21 @@ func unregisterForTesting(name string) {
 	delete(m, name)
 }
 
+// getConnectedAddress returns the connected address for a SubConnState and
+// whether or not it is valid.
+func getConnectedAddress(scs SubConnState) (resolver.Address, bool) {
+	return scs.connectedAddress, scs.ConnectivityState == connectivity.Ready
+}
+
+// setConnectedAddress sets the connected address for a SubConnState.
+func setConnectedAddress(scs *SubConnState, addr resolver.Address) {
+	scs.connectedAddress = addr
+}
+
 func init() {
 	internal.BalancerUnregister = unregisterForTesting
+	internal.GetConnectedAddress = getConnectedAddress
+	internal.SetConnectedAddress = setConnectedAddress
 }
 
 // Get returns the resolver builder registered with the given name.
@@ -410,6 +423,15 @@ type SubConnState struct {
 	// ConnectionError is set if the ConnectivityState is TransientFailure,
 	// describing the reason the SubConn failed.  Otherwise, it is nil.
 	ConnectionError error
+	// connectedAddr contains the connected address when ConnectivityState is
+	// Ready. Otherwise, it is indeterminate.
+	connectedAddress resolver.Address
+}
+
+func (lhs SubConnState) Equal(rhs SubConnState) bool {
+	return lhs.ConnectivityState == rhs.ConnectivityState &&
+		lhs.ConnectionError == rhs.ConnectionError &&
+		lhs.connectedAddress.Addr == rhs.connectedAddress.Addr
 }
 
 // ClientConnState describes the state of a ClientConn relevant to the

balancer_wrapper.go

@@ -25,6 +25,7 @@ import (
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/internal/channelz"
 	"google.golang.org/grpc/internal/grpcsync"
@@ -252,15 +253,21 @@ type acBalancerWrapper struct {
 
 // updateState is invoked by grpc to push a subConn state update to the
 // underlying balancer.
-func (acbw *acBalancerWrapper) updateState(s connectivity.State, err error) {
+func (acbw *acBalancerWrapper) updateState(s connectivity.State, curAddr resolver.Address, err error) {
 	acbw.ccb.serializer.Schedule(func(ctx context.Context) {
 		if ctx.Err() != nil || acbw.ccb.balancer == nil {
 			return
 		}
 		// Even though it is optional for balancers, gracefulswitch ensures
 		// opts.StateListener is set, so this cannot ever be nil.
 		// TODO: delete this comment when UpdateSubConnState is removed.
-		acbw.stateListener(balancer.SubConnState{ConnectivityState: s, ConnectionError: err})
+		scs := balancer.SubConnState{ConnectivityState: s, ConnectionError: err}
+		if s == connectivity.Ready {
+			if sca, ok := internal.SetConnectedAddress.(func(*balancer.SubConnState, resolver.Address)); ok {
+				sca(&scs, curAddr)
+			}
+		}
+		acbw.stateListener(scs)
 	})
 }
 

clientconn.go

@@ -812,17 +812,11 @@ func (cc *ClientConn) applyFailingLBLocked(sc *serviceconfig.ParseResult) {
 	cc.csMgr.updateState(connectivity.TransientFailure)
 }
 
-// Makes a copy of the input addresses slice and clears out the balancer
-// attributes field. Addresses are passed during subconn creation and address
-// update operations. In both cases, we will clear the balancer attributes by
-// calling this function, and therefore we will be able to use the Equal method
-// provided by the resolver.Address type for comparison.
-func copyAddressesWithoutBalancerAttributes(in []resolver.Address) []resolver.Address {
+// Makes a copy of the input addresses slice. Addresses are passed during
+// subconn creation and address update operations.
+func copyAddresses(in []resolver.Address) []resolver.Address {
 	out := make([]resolver.Address, len(in))
-	for i := range in {
-		out[i] = in[i]
-		out[i].BalancerAttributes = nil
-	}
+	copy(out, in)
 	return out
 }
 
@@ -837,7 +831,7 @@ func (cc *ClientConn) newAddrConnLocked(addrs []resolver.Address, opts balancer.
 	ac := &addrConn{
 		state:        connectivity.Idle,
 		cc:           cc,
-		addrs:        copyAddressesWithoutBalancerAttributes(addrs),
+		addrs:        copyAddresses(addrs),
 		scopts:       opts,
 		dopts:        cc.dopts,
 		channelz:     channelz.RegisterSubChannel(cc.channelz, ""),
@@ -924,12 +918,18 @@ func (ac *addrConn) connect() error {
 	return nil
 }
 
+func equalAddress(a, b resolver.Address) bool {
+	return a.Addr == b.Addr && a.ServerName == b.ServerName &&
+		a.Attributes.Equal(b.Attributes) &&
+		a.Metadata == b.Metadata
+}
+
 func equalAddresses(a, b []resolver.Address) bool {
 	if len(a) != len(b) {
 		return false
 	}
 	for i, v := range a {
-		if !v.Equal(b[i]) {
+		if !equalAddress(v, b[i]) {
 			return false
 		}
 	}
@@ -939,7 +939,7 @@ func equalAddresses(a, b []resolver.Address) bool {
 // updateAddrs updates ac.addrs with the new addresses list and handles active
 // connections or connection attempts.
 func (ac *addrConn) updateAddrs(addrs []resolver.Address) {
-	addrs = copyAddressesWithoutBalancerAttributes(addrs)
+	addrs = copyAddresses(addrs)
 	limit := len(addrs)
 	if limit > 5 {
 		limit = 5
@@ -966,7 +966,7 @@ func (ac *addrConn) updateAddrs(addrs []resolver.Address) {
 		// Try to find the connected address.
 		for _, a := range addrs {
 			a.ServerName = ac.cc.getServerName(a)
-			if a.Equal(ac.curAddr) {
+			if equalAddress(a, ac.curAddr) {
 				// We are connected to a valid address, so do nothing but
 				// update the addresses.
 				ac.mu.Unlock()
@@ -1214,7 +1214,7 @@ func (ac *addrConn) updateConnectivityState(s connectivity.State, lastErr error)
 	} else {
 		channelz.Infof(logger, ac.channelz, "Subchannel Connectivity change to %v, last error: %s", s, lastErr)
 	}
-	ac.acbw.updateState(s, lastErr)
+	ac.acbw.updateState(s, ac.curAddr, lastErr)
 }
 
 // adjustParams updates parameters used to create transports upon

internal/internal.go

@@ -208,6 +208,13 @@ var (
 	// ShuffleAddressListForTesting pseudo-randomizes the order of addresses.  n
 	// is the number of elements.  swap swaps the elements with indexes i and j.
 	ShuffleAddressListForTesting any // func(n int, swap func(i, j int))
+
+	// GetConnectedAddress returns the connected address for a SubConnState and
+	// whether the address is valid based on the state.
+	GetConnectedAddress any // func (scs SubConnState) (resolver.Address, bool)
+
+	// SetConnectedAddress sets the connected address for a SubConnState.
+	SetConnectedAddress any // func(scs *SubConnState, addr resolver.Address)
 )
 
 // HealthChecker defines the signature of the client-side LB channel health

internal/testutils/balancer.go

@@ -26,6 +26,7 @@ import (
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/grpcsync"
 	"google.golang.org/grpc/resolver"
 )
@@ -76,6 +77,21 @@ func (tsc *TestSubConn) UpdateState(state balancer.SubConnState) {
 	}
 }
 
+// UpdateStateAndConnectedAddress saves the connected address to state if the
+// connectivity state is Ready and pushes the state to the listener if one is
+// registered.
+func (tsc *TestSubConn) UpdateStateAndConnectedAddress(state balancer.SubConnState, addr resolver.Address) {
+	<-tsc.connectCalled.Done()
+	if state.ConnectivityState == connectivity.Ready {
+		sca := internal.SetConnectedAddress.(func(*balancer.SubConnState, resolver.Address))
+		sca(&state, addr)
+	}
+	if tsc.stateListener != nil {
+		tsc.stateListener(state)
+		return
+	}
+}
+
 // Shutdown pushes the SubConn to the ShutdownSubConn channel in the parent
 // TestClientConn.
 func (tsc *TestSubConn) Shutdown() {

xds/internal/balancer/clusterimpl/balancer_test.go

@@ -637,7 +637,7 @@ func (s) TestLoadReporting(t *testing.T) {
 		t.Fatal(err.Error())
 	}
 
-	sc1.UpdateState(balancer.SubConnState{ConnectivityState: connectivity.Ready})
+	sc1.UpdateStateAndConnectedAddress(balancer.SubConnState{ConnectivityState: connectivity.Ready}, addrs[0])
 	// Test pick with one backend.
 	const successCount = 5
 	const errorCount = 5

xds/internal/balancer/clusterimpl/clusterimpl.go

@@ -31,6 +31,7 @@
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
+	"google.golang.org/grpc/internal"
 	"google.golang.org/grpc/internal/balancer/gracefulswitch"
 	"google.golang.org/grpc/internal/buffer"
 	"google.golang.org/grpc/internal/grpclog"
@@ -360,22 +361,33 @@
 func (b *clusterImplBalancer) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
 	clusterName := b.getClusterName()
 	newAddrs := make([]resolver.Address, len(addrs))
-	var lID xdsinternal.LocalityID
 	for i, addr := range addrs {
 		newAddrs[i] = xds.SetXDSHandshakeClusterName(addr, clusterName)
-		lID = xdsinternal.GetLocalityID(newAddrs[i])
 	}
 	var sc balancer.SubConn
+	scw := &scWrapper{}
 	oldListener := opts.StateListener
-	opts.StateListener = func(state balancer.SubConnState) { b.updateSubConnState(sc, state, oldListener) }
+	opts.StateListener = func(state balancer.SubConnState) {
+		b.updateSubConnState(sc, state, oldListener)
+		// Read connected address and call updateLocalityID() based on the connected
+		// address's locality. https://github.com/grpc/grpc-go/issues/7339
+		if gca, ok := internal.GetConnectedAddress.(func(balancer.SubConnState) (resolver.Address, bool)); ok {
+			if addr, ok := gca(state); ok {
+				lID := xdsinternal.GetLocalityID(addr)
+				if !lID.Empty() {
+					scw.updateLocalityID(lID)
+				} else if b.logger.V(2) {
+					b.logger.Infof("Locality ID for %v unexpectedly empty", addr)
+				}
+			}
+		}
+	}
 	sc, err := b.ClientConn.NewSubConn(newAddrs, opts)
 	if err != nil {
 		return nil, err
 	}
-	// Wrap this SubConn in a wrapper, and add it to the map.
-	ret := &scWrapper{SubConn: sc}
-	ret.updateLocalityID(lID)
-	return ret, nil
+	scw.SubConn = sc
+	return scw, nil
 }
 
 func (b *clusterImplBalancer) RemoveSubConn(sc balancer.SubConn) {

xds/internal/balancer/clusterimpl/tests/balancer_test.go

@@ -310,14 +310,9 @@ func (s) TestLoadReportingPickFirstMultiLocality(t *testing.T) {
 	}
 	mgmtServer.LRSServer.LRSResponseChan <- &resp
 
-	// Wait for load to be reported for locality of server 2.
-	// We (incorrectly) wait for load report for region-2 because presently
-	// pickfirst always reports load for the locality of the last address in the
-	// subconn. This will be fixed by ensuring there is only one address per
-	// subconn.
-	// TODO(#7339): Change region to region-1 once fixed.
-	if err := waitForSuccessfulLoadReport(ctx, mgmtServer.LRSServer, "region-2"); err != nil {
-		t.Fatalf("region-2 did not receive load due to error: %v", err)
+	// Wait for load to be reported for locality of server 1.
+	if err := waitForSuccessfulLoadReport(ctx, mgmtServer.LRSServer, "region-1"); err != nil {
+		t.Fatalf("Server 1 did not receive load due to error: %v", err)
 	}
 
 	// Stop server 1 and send one more rpc. Now the request should go to server 2.

xds/internal/internal.go

@@ -55,6 +55,11 @@ func (l LocalityID) Equal(o any) bool {
 	return l.Region == ol.Region && l.Zone == ol.Zone && l.SubZone == ol.SubZone
 }
 
+// Empty returns whether or not the locality ID is empty.
+func (l LocalityID) Empty() bool {
+	return len(l.Region) == 0 && len(l.Zone) == 0 && len(l.SubZone) == 0
+}
+
 // LocalityIDFromString converts a json representation of locality, into a
 // LocalityID struct.
 func LocalityIDFromString(s string) (ret LocalityID, _ error) {