balancer: set RPC metadata in address attributes, instead of Metadata field

balancer/grpclb/grpclb_remote_balancer.go

@@ -161,19 +161,19 @@ func (lb *lbBalancer) refreshSubConns(backendAddrs []resolver.Address, fallback
 	addrsSet := make(map[resolver.Address]struct{})
 	// Create new SubConns.
 	for _, addr := range backendAddrs {
-		addrWithoutMD := addr
-		addrWithoutMD.Attributes = nil
-		addrsSet[addrWithoutMD] = struct{}{}
-		lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutMD)
+		addrWithoutAttrs := addr
+		addrWithoutAttrs.Attributes = nil
+		addrsSet[addrWithoutAttrs] = struct{}{}
+		lb.backendAddrsWithoutMetadata = append(lb.backendAddrsWithoutMetadata, addrWithoutAttrs)
 
-		if _, ok := lb.subConns[addrWithoutMD]; !ok {
+		if _, ok := lb.subConns[addrWithoutAttrs]; !ok {
 			// Use addrWithMD to create the SubConn.
 			sc, err := lb.cc.NewSubConn([]resolver.Address{addr}, opts)
 			if err != nil {
 				logger.Warningf("grpclb: failed to create new SubConn: %v", err)
 				continue
 			}
-			lb.subConns[addrWithoutMD] = sc // Use the addr without MD as key for the map.
+			lb.subConns[addrWithoutAttrs] = sc // Use the addr without MD as key for the map.
 			if _, ok := lb.scStates[sc]; !ok {
 				// Only set state of new sc to IDLE. The state could already be
 				// READY for cached SubConns.

balancer/grpclb/grpclb_util.go

@@ -124,16 +124,16 @@ func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer
 	if len(addrs) != 1 {
 		return nil, fmt.Errorf("grpclb calling NewSubConn with addrs of length %v", len(addrs))
 	}
-	addrWithoutMD := addrs[0]
-	addrWithoutMD.Attributes = nil
+	addrWithoutAttrs := addrs[0]
+	addrWithoutAttrs.Attributes = nil
 
 	ccc.mu.Lock()
 	defer ccc.mu.Unlock()
-	if entry, ok := ccc.subConnCache[addrWithoutMD]; ok {
+	if entry, ok := ccc.subConnCache[addrWithoutAttrs]; ok {
 		// If entry is in subConnCache, the SubConn was being deleted.
 		// cancel function will never be nil.
 		entry.cancel()
-		delete(ccc.subConnCache, addrWithoutMD)
+		delete(ccc.subConnCache, addrWithoutAttrs)
 		return entry.sc, nil
 	}
 
@@ -142,7 +142,7 @@ func (ccc *lbCacheClientConn) NewSubConn(addrs []resolver.Address, opts balancer
 		return nil, err
 	}
 
-	ccc.subConnToAddr[scNew] = addrWithoutMD
+	ccc.subConnToAddr[scNew] = addrWithoutAttrs
 	return scNew, nil
 }
 

internal/hierarchy/hierarchy.go

@@ -23,7 +23,6 @@
 package hierarchy
 
 import (
-	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/resolver"
 )
 
@@ -37,19 +36,12 @@ func Get(addr resolver.Address) []string {
 	if attrs == nil {
 		return nil
 	}
-	path, ok := attrs.Value(pathKey).([]string)
-	if !ok {
-		return nil
-	}
+	path, _ := attrs.Value(pathKey).([]string)
 	return path
 }
 
 // Set overrides the hierarchical path in addr with path.
 func Set(addr resolver.Address, path []string) resolver.Address {
-	if addr.Attributes == nil {
-		addr.Attributes = attributes.New(pathKey, path)
-		return addr
-	}
 	addr.Attributes = addr.Attributes.WithValues(pathKey, path)
 	return addr
 }

internal/metadata/metadata.go

@@ -22,7 +22,6 @@
 package metadata
 
 import (
-	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/resolver"
 )
@@ -37,10 +36,7 @@ func Get(addr resolver.Address) metadata.MD {
 	if attrs == nil {
 		return nil
 	}
-	md, ok := attrs.Value(mdKey).(metadata.MD)
-	if !ok {
-		return nil
-	}
+	md, _ := attrs.Value(mdKey).(metadata.MD)
 	return md
 }
 
@@ -49,10 +45,6 @@ func Get(addr resolver.Address) metadata.MD {
 // When a SubConn is created with this address, the RPCs sent on it will all
 // have this metadata.
 func Set(addr resolver.Address, md metadata.MD) resolver.Address {
-	if addr.Attributes == nil {
-		addr.Attributes = attributes.New(mdKey, md)
-		return addr
-	}
 	addr.Attributes = addr.Attributes.WithValues(mdKey, md)
 	return addr
 }

test/balancer_test.go

@@ -579,19 +579,16 @@ func (s) TestMetadataInAddressAttributes(t *testing.T) {
 	stub.Register(mdBalancerName, bf)
 	t.Logf("Registered balancer %s...", mdBalancerName)
 
-	r := manual.NewBuilderWithScheme("whatever")
-	t.Logf("Registered manual resolver with scheme %s...", r.Scheme())
-
 	testMDChan := make(chan []string, 1)
 	ss := &stubServer{
 		emptyCall: func(ctx context.Context, _ *testpb.Empty) (*testpb.Empty, error) {
 			md, ok := metadata.FromIncomingContext(ctx)
 			if ok {
 				select {
-				case <-testMDChan:
-				default:
+				case testMDChan <- md[testMDKey]:
+				case <-ctx.Done():
+					return nil, ctx.Err()
 				}
-				testMDChan <- md[testMDKey]
 			}
 			return &testpb.Empty{}, nil
 		},
@@ -601,6 +598,7 @@ func (s) TestMetadataInAddressAttributes(t *testing.T) {
 	}
 	defer ss.Stop()
 
+	r := manual.NewBuilderWithScheme("whatever")
 	dopts := []grpc.DialOption{
 		grpc.WithInsecure(),
 		grpc.WithResolvers(r),