xds: use locality from the connected address for load reporting

balancer/balancer.go

@@ -72,8 +72,20 @@ func unregisterForTesting(name string) {
  delete(m, name)
 }
 
+// getConnectedAddress returns the connected address for a SubConnState.
+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 +422,9 @@ 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
 }
 
 // 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"
+ grpcinternal "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 SetConnectedAddress, ok := grpcinternal.SetConnectedAddress.(func(state *balancer.SubConnState, addr resolver.Address)); ok {
+ SetConnectedAddress(&scs, curAddr)
+ }
+ }
+ acbw.stateListener(scs)
  })
 }
 

clientconn.go

@@ -53,7 +53,9 @@ import (
 
 const (
  // minimum time to give a connection to complete
- minConnectTimeout = 20 * time.Second
+ minConnectTimeout = 20 * time.Second
+ withBalancerAttributes = true
+ withoutBalancerAttributes = false
 )
 
 var (
@@ -812,16 +814,26 @@ 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 {
+// addressWithoutBalancerAttributes returns a copy of the input address with
+// the BalancerAttributes field cleared.
+func addressWithoutBalancerAttributes(a resolver.Address) resolver.Address {
+ a.BalancerAttributes = nil
+ return a
+}
+
+// Makes a copy of the input addresses slice and optionally clears out the
+// balancer attributes field. Addresses are passed during subconn creation and
+// address update operations. In both cases, we may clear the balancer
+// attributes by calling this function, which would therefore allow us to use
+// the Equal method provided by the resolver.Address type for comparison.
+func copyAddresses(in []resolver.Address, includeBalancerAttributes bool) []resolver.Address {
  out := make([]resolver.Address, len(in))
  for i := range in {
- out[i] = in[i]
- out[i].BalancerAttributes = nil
+ if includeBalancerAttributes {
+ out[i] = in[i]
+ } else {
+ out[i] = addressWithoutBalancerAttributes(in[i])
+ }
  }
  return out
 }
@@ -837,7 +849,7 @@ func (cc *ClientConn) newAddrConnLocked(addrs []resolver.Address, opts balancer.
  ac := &addrConn{
  state: connectivity.Idle,
  cc: cc,
- addrs: copyAddressesWithoutBalancerAttributes(addrs),
+ addrs: copyAddresses(addrs, withBalancerAttributes),
  scopts: opts,
  dopts: cc.dopts,
  channelz: channelz.RegisterSubChannel(cc.channelz, ""),
@@ -924,12 +936,18 @@ func (ac *addrConn) connect() error {
  return nil
 }
 
-func equalAddresses(a, b []resolver.Address) bool {
+func equalAddressIgnoreBalancerAttributes(a, b resolver.Address) bool {
+ return a.Addr == b.Addr && a.ServerName == b.ServerName &&
+ a.Attributes.Equal(b.Attributes) &&
+ a.Metadata == b.Metadata
+}
+
+func equalAddressesIgnoreBalancerAttributes(a, b []resolver.Address) bool {
  if len(a) != len(b) {
  return false
  }
  for i, v := range a {
- if !v.Equal(b[i]) {
+ if !equalAddressIgnoreBalancerAttributes(v, b[i]) {
  return false
  }
  }
@@ -939,15 +957,15 @@ 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, withBalancerAttributes)
  limit := len(addrs)
  if limit > 5 {
  limit = 5
  }
- channelz.Infof(logger, ac.channelz, "addrConn: updateAddrs addrs (%d of %d): %v", limit, len(addrs), addrs[:limit])
+ channelz.Infof(logger, ac.channelz, "addrConn: updateAddrs addrs (%d of %d): %v", limit, len(addrs), copyAddresses(addrs[:limit], withoutBalancerAttributes))
 
  ac.mu.Lock()
- if equalAddresses(ac.addrs, addrs) {
+ if equalAddressesIgnoreBalancerAttributes(ac.addrs, addrs) {
  ac.mu.Unlock()
  return
  }
@@ -966,7 +984,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 equalAddressIgnoreBalancerAttributes(a, ac.curAddr) {
  // We are connected to a valid address, so do nothing but
  // update the addresses.
  ac.mu.Unlock()
@@ -1214,7 +1232,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
@@ -1347,6 +1365,7 @@ func (ac *addrConn) tryAllAddrs(ctx context.Context, addrs []resolver.Address, c
 // new transport.
 func (ac *addrConn) createTransport(ctx context.Context, addr resolver.Address, copts transport.ConnectOptions, connectDeadline time.Time) error {
  addr.ServerName = ac.cc.getServerName(addr)
+ addrWithoutBalancerAttributes := addressWithoutBalancerAttributes(addr)
  hctx, hcancel := context.WithCancel(ctx)
 
  onClose := func(r transport.GoAwayReason) {
@@ -1381,14 +1400,14 @@ func (ac *addrConn) createTransport(ctx context.Context, addr resolver.Address,
  defer cancel()
  copts.ChannelzParent = ac.channelz
 
- newTr, err := transport.NewClientTransport(connectCtx, ac.cc.ctx, addr, copts, onClose)
+ newTr, err := transport.NewClientTransport(connectCtx, ac.cc.ctx, addrWithoutBalancerAttributes, copts, onClose)
  if err != nil {
  if logger.V(2) {
- logger.Infof("Creating new client transport to %q: %v", addr, err)
+ logger.Infof("Creating new client transport to %q: %v", addrWithoutBalancerAttributes, err)
  }
  // newTr is either nil, or closed.
  hcancel()
- channelz.Warningf(logger, ac.channelz, "grpc: addrConn.createTransport failed to connect to %s. Err: %v", addr, err)
+ channelz.Warningf(logger, ac.channelz, "grpc: addrConn.createTransport failed to connect to %s. Err: %v", addrWithoutBalancerAttributes, err)
  return err
  }
 

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

xds/internal/balancer/clusterimpl/clusterimpl.go

@@ -31,6 +31,7 @@ import (
 
  "google.golang.org/grpc/balancer"
  "google.golang.org/grpc/connectivity"
+ grpcinternal "google.golang.org/grpc/internal"
  "google.golang.org/grpc/internal/balancer/gracefulswitch"
  "google.golang.org/grpc/internal/buffer"
  "google.golang.org/grpc/internal/grpclog"
@@ -366,14 +367,27 @@ func (b *clusterImplBalancer) NewSubConn(addrs []resolver.Address, opts balancer
  lID = xdsinternal.GetLocalityID(newAddrs[i])
  }
  var sc balancer.SubConn
+ ret := &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 GetConnectedAddress, ok := grpcinternal.GetConnectedAddress.(func(state balancer.SubConnState) (resolver.Address, bool)); ok {
+ if addr, ok := GetConnectedAddress(state); ok {
+ // TODO: Why is lID empty when running the test? The locality info is being lost somehow.
+ lID := xdsinternal.GetLocalityID(addr)
+ if !lID.Equal(xdsinternal.LocalityID{}) {
+ ret.updateLocalityID(lID)
+ }
+ }
+ }
+ }
  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.SubConn = sc
  ret.updateLocalityID(lID)
  return ret, nil
 }

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.