subset lb: allow ring hash/maglev LB to work with subsets

docs/root/intro/arch_overview/upstream/load_balancing/subsets.rst

@@ -35,12 +35,20 @@ therefore, contain a definition that has the same keys as a given route in order
 balancing to occur.
 
 This feature can only be enabled using the V2 configuration API. Furthermore, host metadata is only
-supported when using the EDS discovery type for clusters. Host metadata for subset load balancing
-must be placed under the filter name ``"envoy.lb"``. Similarly, route metadata match criteria use
-the ``"envoy.lb"`` filter name. Host metadata may be hierarchical (e.g., the value for a top-level
-key may be a structured value or list), but the subset load balancer only compares top-level keys
-and values. Therefore when using structured values, a route's match criteria will only match if an
-identical structured value appears in the host's metadata.
+supported when hosts are defined using
+:ref:`ClusterLoadAssignments <envoy_api_msg_ClusterLoadAssignment>`. ClusterLoadAssignments are
+available via EDS or the Cluster :ref:`load_assignment <envoy_api_field_Cluster.load_assignment>`
+field. Host metadata for subset load balancing must be placed under the filter name ``"envoy.lb"``.
+Similarly, route metadata match criteria use ``"envoy.lb"`` filter name. Host metadata may be
+hierarchical (e.g., the value for a top-level key may be a structured value or list), but the
+subset load balancer only compares top-level keys and values. Therefore when using structured
+values, a route's match criteria will only match if an identical structured value appears in the
+host's metadata.
+
+Finally, note that subset load balancing is not available for the
+:ref:`ORIGINAL_DST_LB <envoy_api_enum_value_Cluster.LbPolicy.ORIGINAL_DST_LB>` or
+:ref:`CLUSTER_PROVIDED <envoy_api_enum_value_Cluster.LbPolicy.CLUSTER_PROVIDED>` load balancer
+policies.
 
 Examples
 ^^^^^^^^

source/common/upstream/cluster_manager_impl.cc

@@ -667,17 +667,22 @@ void ClusterManagerImpl::loadCluster(const envoy::api::v2::Cluster& cluster,
   const auto cluster_entry_it = cluster_map.find(cluster_reference.info()->name());
 
   // If an LB is thread aware, create it here. The LB is not initialized until cluster pre-init
-  // finishes.
+  // finishes. For RingHash/Maglev don't create the LB here if subset balancing is enabled,
+  // because the thread_aware_lb_ field takes precedence over the subset lb).
   if (cluster_reference.info()->lbType() == LoadBalancerType::RingHash) {
-    cluster_entry_it->second->thread_aware_lb_ = std::make_unique<RingHashLoadBalancer>(
-        cluster_reference.prioritySet(), cluster_reference.info()->stats(),
-        cluster_reference.info()->statsScope(), runtime_, random_,
-        cluster_reference.info()->lbRingHashConfig(), cluster_reference.info()->lbConfig());
+    if (!cluster_reference.info()->lbSubsetInfo().isEnabled()) {
+      cluster_entry_it->second->thread_aware_lb_ = std::make_unique<RingHashLoadBalancer>(
+          cluster_reference.prioritySet(), cluster_reference.info()->stats(),
+          cluster_reference.info()->statsScope(), runtime_, random_,
+          cluster_reference.info()->lbRingHashConfig(), cluster_reference.info()->lbConfig());
+    }
   } else if (cluster_reference.info()->lbType() == LoadBalancerType::Maglev) {
-    cluster_entry_it->second->thread_aware_lb_ = std::make_unique<MaglevLoadBalancer>(
-        cluster_reference.prioritySet(), cluster_reference.info()->stats(),
-        cluster_reference.info()->statsScope(), runtime_, random_,
-        cluster_reference.info()->lbConfig());
+    if (!cluster_reference.info()->lbSubsetInfo().isEnabled()) {
+      cluster_entry_it->second->thread_aware_lb_ = std::make_unique<MaglevLoadBalancer>(
+          cluster_reference.prioritySet(), cluster_reference.info()->stats(),
+          cluster_reference.info()->statsScope(), runtime_, random_,
+          cluster_reference.info()->lbConfig());
+    }
   } else if (cluster_reference.info()->lbType() == LoadBalancerType::ClusterProvided) {
     cluster_entry_it->second->thread_aware_lb_ = std::move(new_cluster_pair.second);
   }

source/common/upstream/original_dst_cluster.cc

@@ -184,10 +184,6 @@ OriginalDstClusterFactory::createClusterImpl(
         envoy::api::v2::Cluster_LbPolicy_Name(cluster.lb_policy()),
         envoy::api::v2::Cluster_DiscoveryType_Name(cluster.type())));
   }
