Make zone load balancing logic depend on local_cluster host distribution.

include/envoy/upstream/upstream.h

@@ -192,10 +192,12 @@ class HostSet {
   COUNTER(update_success)                                                                          \
   COUNTER(update_failure)                                                                          \
   COUNTER(zone_cluster_too_small)                                                                  \
-  COUNTER(zone_over_percentage)                                                                    \
+  COUNTER(zone_routing_all_directly)                                                               \
   COUNTER(zone_routing_sampled)                                                                    \
-  COUNTER(zone_routing_no_sampled)                                                                 \
-  GAUGE  (max_host_weight)
+  COUNTER(zone_routing_cross_zone)                                                                 \
+  GAUGE  (max_host_weight)                                                                         \
+  COUNTER(local_cluster_not_ok)                                                                    \
+  COUNTER(zone_number_differs)
 // clang-format on
 
 /**

source/common/upstream/load_balancer_impl.cc

@@ -29,13 +29,26 @@ bool LoadBalancerBase::earlyExitNonZoneRouting() {
     return true;
   }
 
+  // If local cluster is not set, or we are in panic mode for it.
+  if (local_host_set_ == nullptr || local_host_set_->hosts().empty() ||
+      isGlobalPanic(*local_host_set_)) {
+    stats_.local_cluster_not_ok_.inc();
+    return true;
+  }
+
+  // Same number of zones should be for local and upstream cluster.
+  if (host_set_.healthyHostsPerZone().size() != local_host_set_->healthyHostsPerZone().size()) {
+    stats_.zone_number_differs_.inc();
+    return true;
+  }
+
   return false;
 }
 
-bool LoadBalancerBase::isGlobalPanic() {
+bool LoadBalancerBase::isGlobalPanic(const HostSet& host_set) {
   uint64_t global_panic_threshold =
       std::min(100UL, runtime_.snapshot().getInteger("upstream.healthy_panic_threshold", 50));
-  double healthy_percent = 100.0 * host_set_.healthyHosts().size() / host_set_.hosts().size();
+  double healthy_percent = 100.0 * host_set.healthyHosts().size() / host_set.hosts().size();
 
   // If the % of healthy hosts in the cluster is less than our panic threshold, we use all hosts.
   if (healthy_percent < global_panic_threshold) {
@@ -46,48 +59,102 @@ bool LoadBalancerBase::isGlobalPanic() {
   return false;
 }
 
-const std::vector<HostPtr>& LoadBalancerBase::hostsToUse() {
-  ASSERT(host_set_.healthyHosts().size() <= host_set_.hosts().size());
-
-  if (host_set_.hosts().empty() || isGlobalPanic()) {
-    return host_set_.hosts();
+void LoadBalancerBase::calculateZonePercentage(
+    const std::vector<std::vector<HostPtr>>& hosts_per_zone, uint64_t* ret) {
+  uint64_t total_hosts = 0;
+  for (const auto& zone_hosts : hosts_per_zone) {
+    total_hosts += zone_hosts.size();
   }
 
-  if (earlyExitNonZoneRouting()) {
-    return host_set_.healthyHosts();
+  if (total_hosts != 0) {
+    size_t i = 0;
+    for (const auto& zone_hosts : hosts_per_zone) {
+      ret[i++] = 10000ULL * zone_hosts.size() / total_hosts;
+    }
   }
+}
 
+const std::vector<HostPtr>& LoadBalancerBase::tryChooseLocalZoneHosts() {
   // At this point it's guaranteed to be at least 2 zones.
-  uint32_t number_of_zones = host_set_.healthyHostsPerZone().size();
+  size_t number_of_zones = host_set_.healthyHostsPerZone().size();
+
   ASSERT(number_of_zones >= 2U);
-  const std::vector<HostPtr>& local_zone_healthy_hosts = host_set_.healthyHostsPerZone()[0];
 
-  // If number of hosts in a local zone big enough then route all requests to the same zone.
-  if (local_zone_healthy_hosts.size() * number_of_zones >= host_set_.healthyHosts().size()) {
-    stats_.zone_over_percentage_.inc();
-    return local_zone_healthy_hosts;
+  uint64_t local_percentage[number_of_zones];
+  calculateZonePercentage(local_host_set_->healthyHostsPerZone(), local_percentage);
+
+  uint64_t upstream_percentage[number_of_zones];
+  calculateZonePercentage(host_set_.healthyHostsPerZone(), upstream_percentage);
+
+  // Try to push all of the requests to the same zone first.
+  // If we have lower percent of hosts in the local cluster in the same zone,
+  // we can push all of the requests directly to upstream cluster in the same zone.
+  if (upstream_percentage[0] >= local_percentage[0]) {
+    stats_.zone_routing_all_directly_.inc();
+    return host_set_.healthyHostsPerZone()[0];
   }
 
-  // If local zone ratio is lower than expected we should only partially route requests from the
-  // same zone.
-  double zone_host_ratio = 1.0 * local_zone_healthy_hosts.size() / host_set_.healthyHosts().size();
-  double ratio_to_route = zone_host_ratio * number_of_zones;
+  // If we cannot route all requests to the same zone, calculate what percentage can be routed.
+  // For example, if local percentage is 20% and upstream is 10%
+  // we can route only 50% of requests directly.
+  uint64_t local_percent_route = upstream_percentage[0] * 10000 / local_percentage[0];
+  if (random_.random() % 10000 < local_percent_route) {
+    stats_.zone_routing_sampled_.inc();
+    return host_set_.healthyHostsPerZone()[0];
+  }
+
+  // At this point we must route cross zone as we cannot route to the local zone.
+  stats_.zone_routing_cross_zone_.inc();
+
+  // Local zone does not have additional capacity (we have already routed what we could).
+  // Now we need to figure out how much traffic we can route cross zone and to which exact zone
+  // we should route. Percentage of requests routed cross zone to a specific zone needed be
+  // proportional to the residual capacity upstream zone has.
+  //
+  // residual_capacity contains capacity left in a given zone, we keep accumulating residual
+  // capacity to make search for sampled value easier.
+  uint64_t residual_capacity[number_of_zones];
+
+  // Local zone (index 0) does not have residual capacity as we have routed all we could.
+  residual_capacity[0] = 0;
+  for (size_t i = 1; i < number_of_zones; ++i) {
+    // Only route to the zones that have additional capacity.
+    if (upstream_percentage[i] > local_percentage[i]) {
+      residual_capacity[i] =
+          residual_capacity[i - 1] + upstream_percentage[i] - local_percentage[i];
+    } else {
+      // Zone with index "i" does not have residual capacity, but we keep accumulating previous
+      // values to make search easier on the next step.
+      residual_capacity[i] = residual_capacity[i - 1];
+    }
+  }
 
-  // Not zone routed requests will be distributed between all hosts and hence
-  // we need to route only fraction of req_percent_to_route to the local zone.
-  double actual_routing_ratio = (ratio_to_route - zone_host_ratio) / (1 - zone_host_ratio);
+  // Random sampling to select specific zone for cross zone traffic based on the additional
+  // capacity in zones.
+  uint64_t threshold = random_.random() % residual_capacity[number_of_zones - 1];
 
-  // Scale actual_routing_ratio to improve precision.
-  const uint64_t scale_factor = 10000;
-  uint64_t zone_routing_threshold = scale_factor * actual_routing_ratio;
+  // This potentially can be optimized to be O(log(N)) where N is the number of zones.
+  // Linear scan should be faster for smaller N, in most of the scenarios N will be small.
+  int i = 0;
+  while (threshold > residual_capacity[i]) {
+    i++;
+  }
 
-  if (random_.random() % 10000 < zone_routing_threshold) {
-    stats_.zone_routing_sampled_.inc();
-    return local_zone_healthy_hosts;
-  } else {
-    stats_.zone_routing_no_sampled_.inc();
+  return host_set_.healthyHostsPerZone()[i];
+}
+
+const std::vector<HostPtr>& LoadBalancerBase::hostsToUse() {
+  ASSERT(host_set_.healthyHosts().size() <= host_set_.hosts().size());
+
+  if (host_set_.hosts().empty() || isGlobalPanic(host_set_)) {
+    return host_set_.hosts();
+  }
+
+  if (earlyExitNonZoneRouting()) {
     return host_set_.healthyHosts();
   }
+
+  return tryChooseLocalZoneHosts();
 }
 
 ConstHostPtr RoundRobinLoadBalancer::chooseHost() {

source/common/upstream/load_balancer_impl.h

@@ -26,8 +26,26 @@ class LoadBalancerBase {
   Runtime::RandomGenerator& random_;
 
 private:
+  /*
+   * @return decision on quick exit from zone aware host selection.
+   */
   bool earlyExitNonZoneRouting();
-  bool isGlobalPanic();
+
+  /**
+   * For the given host_set it @return if we should be in a panic mode or not.
+   * For example, if majority of hosts are unhealthy we'll be likely in a panic mode.
+   * In this case we'll route requests to hosts no matter if they are healthy or not.
+   */
+  bool isGlobalPanic(const HostSet& host_set);
+  const std::vector<HostPtr>& tryChooseLocalZoneHosts();
+
+  /**
+   * @return (number of hosts in a given zone)/(total number of hosts) in ret param.
+   * The result is stored as integer number and scaled by 10000 multiplier for better precision.
+   * Caller is responsible for allocation/de-allocation of ret.
+   */
+  void calculateZonePercentage(const std::vector<std::vector<HostPtr>>& hosts_per_zone,
+                               uint64_t* ret);
 
   const HostSet& host_set_;
   const HostSet* local_host_set_;

source/common/upstream/upstream_impl.h

@@ -105,7 +105,7 @@ typedef std::shared_ptr<std::vector<std::vector<HostPtr>>> HostListsPtr;
 typedef std::shared_ptr<const std::vector<std::vector<HostPtr>>> ConstHostListsPtr;
 
 /**
- * Base clase for all clusters as well as thread local host sets.
+ * Base class for all clusters as well as thread local host sets.
  */
 class HostSetImpl : public virtual HostSet {
 public: