upstream: partition hosts in a single pass

include/envoy/upstream/upstream.h

@@ -229,6 +229,16 @@ class HostsPerLocality {
   virtual std::shared_ptr<const HostsPerLocality>
   filter(std::function<bool(const Host&)> predicate) const PURE;
 
+  /**
+   * Clone object with multiple filter predicates. Returns a vector of clones, each with host that
+   * match the provided predicates.
+   * @param predicates vector of predicates on Host entries.
+   * @return vector of HostsPerLocalityConstSharedPtr clones of the HostsPerLocality that match
+   *         hosts according to predicates.
+   */
+  virtual std::vector<std::shared_ptr<const HostsPerLocality>>
+  filter(const std::vector<std::function<bool(const Host&)>>& predicates) const PURE;
+
   /**
    * Clone object.
    * @return HostsPerLocalityConstSharedPtr clone of the HostsPerLocality.

source/common/upstream/upstream_impl.cc

@@ -272,21 +272,42 @@ void HostImpl::weight(uint32_t new_weight) { weight_ = std::max(1U, std::min(128
 
 HostsPerLocalityConstSharedPtr
 HostsPerLocalityImpl::filter(std::function<bool(const Host&)> predicate) const {
-  auto* filtered_clone = new HostsPerLocalityImpl();
-  HostsPerLocalityConstSharedPtr shared_filtered_clone{filtered_clone};
+  return filter(std::vector<std::function<bool(const Host&)>>{predicate})[0];
+}
+
+std::vector<HostsPerLocalityConstSharedPtr> HostsPerLocalityImpl::filter(
+    const std::vector<std::function<bool(const Host&)>>& predicates) const {
+  // We keep two lists: one for being able to mutate the clone and one for returning to the caller.
+  // Creating them both at the start avoids iterating over the mutable values at the end to convert
+  // them to a const pointer.
+  std::vector<HostsPerLocalityImpl*> mutable_clones;
+  std::vector<HostsPerLocalityConstSharedPtr> filtered_clones;
+
+  for (size_t i = 0; i < predicates.size(); ++i) {
+    mutable_clones.emplace_back(new HostsPerLocalityImpl());
+    filtered_clones.emplace_back(mutable_clones.back());
+    mutable_clones.back()->local_ = local_;
+  }
 
-  filtered_clone->local_ = local_;
   for (const auto& hosts_locality : hosts_per_locality_) {
-    HostVector current_locality_hosts;
+    std::vector<HostVector> current_locality_hosts;
+    current_locality_hosts.resize(predicates.size());
+
+    // Since # of hosts >> # of predicates, we iterate over the hosts in the outer loop.
     for (const auto& host : hosts_locality) {
-      if (predicate(*host)) {
-        current_locality_hosts.emplace_back(host);
+      for (size_t i = 0; i < predicates.size(); ++i) {
+        if (predicates[i](*host)) {
+          current_locality_hosts[i].emplace_back(host);
+        }
       }
     }
-    filtered_clone->hosts_per_locality_.push_back(std::move(current_locality_hosts));
+
+    for (size_t i = 0; i < predicates.size(); ++i) {
+      mutable_clones[i]->hosts_per_locality_.push_back(std::move(current_locality_hosts[0]));
+    }
   }
 
-  return shared_filtered_clone;
+  return filtered_clones;
 }
 
 void HostSetImpl::updateHosts(PrioritySet::UpdateHostsParams&& update_hosts_params,
@@ -414,16 +435,15 @@ HostSetImpl::updateHostsParams(HostVectorConstSharedPtr hosts,
 PrioritySet::UpdateHostsParams
 HostSetImpl::partitionHosts(HostVectorConstSharedPtr hosts,
                             HostsPerLocalityConstSharedPtr hosts_per_locality) {
-  auto healthy_hosts = ClusterImplBase::createHostList(*hosts, Host::Health::Healthy);
-  auto degraded_hosts = ClusterImplBase::createHostList(*hosts, Host::Health::Degraded);
-  auto healthy_hosts_per_locality =
-      ClusterImplBase::createHostLists(*hosts_per_locality, Host::Health::Healthy);
-  auto degraded_hosts_per_locality =
-      ClusterImplBase::createHostLists(*hosts_per_locality, Host::Health::Degraded);
+  auto healthy_and_degraded_hosts = ClusterImplBase::partitionHostList(*hosts);
+  auto healthy_and_degraded_hosts_per_locality =
+      ClusterImplBase::partitionHostsPerLocality(*hosts_per_locality);
 
   return updateHostsParams(std::move(hosts), std::move(hosts_per_locality),
-                           std::move(healthy_hosts), std::move(healthy_hosts_per_locality),
-                           std::move(degraded_hosts), std::move(degraded_hosts_per_locality));
+                           std::move(healthy_and_degraded_hosts.first),
+                           std::move(healthy_and_degraded_hosts_per_locality.first),
+                           std::move(healthy_and_degraded_hosts.second),
+                           std::move(healthy_and_degraded_hosts_per_locality.second));
 }
 
 double HostSetImpl::effectiveLocalityWeight(uint32_t index,
@@ -673,21 +693,30 @@ ClusterImplBase::ClusterImplBase(
       });
 }
 
-HostVectorConstSharedPtr ClusterImplBase::createHostList(const HostVector& hosts,
-                                                         Host::Health health) {
-  HostVectorSharedPtr healthy_list(new HostVector());
+std::pair<HostVectorConstSharedPtr, HostVectorConstSharedPtr>
+ClusterImplBase::partitionHostList(const HostVector& hosts) {
+  auto healthy_list = std::make_shared<HostVector>();
+  auto degraded_list = std::make_shared<HostVector>();
+
   for (const auto& host : hosts) {
-    if (host->health() == health) {
+    if (host->health() == Host::Health::Healthy) {
       healthy_list->emplace_back(host);
     }
+    if (host->health() == Host::Health::Degraded) {
+      degraded_list->emplace_back(host);
+    }
   }
 
-  return healthy_list;
+  return {healthy_list, degraded_list};
 }
 
-HostsPerLocalityConstSharedPtr ClusterImplBase::createHostLists(const HostsPerLocality& hosts,
-                                                                Host::Health health) {
-  return hosts.filter([&health](const Host& host) { return host.health() == health; });
+std::pair<HostsPerLocalityConstSharedPtr, HostsPerLocalityConstSharedPtr>
+ClusterImplBase::partitionHostsPerLocality(const HostsPerLocality& hosts) {
+  auto filtered_clones =
+      hosts.filter({[](const Host& host) { return host.health() == Host::Health::Healthy; },
+                    [](const Host& host) { return host.health() == Host::Health::Degraded; }});
+
+  return {std::move(filtered_clones[0]), std::move(filtered_clones[1])};
 }
 
 bool ClusterInfoImpl::maintenanceMode() const {

source/common/upstream/upstream_impl.h

@@ -257,6 +257,8 @@ class HostsPerLocalityImpl : public HostsPerLocality {
   bool hasLocalLocality() const override { return local_; }
   const std::vector<HostVector>& get() const override { return hosts_per_locality_; }
   HostsPerLocalityConstSharedPtr filter(std::function<bool(const Host&)> predicate) const override;
+  std::vector<HostsPerLocalityConstSharedPtr>
+  filter(const std::vector<std::function<bool(const Host&)>>& predicate) const override;
 
   // The const shared pointer for the empty HostsPerLocalityImpl.
   static HostsPerLocalityConstSharedPtr empty() {
@@ -637,9 +639,14 @@ class ClusterImplBase : public Cluster, protected Logger::Loggable<Logger::Id::u
   const Network::Address::InstanceConstSharedPtr
   resolveProtoAddress(const envoy::api::v2::core::Address& address);
 
-  static HostVectorConstSharedPtr createHostList(const HostVector& hosts, Host::Health health);
-  static HostsPerLocalityConstSharedPtr createHostLists(const HostsPerLocality& hosts,
-                                                        Host::Health);
+  // Partitions the provided list of hosts into two new lists containing the healthy and degraded
+  // hosts respectively.
+  static std::pair<HostVectorConstSharedPtr, HostVectorConstSharedPtr>
+  partitionHostList(const HostVector& hosts);
+  // Partitions the provided list of hosts per locality into two new lists containing the healthy
+  // and degraded hosts respectively.
+  static std::pair<HostsPerLocalityConstSharedPtr, HostsPerLocalityConstSharedPtr>
+  partitionHostsPerLocality(const HostsPerLocality& hosts);
 
   // Upstream::Cluster
   HealthChecker* healthChecker() override { return health_checker_.get(); }