-  if (cluster.has_lb_subset_config() && cluster.lb_subset_config().subset_selectors_size() != 0) {
-    throw EnvoyException(
-        fmt::format("cluster: cluster type 'original_dst' may not be used with lb_subset_config"));
-  }
 
   // TODO(mattklein123): The original DST load balancer type should be deprecated and instead
   //                     the cluster should directly supply the load balancer. This will remove

source/common/upstream/upstream_impl.cc

@@ -643,12 +643,24 @@ ClusterInfoImpl::ClusterInfoImpl(
                       envoy::api::v2::Cluster_LbPolicy_Name(config.lb_policy()),
                       envoy::api::v2::Cluster_DiscoveryType_Name(config.type())));
     }
+    if (config.has_lb_subset_config()) {
+      throw EnvoyException(
+          fmt::format("cluster: LB policy {} cannot be combined with lb_subset_config",
+                      envoy::api::v2::Cluster_LbPolicy_Name(config.lb_policy())));
+    }
+
     lb_type_ = LoadBalancerType::ClusterProvided;
     break;
   case envoy::api::v2::Cluster::MAGLEV:
     lb_type_ = LoadBalancerType::Maglev;
     break;
   case envoy::api::v2::Cluster::CLUSTER_PROVIDED:
+    if (config.has_lb_subset_config()) {
+      throw EnvoyException(
+          fmt::format("cluster: LB policy {} cannot be combined with lb_subset_config",
+                      envoy::api::v2::Cluster_LbPolicy_Name(config.lb_policy())));
+    }
+
     lb_type_ = LoadBalancerType::ClusterProvided;
     break;
   default:

test/common/upstream/BUILD

