Better localityWeights handling.

xds/src/main/java/io/grpc/xds/ClusterResolverLoadBalancer.java

@@ -152,7 +152,6 @@ private final class ClusterResolverLbState extends LoadBalancer {
     private final Helper helper;
     private final List<String> clusters = new ArrayList<>();
     private final Map<String, ClusterState> clusterStates = new HashMap<>();
-    private final Map<Locality, Integer> localityWeights = new HashMap<>();
     private PolicySelection endpointLbPolicy;
     private ResolvedAddresses resolvedAddresses;
     private LoadBalancer childLb;
@@ -206,6 +205,8 @@ private void handleEndpointResourceUpdate() {
       List<EquivalentAddressGroup> addresses = new ArrayList<>();
       Map<String, PriorityChildConfig> priorityChildConfigs = new HashMap<>();
       List<String> priorities = new ArrayList<>();  // totally ordered priority list
+      Map<Locality, Integer> localityWeights = new HashMap<>();
+
       Status endpointNotFound = Status.OK;
       for (String cluster : clusters) {
         ClusterState state = clusterStates.get(cluster);
@@ -217,6 +218,7 @@ private void handleEndpointResourceUpdate() {
           addresses.addAll(state.result.addresses);
           priorityChildConfigs.putAll(state.result.priorityChildConfigs);
           priorities.addAll(state.result.priorities);
+          localityWeights.putAll(state.result.localityWeights);
         } else {
           endpointNotFound = state.status;
         }
@@ -248,7 +250,8 @@ private void handleEndpointResourceUpdate() {
               .setLoadBalancingPolicyConfig(childConfig)
               .setAddresses(Collections.unmodifiableList(addresses))
               .setAttributes(resolvedAddresses.getAttributes().toBuilder()
-                  .set(InternalXdsAttributes.ATTR_LOCALITY_WEIGHTS, localityWeights).build())
+                  .set(InternalXdsAttributes.ATTR_LOCALITY_WEIGHTS,
+                      Collections.unmodifiableMap(localityWeights)).build())
               .build());
     }
 
@@ -318,6 +321,7 @@ private abstract class ClusterState {
       // Most recently resolved addresses and config, or null if resource not exists.
       @Nullable
       protected ClusterResolutionResult result;
+
       protected boolean shutdown;
 
       private ClusterState(String name, @Nullable ServerInfo lrsServerInfo,
@@ -377,6 +381,7 @@ public void run() {
             }
             Map<Locality, LocalityLbEndpoints> localityLbEndpoints =
                 update.localityLbEndpointsMap;
+            Map<Locality, Integer> localityWeights = new HashMap<>();
             List<DropOverload> dropOverloads = update.dropPolicies;
             List<EquivalentAddressGroup> addresses = new ArrayList<>();
             Map<String, Map<Locality, Integer>> prioritizedLocalityWeights = new HashMap<>();
@@ -429,7 +434,8 @@ public void run() {
                     endpointLbPolicy, lbRegistry, prioritizedLocalityWeights, dropOverloads);
             status = Status.OK;
             resolved = true;
-            result = new ClusterResolutionResult(addresses, priorityChildConfigs, priorities);
+            result = new ClusterResolutionResult(addresses, priorityChildConfigs, priorities,
+                localityWeights);
             handleEndpointResourceUpdate();
           }
         }
@@ -635,18 +641,23 @@ private static class ClusterResolutionResult {
     private final Map<String, PriorityChildConfig> priorityChildConfigs;
     // List of priority names ordered in descending priorities.
     private final List<String> priorities;
+    // Most recent view on how localities in the cluster should be wighted. Only set for EDS
+    // clusters that support the concept.
+    private final Map<Locality, Integer> localityWeights;
 
     ClusterResolutionResult(List<EquivalentAddressGroup> addresses, String priority,
         PriorityChildConfig config) {
       this(addresses, Collections.singletonMap(priority, config),
-          Collections.singletonList(priority));
+          Collections.singletonList(priority), Collections.emptyMap());
     }
 
     ClusterResolutionResult(List<EquivalentAddressGroup> addresses,
-        Map<String, PriorityChildConfig> configs, List<String> priorities) {
+        Map<String, PriorityChildConfig> configs, List<String> priorities,
+        Map<Locality, Integer> localityWeights) {
       this.addresses = addresses;
       this.priorityChildConfigs = configs;
       this.priorities = priorities;
+      this.localityWeights = localityWeights;
     }
   }