Add support for xDS based EndpointGroup

xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/ClusterEntry.java

@@ -30,16 +30,27 @@
 import com.linecorp.armeria.xds.ClusterSnapshot;
 import com.linecorp.armeria.xds.EndpointSnapshot;
 
+import io.envoyproxy.envoy.config.cluster.v3.Cluster;
+import io.envoyproxy.envoy.config.endpoint.v3.ClusterLoadAssignment;
+
 final class ClusterEntry implements Consumer<List<Endpoint>>, AsyncCloseable {
 
     private final EndpointGroup endpointGroup;
+    private final Cluster cluster;
+    private final ClusterLoadAssignment clusterLoadAssignment;
     private final LoadBalancer loadBalancer;
     private List<Endpoint> endpoints = ImmutableList.of();
 
     ClusterEntry(ClusterSnapshot clusterSnapshot, ClusterManager clusterManager) {
         final EndpointSnapshot endpointSnapshot = clusterSnapshot.endpointSnapshot();
         assert endpointSnapshot != null;
-        loadBalancer = new SubsetLoadBalancer(clusterSnapshot);
+        cluster = clusterSnapshot.xdsResource().resource();
+        clusterLoadAssignment = endpointSnapshot.xdsResource().resource();
+        if (cluster.hasLbSubsetConfig()) {
+            loadBalancer = new SubsetLoadBalancer(clusterSnapshot);
+        } else {
+            loadBalancer = new ZoneAwareLoadBalancer();
+        }
 
         // The order of adding listeners is important
         endpointGroup = XdsEndpointUtil.convertEndpointGroup(clusterSnapshot);
@@ -55,7 +66,12 @@ Endpoint selectNow(ClientRequestContext ctx) {
     @Override
     public void accept(List<Endpoint> endpoints) {
         this.endpoints = ImmutableList.copyOf(endpoints);
-        final PrioritySet prioritySet = new PrioritySet(endpoints);
+        final PriorityStateManager priorityStateManager =
+                new PriorityStateManager(cluster, clusterLoadAssignment);
+        for (Endpoint endpoint: endpoints) {
+            priorityStateManager.registerEndpoint(endpoint);
+        }
+        final PrioritySet prioritySet = priorityStateManager.build();
         loadBalancer.prioritySetUpdated(prioritySet);
     }
 

xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/EndpointGroupUtil.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2024 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.xds.client.endpoint;
+
+import static com.linecorp.armeria.xds.client.endpoint.EndpointUtil.locality;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Map.Entry;
+import java.util.function.Predicate;
+import java.util.stream.Collectors;
+
+import com.google.common.collect.ImmutableMap;
+
+import com.linecorp.armeria.client.Endpoint;
+import com.linecorp.armeria.client.endpoint.EndpointGroup;
+import com.linecorp.armeria.client.endpoint.EndpointSelectionStrategy;
+
+import io.envoyproxy.envoy.config.core.v3.Locality;
+
+final class EndpointGroupUtil {
+
+    static Map<Locality, List<Endpoint>> endpointsByLocality(List<Endpoint> endpoints) {
+        final Map<Locality, List<Endpoint>> endpointsPerLocality = new HashMap<>();
+        for (Endpoint endpoint : endpoints) {
+            endpointsPerLocality.computeIfAbsent(locality(endpoint), ignored -> new ArrayList<>())
+                                .add(endpoint);
+        }
+        return ImmutableMap.copyOf(endpointsPerLocality);
+    }
+
+    static EndpointGroup filter(List<Endpoint> endpoints, EndpointSelectionStrategy strategy,
+                                Predicate<Endpoint> predicate) {
+        final List<Endpoint> filteredEndpoints =
+                endpoints.stream().filter(predicate).collect(Collectors.toList());
+        return EndpointGroup.of(strategy, filteredEndpoints);
+    }
+
+    static EndpointGroup filter(EndpointGroup origEndpointGroup, Predicate<Endpoint> predicate) {
+        return filter(origEndpointGroup.endpoints(), origEndpointGroup.selectionStrategy(), predicate);
+    }
+
+    static Map<Locality, EndpointGroup> filterByLocality(Map<Locality, List<Endpoint>> endpointsMap,
+                                                         EndpointSelectionStrategy strategy,
+                                                         Predicate<Endpoint> predicate) {
+        final ImmutableMap.Builder<Locality, EndpointGroup> filteredLocality = ImmutableMap.builder();
+        for (Entry<Locality, List<Endpoint>> entry: endpointsMap.entrySet()) {
+            final EndpointGroup endpointGroup = filter(entry.getValue(), strategy, predicate);
+            if (endpointGroup.endpoints().isEmpty()) {
+                continue;
+            }
+            filteredLocality.put(entry.getKey(), endpointGroup);
+        }
+        return filteredLocality.build();
+    }
+
+    static Map<Locality, EndpointGroup> filterByLocality(Map<Locality, EndpointGroup> origLocality,
+                                                         Predicate<Endpoint> predicate) {
+        final ImmutableMap.Builder<Locality, EndpointGroup> filteredLocality = ImmutableMap.builder();
+        for (Entry<Locality, EndpointGroup> entry: origLocality.entrySet()) {
+            final EndpointGroup endpointGroup = filter(entry.getValue(), predicate);
+            if (endpointGroup.endpoints().isEmpty()) {
+                continue;
+            }
+            filteredLocality.put(entry.getKey(), endpointGroup);
+        }
+        return filteredLocality.build();
+    }
+
+    private EndpointGroupUtil() {}
+}

xds/src/main/java/com/linecorp/armeria/xds/client/endpoint/EndpointUtil.java

@@ -0,0 +1,137 @@
+/*
+ * Copyright 2024 LINE Corporation
+ *
+ * LINE Corporation licenses this file to you under the Apache License,
+ * version 2.0 (the "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at:
+ *
+ *   https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package com.linecorp.armeria.xds.client.endpoint;
+
+import java.util.concurrent.ThreadLocalRandom;
+
+import com.linecorp.armeria.client.ClientRequestContext;
+import com.linecorp.armeria.client.Endpoint;
+import com.linecorp.armeria.client.endpoint.EndpointSelectionStrategy;
+
+import io.envoyproxy.envoy.config.cluster.v3.Cluster;
+import io.envoyproxy.envoy.config.cluster.v3.Cluster.CommonLbConfig;
+import io.envoyproxy.envoy.config.core.v3.Locality;
+import io.envoyproxy.envoy.config.core.v3.Metadata;
+import io.envoyproxy.envoy.config.endpoint.v3.ClusterLoadAssignment;
+import io.envoyproxy.envoy.config.endpoint.v3.ClusterLoadAssignment.Policy;
+import io.envoyproxy.envoy.config.endpoint.v3.LbEndpoint;
+import io.envoyproxy.envoy.config.endpoint.v3.LocalityLbEndpoints;
+
+final class EndpointUtil {
+
+    static Metadata metadata(Endpoint endpoint) {
+        return lbEndpoint(endpoint).getMetadata();
+    }
+
+    static Locality locality(Endpoint endpoint) {
+        final LocalityLbEndpoints localityLbEndpoints = localityLbEndpoints(endpoint);
+        return localityLbEndpoints.hasLocality() ? localityLbEndpoints.getLocality()
+                                                 : Locality.getDefaultInstance();
+    }
+
+    static CoarseHealth coarseHealth(Endpoint endpoint) {
+        final LbEndpoint lbEndpoint = lbEndpoint(endpoint);
+        switch (lbEndpoint.getHealthStatus()) {
+            // Assume UNKNOWN means health check wasn't performed
+            case UNKNOWN:
+            case HEALTHY:
+                return CoarseHealth.HEALTHY;
+            case DEGRADED:
+                return CoarseHealth.DEGRADED;
+            default:
+                return CoarseHealth.UNHEALTHY;
+        }
+    }
+
+    static int hash(ClientRequestContext ctx) {
+        if (ctx.hasAttr(XdsAttributesKeys.SELECTION_HASH)) {
+            return ctx.attr(XdsAttributesKeys.SELECTION_HASH);
+        }
+        return ThreadLocalRandom.current().nextInt(0, Integer.MAX_VALUE);
+    }
+
+    static int priority(Endpoint endpoint) {
+        return localityLbEndpoints(endpoint).getPriority();
+    }
+
+    static boolean hasLoadBalancingWeight(Endpoint endpoint) {
+        return lbEndpoint(endpoint).hasLoadBalancingWeight();
+    }
+
+    private static LbEndpoint lbEndpoint(Endpoint endpoint) {
+        final LbEndpoint lbEndpoint = endpoint.attr(XdsAttributesKeys.LB_ENDPOINT_KEY);
+        assert lbEndpoint != null;
+        return lbEndpoint;
+    }
+
+    private static LocalityLbEndpoints localityLbEndpoints(Endpoint endpoint) {
+        final LocalityLbEndpoints localityLbEndpoints = endpoint.attr(
+                XdsAttributesKeys.LOCALITY_LB_ENDPOINTS_KEY);
+        assert localityLbEndpoints != null;
+        return localityLbEndpoints;
+    }
+
+    static EndpointSelectionStrategy selectionStrategy(Cluster cluster) {
+        switch (cluster.getLbPolicy()) {
+            case ROUND_ROBIN:
+                return EndpointSelectionStrategy.weightedRoundRobin();
+            case RANDOM:
+                return EndpointSelectionStrategy.roundRobin();
+            case RING_HASH:
+                // implementing this is trivial so it will be done separately
+            default:
+                return EndpointSelectionStrategy.weightedRoundRobin();
+        }
+    }
+
+    static int overProvisionFactor(ClusterLoadAssignment clusterLoadAssignment) {
+        if (!clusterLoadAssignment.hasPolicy()) {
+            return 140;
+        }
+        final Policy policy = clusterLoadAssignment.getPolicy();
+        final int overProvisionFactor;
+        overProvisionFactor =
+                policy.hasOverprovisioningFactor() ? policy.getOverprovisioningFactor().getValue() : 140;
+        return overProvisionFactor;
+    }
+
+    static boolean weightedPriorityHealth(ClusterLoadAssignment clusterLoadAssignment) {
+        final boolean weightedPriorityHealth;
+        final Policy policy = clusterLoadAssignment.getPolicy();
+        weightedPriorityHealth = policy.getWeightedPriorityHealth();
+        return weightedPriorityHealth;
+    }
+
+    static int panicThreshold(Cluster cluster) {
+        if (!cluster.hasCommonLbConfig()) {
+            return 50;
+        }
+        final CommonLbConfig commonLbConfig = cluster.getCommonLbConfig();
+        if (!commonLbConfig.hasHealthyPanicThreshold()) {
+            return 50;
+        }
+        return Math.min((int) Math.round(commonLbConfig.getHealthyPanicThreshold().getValue()), 100);
+    }
+
+    enum CoarseHealth {
+        HEALTHY,
+        DEGRADED,
+        UNHEALTHY,
+    }
+
+    private EndpointUtil() {}
+}