1
#include "source/extensions/load_balancing_policies/common/locality_wrr.h"
2

            
3
namespace Envoy {
4
namespace Upstream {
5

            
6
91
LocalityWrr::LocalityWrr(const HostSet& host_set, uint64_t seed) {
7
91
  rebuildLocalityScheduler(healthy_locality_scheduler_, healthy_locality_entries_,
8
91
                           host_set.healthyHostsPerLocality(), host_set.healthyHosts(),
9
91
                           host_set.hostsPerLocalityPtr(), host_set.excludedHostsPerLocalityPtr(),
10
91
                           host_set.localityWeights(), host_set.overprovisioningFactor(), seed);
11
91
  rebuildLocalityScheduler(degraded_locality_scheduler_, degraded_locality_entries_,
12
91
                           host_set.degradedHostsPerLocality(), host_set.degradedHosts(),
13
91
                           host_set.hostsPerLocalityPtr(), host_set.excludedHostsPerLocalityPtr(),
14
91
                           host_set.localityWeights(), host_set.overprovisioningFactor(), seed);
15
91
}
16

            
17
1978
absl::optional<uint32_t> LocalityWrr::chooseHealthyLocality() {
18
1978
  return chooseLocality(healthy_locality_scheduler_.get());
19
1978
}
20

            
21
5
absl::optional<uint32_t> LocalityWrr::chooseDegradedLocality() {
22
5
  return chooseLocality(degraded_locality_scheduler_.get());
23
5
}
24

            
25
void LocalityWrr::rebuildLocalityScheduler(
26
    std::unique_ptr<EdfScheduler<LocalityEntry>>& locality_scheduler,
27
    std::vector<std::shared_ptr<LocalityEntry>>& locality_entries,
28
    const HostsPerLocality& eligible_hosts_per_locality, const HostVector& eligible_hosts,
29
    HostsPerLocalityConstSharedPtr all_hosts_per_locality,
30
    HostsPerLocalityConstSharedPtr excluded_hosts_per_locality,
31
    LocalityWeightsConstSharedPtr locality_weights, uint32_t overprovisioning_factor,
32
182
    uint64_t seed) {
33
  // Rebuild the locality scheduler by computing the effective weight of each
34
  // locality in this priority. The scheduler is reset by default, and is rebuilt only if we have
35
  // locality weights (i.e. using EDS) and there is at least one eligible host in this priority.
36
  //
37
  // We omit building a scheduler when there are zero eligible hosts in the priority as
38
  // all the localities will have zero effective weight. At selection time, we'll either select
39
  // from a different scheduler or there will be no available hosts in the priority. At that point
40
  // we'll rely on other mechanisms such as panic mode to select a host, none of which rely on the
41
  // scheduler.
42
  //
43
  // TODO(htuch): if the underlying locality index ->
44
  // envoy::config::core::v3::Locality hasn't changed in hosts_/healthy_hosts_/degraded_hosts_, we
45
  // could just update locality_weight_ without rebuilding. Similar to how host
46
  // level WRR works, we would age out the existing entries via picks and lazily
47
  // apply the new weights.
48
182
  locality_scheduler = nullptr;
49
182
  if (all_hosts_per_locality != nullptr && locality_weights != nullptr &&
50
182
      !locality_weights->empty() && !eligible_hosts.empty()) {
51
58
    locality_entries.clear();
52
178
    for (uint32_t i = 0; i < all_hosts_per_locality->get().size(); ++i) {
53
120
      const double effective_weight = effectiveLocalityWeight(
54
120
          i, eligible_hosts_per_locality, *excluded_hosts_per_locality, *all_hosts_per_locality,
55
120
          *locality_weights, overprovisioning_factor);
56
120
      if (effective_weight > 0) {
57
112
        locality_entries.emplace_back(std::make_shared<LocalityEntry>(i, effective_weight));
58
112
      }
59
120
    }
60
    // If not all effective weights were zero, create the scheduler.
61
58
    if (!locality_entries.empty()) {
62
55
      locality_scheduler = std::make_unique<EdfScheduler<LocalityEntry>>(
63
55
          EdfScheduler<LocalityEntry>::createWithPicks(
64
150
              locality_entries, [](const LocalityEntry& entry) { return entry.effective_weight_; },
65
55
              seed));
66
55
    }
67
58
  }
68
182
}
69

            
70
absl::optional<uint32_t>
71
1983
LocalityWrr::chooseLocality(EdfScheduler<LocalityEntry>* locality_scheduler) {
72
1983
  if (locality_scheduler == nullptr) {
73
22
    return {};
74
22
  }
75
1961
  const std::shared_ptr<LocalityEntry> locality = locality_scheduler->pickAndAdd(
76
1961
      [](const LocalityEntry& locality) { return locality.effective_weight_; });
77
  // We don't build a schedule if there are no weighted localities, so we should always succeed.
78
1961
  ASSERT(locality != nullptr);
79
  // If we picked it before, its weight must have been positive.
80
1961
  ASSERT(locality->effective_weight_ > 0);
81
1961
  return locality->index_;
82
1983
}
83

            
84
double LocalityWrr::effectiveLocalityWeight(uint32_t index,
85
                                            const HostsPerLocality& eligible_hosts_per_locality,
86
                                            const HostsPerLocality& excluded_hosts_per_locality,
87
                                            const HostsPerLocality& all_hosts_per_locality,
88
                                            const LocalityWeights& locality_weights,
89
120
                                            uint32_t overprovisioning_factor) {
90
120
  const auto& locality_eligible_hosts = eligible_hosts_per_locality.get()[index];
91
120
  const uint32_t excluded_count = excluded_hosts_per_locality.get().size() > index
92
120
                                      ? excluded_hosts_per_locality.get()[index].size()
93
120
                                      : 0;
94
120
  const auto host_count = all_hosts_per_locality.get()[index].size() - excluded_count;
95
120
  if (host_count == 0) {
96
    return 0.0;
97
  }
98
120
  const double locality_availability_ratio = 1.0 * locality_eligible_hosts.size() / host_count;
99
120
  const uint32_t weight = locality_weights[index];
100
  // Availability ranges from 0-1.0, and is the ratio of eligible hosts to total hosts, modified
101
  // by the overprovisioning factor.
102
120
  const double effective_locality_availability_ratio =
103
120
      std::min(1.0, (overprovisioning_factor / 100.0) * locality_availability_ratio);
104
120
  return weight * effective_locality_availability_ratio;
105
120
}
106

            
107
} // namespace Upstream
108
} // namespace Envoy