@@ -29,6 +29,9 @@ envoy_cc_test(
 envoy_cc_test(
     name = "cluster_manager_impl_test",
     srcs = ["cluster_manager_impl_test.cc"],
+    external_deps = [
+        "abseil_optional",
+    ],
     deps = [
         ":utility_lib",
         "//include/envoy/stats:stats_interface",
@@ -43,8 +46,10 @@ envoy_cc_test(
         "//source/common/stats:stats_lib",
         "//source/common/upstream:cluster_factory_lib",
         "//source/common/upstream:cluster_manager_lib",
+        "//source/common/upstream:subset_lb_lib",
         "//source/extensions/transport_sockets/raw_buffer:config",
         "//source/extensions/transport_sockets/tls:context_lib",
+        "//test/integration/clusters:custom_static_cluster",
         "//test/mocks/access_log:access_log_mocks",
         "//test/mocks/api:api_mocks",
         "//test/mocks/http:http_mocks",
@@ -61,6 +66,7 @@ envoy_cc_test(
         "//test/test_common:simulated_time_system_lib",
         "//test/test_common:threadsafe_singleton_injector_lib",
         "//test/test_common:utility_lib",
+        "@envoy_api//envoy/api/v2:cds_cc",
     ],
 )
 

test/common/upstream/cluster_manager_impl_test.cc

@@ -2,6 +2,7 @@
 #include <string>
 
 #include "envoy/admin/v2alpha/config_dump.pb.h"
+#include "envoy/api/v2/cds.pb.h"
 #include "envoy/api/v2/core/base.pb.h"
 #include "envoy/network/listen_socket.h"
 #include "envoy/upstream/upstream.h"
@@ -17,10 +18,12 @@
 #include "common/singleton/manager_impl.h"
 #include "common/upstream/cluster_factory_impl.h"
 #include "common/upstream/cluster_manager_impl.h"
+#include "common/upstream/subset_lb.h"
 
 #include "extensions/transport_sockets/tls/context_manager_impl.h"
 
 #include "test/common/upstream/utility.h"
+#include "test/integration/clusters/custom_static_cluster.h"
 #include "test/mocks/access_log/mocks.h"
 #include "test/mocks/api/mocks.h"
 #include "test/mocks/http/mocks.h"
@@ -38,6 +41,7 @@
 #include "test/test_common/threadsafe_singleton_injector.h"
 #include "test/test_common/utility.h"
 
+#include "absl/strings/str_replace.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
@@ -575,14 +579,48 @@ TEST_F(ClusterManagerImplTest, OriginalDstLbRestriction2) {
       "'ORIGINAL_DST_LB' is allowed only with cluster type 'ORIGINAL_DST'");
 }
 
-TEST_F(ClusterManagerImplTest, SubsetLoadBalancerInitialization) {
-  const std::string yaml = R"EOF(
+class ClusterManagerSubsetInitializationTest
+    : public ClusterManagerImplTest,
+      public testing::WithParamInterface<envoy::api::v2::Cluster_LbPolicy> {
+public:
+  ClusterManagerSubsetInitializationTest() = default;
+
+  static std::vector<envoy::api::v2::Cluster_LbPolicy> lbPolicies() {
+    int first = static_cast<int>(envoy::api::v2::Cluster_LbPolicy_LbPolicy_MIN);
+    int last = static_cast<int>(envoy::api::v2::Cluster_LbPolicy_LbPolicy_MAX);
+    ASSERT(first < last);
+
+    std::vector<envoy::api::v2::Cluster_LbPolicy> policies;
+    for (int i = first; i <= last; i++) {
+      if (envoy::api::v2::Cluster_LbPolicy_IsValid(i)) {
+        auto policy = static_cast<envoy::api::v2::Cluster_LbPolicy>(i);
+        if (policy != envoy::api::v2::Cluster_LbPolicy_LOAD_BALANCING_POLICY_CONFIG) {
+          policies.push_back(policy);
+        }
+      }
+    }
+    return policies;
+  }
+
+  static std::string paramName(const testing::TestParamInfo<ParamType>& info) {
+    const std::string& name = envoy::api::v2::Cluster_LbPolicy_Name(info.param);
+    return absl::StrReplaceAll(name, {{"_", ""}});
+  }
+};
+
+// Test initialization of subset load balancer with every possible load balancer policy.
+TEST_P(ClusterManagerSubsetInitializationTest, SubsetLoadBalancerInitialization) {
+  const std::string yamlPattern = R"EOF(
 static_resources:
   clusters:
   - name: cluster_1
     connect_timeout: 0.250s
-    type: static
-    lb_policy: round_robin
+    {}
+    lb_policy: "{}"
+    lb_subset_config:
+      fallback_policy: ANY_ENDPOINT
+      subset_selectors:
+        - keys: [ "x" ]
     load_assignment:
       endpoints:
         - lb_endpoints:
@@ -598,37 +636,82 @@ TEST_F(ClusterManagerImplTest, SubsetLoadBalancerInitialization) {
                   port_value: 8001
   )EOF";
 
-  envoy::config::bootstrap::v2::Bootstrap bootstrap = parseBootstrapFromV2Yaml(yaml);
-  envoy::api::v2::Cluster::LbSubsetConfig* subset_config =
-      bootstrap.mutable_static_resources()->mutable_clusters(0)->mutable_lb_subset_config();
-  subset_config->set_fallback_policy(envoy::api::v2::Cluster::LbSubsetConfig::ANY_ENDPOINT);
-  subset_config->add_subset_selectors()->add_keys("x");
+  const std::string& policy_name = envoy::api::v2::Cluster_LbPolicy_Name(GetParam());
 
-  create(bootstrap);
-  checkStats(1 /*added*/, 0 /*modified*/, 0 /*removed*/, 1 /*active*/, 0 /*warming*/);
+  std::string cluster_type = "type: STATIC";
+  if (GetParam() == envoy::api::v2::Cluster_LbPolicy_ORIGINAL_DST_LB) {
+    cluster_type = "type: ORIGINAL_DST";
+  } else if (GetParam() == envoy::api::v2::Cluster_LbPolicy_CLUSTER_PROVIDED) {
+    // This custom cluster type is registered by linking test/integration/custom/static_cluster.cc.
+    cluster_type = "cluster_type: { name: envoy.clusters.custom_static_with_lb }";
+  }
 
-  factory_.tls_.shutdownThread();
+  const std::string yaml = fmt::format(yamlPattern, cluster_type, policy_name);
+
+  if (GetParam() == envoy::api::v2::Cluster_LbPolicy_ORIGINAL_DST_LB ||
+      GetParam() == envoy::api::v2::Cluster_LbPolicy_CLUSTER_PROVIDED) {
+    EXPECT_THROW_WITH_MESSAGE(
+        create(parseBootstrapFromV2Yaml(yaml)), EnvoyException,
+        fmt::format("cluster: LB policy {} cannot be combined with lb_subset_config",
+                    envoy::api::v2::Cluster_LbPolicy_Name(GetParam())));
+
+  } else {
+    create(parseBootstrapFromV2Yaml(yaml));
+    checkStats(1 /*added*/, 0 /*modified*/, 0 /*removed*/, 1 /*active*/, 0 /*warming*/);
+
+    Upstream::ThreadLocalCluster* tlc = cluster_manager_->get("cluster_1");
+    EXPECT_NE(nullptr, tlc);
+
+    if (tlc) {
+      Upstream::LoadBalancer& lb = tlc->loadBalancer();
+      EXPECT_NE(nullptr, dynamic_cast<Upstream::SubsetLoadBalancer*>(&lb));
+    }
+
+    factory_.tls_.shutdownThread();
+  }
 }
 
-TEST_F(ClusterManagerImplTest, SubsetLoadBalancerRestriction) {
+INSTANTIATE_TEST_SUITE_P(ClusterManagerSubsetInitializationTest,
+                         ClusterManagerSubsetInitializationTest,
+                         testing::ValuesIn(ClusterManagerSubsetInitializationTest::lbPolicies()),
+                         ClusterManagerSubsetInitializationTest::paramName);
+
+TEST_F(ClusterManagerImplTest, SubsetLoadBalancerOriginalDstRestriction) {
   const std::string yaml = R"EOF(
 static_resources:
   clusters:
   - name: cluster_1
     connect_timeout: 0.250s
     type: original_dst
     lb_policy: original_dst_lb
+    lb_subset_config:
+      fallback_policy: ANY_ENDPOINT
+      subset_selectors:
+        - keys: [ "x" ]
   )EOF";
 
-  envoy::config::bootstrap::v2::Bootstrap bootstrap = parseBootstrapFromV2Yaml(yaml);
-  envoy::api::v2::Cluster::LbSubsetConfig* subset_config =
-      bootstrap.mutable_static_resources()->mutable_clusters(0)->mutable_lb_subset_config();
-  subset_config->set_fallback_policy(envoy::api::v2::Cluster::LbSubsetConfig::ANY_ENDPOINT);
-  subset_config->add_subset_selectors()->add_keys("x");
+  EXPECT_THROW_WITH_MESSAGE(
+      create(parseBootstrapFromV2Yaml(yaml)), EnvoyException,
+      "cluster: LB policy ORIGINAL_DST_LB cannot be combined with lb_subset_config");
+}
+
+TEST_F(ClusterManagerImplTest, SubsetLoadBalancerClusterProvidedLbRestriction) {
+  const std::string yaml = R"EOF(
+static_resources:
+  clusters:
+  - name: cluster_1
+    connect_timeout: 0.250s
+    type: static
+    lb_policy: cluster_provided
+    lb_subset_config:
+      fallback_policy: ANY_ENDPOINT
+      subset_selectors:
+        - keys: [ "x" ]
+  )EOF";
 
   EXPECT_THROW_WITH_MESSAGE(
-      create(bootstrap), EnvoyException,
-      "cluster: cluster type 'original_dst' may not be used with lb_subset_config");
+      create(parseBootstrapFromV2Yaml(yaml)), EnvoyException,
+      "cluster: LB policy CLUSTER_PROVIDED cannot be combined with lb_subset_config");
 }
 
 TEST_F(ClusterManagerImplTest, SubsetLoadBalancerLocalityAware) {
@@ -639,6 +722,11 @@ TEST_F(ClusterManagerImplTest, SubsetLoadBalancerLocalityAware) {
     connect_timeout: 0.250s
     type: STATIC
     lb_policy: ROUND_ROBIN
+    lb_subset_config:
+      fallback_policy: ANY_ENDPOINT
+      subset_selectors:
+        - keys: [ "x" ]
+      locality_weight_aware: true
     load_assignment:
       endpoints:
         - lb_endpoints:
@@ -654,12 +742,7 @@ TEST_F(ClusterManagerImplTest, SubsetLoadBalancerLocalityAware) {
                   port_value: 8001
   )EOF";
 
-  envoy::config::bootstrap::v2::Bootstrap bootstrap = parseBootstrapFromV2Yaml(yaml);
-  envoy::api::v2::Cluster::LbSubsetConfig* subset_config =
-      bootstrap.mutable_static_resources()->mutable_clusters(0)->mutable_lb_subset_config();
-  subset_config->set_locality_weight_aware(true);
-
-  EXPECT_THROW_WITH_MESSAGE(create(bootstrap), EnvoyException,
+  EXPECT_THROW_WITH_MESSAGE(create(parseBootstrapFromV2Yaml(yaml)), EnvoyException,
                             "Locality weight aware subset LB requires that a "
                             "locality_weighted_lb_config be set in cluster_1");
 }

test/integration/BUILD

@@ -252,6 +252,17 @@ envoy_cc_test(
     ],
 )
 
+envoy_cc_test(
+    name = "http_subset_lb_integration_test",
+    srcs = [
+        "http_subset_lb_integration_test.cc",
+    ],
+    deps = [
+        ":http_integration_lib",
+        "//test/common/upstream:utility_lib",
+    ],
+)
+
 envoy_cc_test(
     name = "http_timeout_integration_test",
     srcs = [