diff --git a/servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategy.java b/servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategy.java index e8921bd547..38ebe9ed39 100644 --- a/servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategy.java +++ b/servicetalk-loadbalancer-experimental/src/main/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategy.java @@ -20,6 +20,7 @@ import java.util.ArrayList; import java.util.List; +import java.util.TreeMap; import static io.servicetalk.utils.internal.NumberUtils.ensurePositive; @@ -55,16 +56,15 @@ private List rebuildWithPriorities(final List // remote zones intentionally even if all hosts are well. // https://www.envoyproxy.io/docs/envoy/latest/configuration/upstream/cluster_manager/cluster_runtime // #zone-aware-load-balancing - // The behavior could be structured more as a tree, but it's not obvious how to feed such a tree into the load - // balancer. - List groups = new ArrayList<>(); - // First consolidate our hosts into their respective priority groups. + // Consolidate our hosts into their respective priority groups. Since we're going to use a map we must use + // and ordered map (in this case a TreeMap) so that we can iterate in order of group priority. + TreeMap groups = new TreeMap<>(); for (T host : hosts) { if (host.priority() < 0) { LOGGER.warn("Found illegal priority: {}. Dropping priority grouping data.", host.priority()); return hosts; } - Group group = getGroup(groups, host.priority()); + Group group = groups.computeIfAbsent(host.priority(), i -> new Group()); if (host.isHealthy()) { group.healthyCount++; } @@ -78,7 +78,7 @@ private List rebuildWithPriorities(final List // Compute the health percentage for each group. int totalHealthPercentage = 0; - for (Group group : groups) { + for (Group group : groups.values()) { group.healthPercentage = Math.min(100, overProvisionPercentage * group.healthyCount / group.hosts.size()); totalHealthPercentage = Math.min(100, totalHealthPercentage + group.healthPercentage); } @@ -87,43 +87,29 @@ private List rebuildWithPriorities(final List return hosts; } - // We require that we have a continuous priority set. We could relax this if we wanted by using a tree map to - // traverse in order. However, I think it's also a requirement of other xDS compatible implementations. List weightedResults = new ArrayList<>(); int remainingProbability = 100; - for (int i = 0; i < groups.size() && remainingProbability > 0; i++) { - Group group = groups.get(i); - if (group.hosts.isEmpty()) { - // we don't have a continuous priority group. Warn and rebuild without priorities. - LOGGER.warn("Non-continuous priority groups: {} total groups but missing group {}. " + - "Dropping priority grouping data.", groups.size(), i); - return hosts; - } - - // We need to compute the weight for the group. + for (Group group : groups.values()) { + assert !group.hosts.isEmpty(); final int groupProbability = Math.min(remainingProbability, group.healthPercentage * 100 / totalHealthPercentage); - if (groupProbability == 0) { - // TODO: this means all hosts for this group are unhealthy. This may be worth some logging. - } else { + if (groupProbability > 0) { remainingProbability -= groupProbability; group.addToResults(groupProbability, weightedResults); } + if (remainingProbability == 0) { + break; + } } - // What to do if we don't have any healthy nodes at all? - if (remainingProbability > 0) { - // TODO: this is an awkward situation. Should we panic and just discard priority information? + if (weightedResults.isEmpty()) { + // This is awkward situation can happen if we don't have any healthy groups. + // In that case let's panic and return an un-prioritized set of hosts. + LOGGER.warn("No healthy priority groups found. Returning the un-prioritized set."); + return hosts; } return weightedResults; } - private Group getGroup(List groups, int priority) { - while (groups.size() < priority + 1) { - groups.add(new Group()); - } - return groups.get(priority); - } - private static class Group { final List hosts = new ArrayList<>(); int healthyCount; diff --git a/servicetalk-loadbalancer-experimental/src/test/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategyTest.java b/servicetalk-loadbalancer-experimental/src/test/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategyTest.java index e397ea58c6..83ae8c73a4 100644 --- a/servicetalk-loadbalancer-experimental/src/test/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategyTest.java +++ b/servicetalk-loadbalancer-experimental/src/test/java/io/servicetalk/loadbalancer/DefaultHostPriorityStrategyTest.java @@ -94,6 +94,43 @@ void twoPrioritiesWithWeights() { } } + @Test + void twoPrioritiesWithWeightsButSkipNumbers() { + List hosts = makeHosts(6); + for (int i = 0; i < hosts.size(); i++) { + hosts.get(i).loadBalancingWeight(i + 1d); + if (i >= 3) { + hosts.get(i).priority(2); + } + } + List result = hostPriorityStrategy.prioritize(hosts); + assertThat(result.size(), equalTo(3)); + + // We should only have the first three hosts because they were all healthy, so they are the only group. + for (int i = 0; i < 3; i++) { + assertThat(result.get(i).address(), equalTo(hosts.get(i).address())); + // It doesn't matter what they are exactly so long as all weights are equal. + assertThat(result.get(i).loadBalancedWeight(), approxEqual(result.get(0).loadBalancedWeight() * (i + 1))); + } + } + + @Test + void noHealthyNodesDoesntFilterOutElements() { + List hosts = makeHosts(6); + for (int i = 0; i < hosts.size(); i++) { + hosts.get(i).loadBalancingWeight(i + 1d); + hosts.get(i).isHealthy = false; + if (i >= 3) { + hosts.get(i).priority(1); + } + } + List result = hostPriorityStrategy.prioritize(hosts); + assertThat(result.size(), equalTo(6)); + for (int i = 0; i < hosts.size(); i++) { + assertThat(hosts.get(i).loadBalancedWeight, equalTo(i + 1d)); + } + } + @Test void priorityGroupsWithoutUnhealthyNodes() { List hosts = makeHosts(6);