Wire up panic mode subset to receive updates (#6221)

rgs1 · htuch · commit 1f9e2fd690e3 · 2019-03-12T01:13:20.000-04:00
The initial diff missed handling added/removed hosts, this handles that. While at this, lets move the fallback subset creation to the constructor so that it's consistent with the panic mode subset. Reported-by: @nezdolik Signed-off-by: Raul Gutierrez Segales <rgs@pinterest.com>
diff --git a/source/common/upstream/subset_lb.cc b/source/common/upstream/subset_lb.cc
@@ -35,6 +35,26 @@ SubsetLoadBalancer::SubsetLoadBalancer(
       scale_locality_weight_(subsets.scaleLocalityWeight()) {
   ASSERT(subsets.isEnabled());
 
+  if (fallback_policy_ != envoy::api::v2::Cluster::LbSubsetConfig::NO_FALLBACK) {
+    HostPredicate predicate;
+    if (fallback_policy_ == envoy::api::v2::Cluster::LbSubsetConfig::ANY_ENDPOINT) {
+      predicate = [](const Host&) -> bool { return true; };
+
+      ENVOY_LOG(debug, "subset lb: creating any-endpoint fallback load balancer");
+    } else {
+      predicate = [this](const Host& host) -> bool {
+        return hostMatches(default_subset_metadata_, host);
+      };
+
+      ENVOY_LOG(debug, "subset lb: creating fallback load balancer for {}",
+                describeMetadata(default_subset_metadata_));
+    }
+
+    fallback_subset_ = std::make_unique<LbSubsetEntry>();
+    fallback_subset_->priority_subset_ = std::make_unique<PrioritySubsetImpl>(
+        *this, predicate, locality_weight_aware_, scale_locality_weight_);
+  }
+
   if (subsets.panicModeAny()) {
     HostPredicate predicate = [](const Host&) -> bool { return true; };
 
@@ -186,34 +206,18 @@ SubsetLoadBalancer::LbSubsetEntryPtr SubsetLoadBalancer::findSubset(
 
 void SubsetLoadBalancer::updateFallbackSubset(uint32_t priority, const HostVector& hosts_added,
                                               const HostVector& hosts_removed) {
-  if (fallback_policy_ == envoy::api::v2::Cluster::LbSubsetConfig::NO_FALLBACK) {
+  if (fallback_subset_ == nullptr) {
     ENVOY_LOG(debug, "subset lb: fallback load balancer disabled");
     return;
   }
 
-  if (fallback_subset_ == nullptr) {
-    // First update: create the default host subset.
-    HostPredicate predicate;
-    if (fallback_policy_ == envoy::api::v2::Cluster::LbSubsetConfig::ANY_ENDPOINT) {
-      predicate = [](const Host&) -> bool { return true; };
-
-      ENVOY_LOG(debug, "subset lb: creating any-endpoint fallback load balancer");
-    } else {
-      predicate = std::bind(&SubsetLoadBalancer::hostMatches, this, default_subset_metadata_,
-                            std::placeholders::_1);
-
-      ENVOY_LOG(debug, "subset lb: creating fallback load balancer for {}",
-                describeMetadata(default_subset_metadata_));
-    }
+  // Add/remove hosts.
+  fallback_subset_->priority_subset_->update(priority, hosts_added, hosts_removed);
 
-    fallback_subset_.reset(new LbSubsetEntry());
-    fallback_subset_->priority_subset_.reset(
-        new PrioritySubsetImpl(*this, predicate, locality_weight_aware_, scale_locality_weight_));
-    return;
+  // Same thing for the panic mode subset.
+  if (panic_mode_subset_ != nullptr) {
+    panic_mode_subset_->priority_subset_->update(priority, hosts_added, hosts_removed);
   }
-
-  // Subsequent updates: add/remove hosts.
-  fallback_subset_->priority_subset_->update(priority, hosts_added, hosts_removed);
 }
 
 // Iterates over the added and removed hosts, looking up an LbSubsetEntryPtr for each. For every
@@ -249,9 +253,9 @@ void SubsetLoadBalancer::processSubsets(
           if (entry->initialized()) {
             update_cb(entry);
           } else {
-            HostPredicate predicate =
-                std::bind(&SubsetLoadBalancer::hostMatches, this, kvs, std::placeholders::_1);
-
+            HostPredicate predicate = [this, kvs](const Host& host) -> bool {
+              return hostMatches(kvs, host);
+            };
             new_cb(entry, predicate, kvs, adding_hosts);
           }
         }
diff --git a/test/common/upstream/subset_lb_test.cc b/test/common/upstream/subset_lb_test.cc
@@ -490,6 +490,26 @@ TEST_F(SubsetLoadBalancerTest, FallbackPanicMode) {
   EXPECT_EQ(0U, stats_.lb_subsets_selected_.value());
 }
 
+TEST_P(SubsetLoadBalancerTest, FallbackPanicModeWithUpdates) {
+  EXPECT_CALL(subset_info_, fallbackPolicy())
+      .WillRepeatedly(Return(envoy::api::v2::Cluster::LbSubsetConfig::DEFAULT_SUBSET));
+  EXPECT_CALL(subset_info_, panicModeAny()).WillRepeatedly(Return(true));
+
+  // The default subset will be empty.
+  const ProtobufWkt::Struct default_subset = makeDefaultSubset({{"version", "none"}});
+  EXPECT_CALL(subset_info_, defaultSubset()).WillRepeatedly(ReturnRef(default_subset));
+
+  init({{"tcp://127.0.0.1:80", {{"version", "default"}}}});
+  EXPECT_TRUE(lb_->chooseHost(nullptr) != nullptr);
+
+  // Removing current host, adding a new one.
+  HostSharedPtr added_host = makeHost("tcp://127.0.0.2:8000", {{"version", "new"}});
+  modifyHosts({added_host}, {host_set_.hosts_[0]});
+
+  EXPECT_EQ(1, host_set_.hosts_.size());
+  EXPECT_EQ(added_host, lb_->chooseHost(nullptr));
+}
+
 TEST_P(SubsetLoadBalancerTest, FallbackDefaultSubsetAfterUpdate) {
   EXPECT_CALL(subset_info_, fallbackPolicy())
       .WillRepeatedly(Return(envoy::api::v2::Cluster::LbSubsetConfig::DEFAULT_SUBSET));
