Skip to content

Commit 9dce4f5

Browse files
committed
Extract WeightedClusterUtil for a common operation
Signed-off-by: Venil Noronha <[email protected]>
1 parent 99352da commit 9dce4f5

File tree

7 files changed

+94
-88
lines changed

7 files changed

+94
-88
lines changed

source/common/common/utility.h

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,7 @@ class StringUtil {
387387
};
388388

389389
/**
390-
* Utilities for finding primes
390+
* Utilities for finding primes.
391391
*/
392392
class Primes {
393393
public:
@@ -411,13 +411,58 @@ class RegexUtil {
411411
* Constructs a std::regex, converting any std::regex_error exception into an EnvoyException.
412412
* @param regex std::string containing the regular expression to parse.
413413
* @param flags std::regex::flag_type containing parser flags. Defaults to std::regex::optimize.
414-
* @return std::regex constructed from regex and flags
414+
* @return std::regex constructed from regex and flags.
415415
* @throw EnvoyException if the regex string is invalid.
416416
*/
417417
static std::regex parseRegex(const std::string& regex,
418418
std::regex::flag_type flags = std::regex::optimize);
419419
};
420420

421+
/**
422+
* Utilities for working with weighted clusters.
423+
*/
424+
class WeightedClusterUtil {
425+
public:
426+
/*
427+
* Returns a WeightedClusterEntry from the given weighted clusters based on
428+
* the total cluster weight and a random value.
429+
* @param weighted_clusters a vector of WeightedClusterEntry instances.
430+
* @param total_cluster_weight the total weight of all clusters.
431+
* @param random_value the random value.
432+
* @param ignore_overflow whether to ignore cluster weight overflows.
433+
* @return a WeightedClusterEntry.
434+
*/
435+
template <typename WeightedClusterEntry>
436+
static const WeightedClusterEntry&
437+
pickCluster(const std::vector<WeightedClusterEntry>& weighted_clusters,
438+
const uint64_t total_cluster_weight, const uint64_t random_value,
439+
const bool ignore_overflow) {
440+
uint64_t selected_value = random_value % total_cluster_weight;
441+
uint64_t begin = 0;
442+
uint64_t end = 0;
443+
444+
// Find the right cluster to route to based on the interval in which
445+
// the selected value falls. The intervals are determined as
446+
// [0, cluster1_weight), [cluster1_weight, cluster1_weight+cluster2_weight),..
447+
for (const WeightedClusterEntry& cluster : weighted_clusters) {
448+
end = begin + cluster->clusterWeight();
449+
if (!ignore_overflow) {
450+
// end > total_cluster_weight: This case can only occur with Runtimes,
451+
// when the user specifies invalid weights such that
452+
// sum(weights) > total_cluster_weight.
453+
ASSERT(end <= total_cluster_weight);
454+
}
455+
456+
if (selected_value >= begin && selected_value < end) {
457+
return cluster;
458+
}
459+
begin = end;
460+
}
461+
462+
NOT_REACHED_GCOVR_EXCL_LINE;
463+
}
464+
};
465+
421466
/**
422467
* Maintains sets of numeric intervals. As new intervals are added, existing ones in the
423468
* set are combined so that no overlapping intervals remain in the representation.

source/common/router/config_impl.cc

Lines changed: 2 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -552,24 +552,8 @@ RouteConstSharedPtr RouteEntryImplBase::clusterEntry(const Http::HeaderMap& head
552552
}
553553
}
554554

555-
uint64_t selected_value = random_value % total_cluster_weight_;
556-
uint64_t begin = 0UL;
557-
uint64_t end = 0UL;
558-
559-
// Find the right cluster to route to based on the interval in which
560-
// the selected value falls. The intervals are determined as
561-
// [0, cluster1_weight), [cluster1_weight, cluster1_weight+cluster2_weight),..
562-
for (const WeightedClusterEntrySharedPtr& cluster : weighted_clusters_) {
563-
end = begin + cluster->clusterWeight();
564-
if (((selected_value >= begin) && (selected_value < end)) || (end >= total_cluster_weight_)) {
565-
// end > total_cluster_weight_: This case can only occur with Runtimes, when the user
566-
// specifies invalid weights such that sum(weights) > total_cluster_weight_. In this case,
567-
// terminate the search and just return the cluster whose weight caused the overflow.
568-
return cluster;
569-
}
570-
begin = end;
571-
}
572-
NOT_REACHED_GCOVR_EXCL_LINE;
555+
return WeightedClusterUtil::pickCluster(weighted_clusters_, total_cluster_weight_, random_value,
556+
true);
573557
}
574558

575559
void RouteEntryImplBase::validateClusters(Upstream::ClusterManager& cm) const {

source/common/tcp_proxy/tcp_proxy.cc

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#include "common/common/assert.h"
1616
#include "common/common/empty_string.h"
1717
#include "common/common/fmt.h"
18+
#include "common/common/utility.h"
1819
#include "common/config/well_known_names.h"
1920
#include "common/router/metadatamatchcriteria_impl.h"
2021

@@ -83,8 +84,9 @@ Config::Config(const envoy::config::filter::network::tcp_proxy::v2::TcpProxy& co
8384
total_cluster_weight_ = 0;
8485
for (const envoy::config::filter::network::tcp_proxy::v2::TcpProxy::WeightedCluster::
8586
ClusterWeight& cluster_desc : config.weighted_clusters().clusters()) {
86-
weighted_clusters_.emplace_back(WeightedClusterEntry(cluster_desc));
87-
total_cluster_weight_ += weighted_clusters_.back().cluster_weight_;
87+
std::unique_ptr<WeightedClusterEntry> cluster_entry(new WeightedClusterEntry(cluster_desc));
88+
weighted_clusters_.emplace_back(std::move(cluster_entry));
89+
total_cluster_weight_ += weighted_clusters_.back()->clusterWeight();
8890
}
8991
}
9092

@@ -135,33 +137,13 @@ const std::string& Config::getRegularRouteFromEntries(Network::Connection& conne
135137
return EMPTY_STRING;
136138
}
137139

138-
const std::string& Config::getWeightedClusterRoute(const uint64_t random_value) {
139-
uint64_t selected_value = random_value % total_cluster_weight_;
140-
uint64_t begin = 0UL;
141-
uint64_t end = 0UL;
142-
143-
// Find the right cluster to route to based on the interval in which
144-
// the selected value falls. The intervals are determined as
145-
// [0, cluster1_weight), [cluster1_weight, cluster1_weight+cluster2_weight),..
146-
for (const WeightedClusterEntry& cluster : weighted_clusters_) {
147-
end = begin + cluster.cluster_weight_;
148-
ASSERT(end <= total_cluster_weight_);
149-
150-
if (selected_value >= begin && selected_value < end) {
151-
return cluster.cluster_name_;
152-
}
153-
begin = end;
154-
}
155-
156-
NOT_REACHED_GCOVR_EXCL_LINE;
157-
}
158-
159140
const std::string& Config::getRouteFromEntries(Network::Connection& connection) {
160-
if (!weighted_clusters_.empty()) {
161-
return getWeightedClusterRoute(random_generator_.random());
162-
} else {
141+
if (weighted_clusters_.empty()) {
163142
return getRegularRouteFromEntries(connection);
164143
}
144+
return WeightedClusterUtil::pickCluster(weighted_clusters_, total_cluster_weight_,
145+
random_generator_.random(), false)
146+
->clusterName();
165147
}
166148

167149
UpstreamDrainManager& Config::drainManager() {

source/common/tcp_proxy/tcp_proxy.h

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,6 @@ class Config {
102102
*/
103103
const std::string& getRouteFromEntries(Network::Connection& connection);
104104
const std::string& getRegularRouteFromEntries(Network::Connection& connection);
105-
const std::string& getWeightedClusterRoute(const uint64_t random_value);
106105

107106
const TcpProxyStats& stats() { return shared_config_->stats(); }
108107
const std::vector<AccessLog::InstanceSharedPtr>& accessLogs() { return access_logs_; }
@@ -128,16 +127,22 @@ class Config {
128127
std::string cluster_name_;
129128
};
130129

131-
struct WeightedClusterEntry {
130+
class WeightedClusterEntry {
131+
public:
132132
WeightedClusterEntry(const envoy::config::filter::network::tcp_proxy::v2::TcpProxy::
133133
WeightedCluster::ClusterWeight& config);
134134

135+
const std::string& clusterName() const { return cluster_name_; }
136+
uint64_t clusterWeight() const { return cluster_weight_; }
137+
138+
private:
135139
const std::string cluster_name_;
136140
const uint64_t cluster_weight_;
137141
};
142+
typedef std::shared_ptr<WeightedClusterEntry> WeightedClusterEntrySharedPtr;
138143

139144
std::vector<Route> routes_;
140-
std::vector<WeightedClusterEntry> weighted_clusters_;
145+
std::vector<WeightedClusterEntrySharedPtr> weighted_clusters_;
141146
uint64_t total_cluster_weight_;
142147
std::vector<AccessLog::InstanceSharedPtr> access_logs_;
143148
const uint32_t max_connect_attempts_;

source/extensions/filters/network/thrift_proxy/router/router_impl.cc

Lines changed: 2 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -42,26 +42,8 @@ RouteConstSharedPtr RouteEntryImplBase::clusterEntry(uint64_t random_value) cons
4242
if (weighted_clusters_.empty()) {
4343
return shared_from_this();
4444
}
45-
46-
uint64_t selected_value = random_value % total_cluster_weight_;
47-
uint64_t begin = 0UL;
48-
uint64_t end = 0UL;
49-
50-
// Find the right cluster to route to based on the interval in which
51-
// the selected value falls. The intervals are determined as
52-
// [0, cluster1_weight), [cluster1_weight, cluster1_weight+cluster2_weight),..
53-
for (const WeightedClusterEntrySharedPtr& cluster : weighted_clusters_) {
54-
end = begin + cluster->clusterWeight();
55-
ASSERT(end <= total_cluster_weight_);
56-
57-
if (selected_value >= begin && selected_value < end) {
58-
return cluster;
59-
}
60-
61-
begin = end;
62-
}
63-
64-
NOT_REACHED_GCOVR_EXCL_LINE;
45+
return WeightedClusterUtil::pickCluster(weighted_clusters_, total_cluster_weight_, random_value,
46+
false);
6547
}
6648

6749
bool RouteEntryImplBase::headersMatch(const Http::HeaderMap& headers) const {

test/common/common/utility_test.cc

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,6 +470,33 @@ TEST(RegexUtil, parseRegex) {
470470
}
471471
}
472472

473+
class WeightedClusterEntry {
474+
public:
475+
WeightedClusterEntry(const std::string name, const uint64_t weight)
476+
: name_(name), weight_(weight) {}
477+
478+
const std::string& clusterName() const { return name_; }
479+
uint64_t clusterWeight() const { return weight_; }
480+
481+
private:
482+
const std::string name_;
483+
const uint64_t weight_;
484+
};
485+
typedef std::shared_ptr<WeightedClusterEntry> WeightedClusterEntrySharedPtr;
486+
487+
TEST(WeightedClusterUtil, pickCluster) {
488+
std::vector<WeightedClusterEntrySharedPtr> clusters;
489+
490+
std::unique_ptr<WeightedClusterEntry> cluster1(new WeightedClusterEntry("cluster1", 10));
491+
clusters.emplace_back(std::move(cluster1));
492+
493+
std::unique_ptr<WeightedClusterEntry> cluster2(new WeightedClusterEntry("cluster2", 90));
494+
clusters.emplace_back(std::move(cluster2));
495+
496+
EXPECT_EQ("cluster1", WeightedClusterUtil::pickCluster(clusters, 100, 5, false)->clusterName());
497+
EXPECT_EQ("cluster2", WeightedClusterUtil::pickCluster(clusters, 80, 79, true)->clusterName());
498+
}
499+
473500
static std::string intervalSetIntToString(const IntervalSetImpl<int>& interval_set) {
474501
std::string out;
475502
const char* prefix = "";

test/common/tcp_proxy/tcp_proxy_test.cc

Lines changed: 0 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -481,25 +481,6 @@ class TcpProxyTest : public testing::Test {
481481
Network::Address::InstanceConstSharedPtr upstream_remote_address_;
482482
};
483483

484-
TEST_F(TcpProxyTest, WeightedClusters) {
485-
envoy::config::filter::network::tcp_proxy::v2::TcpProxy config;
486-
487-
envoy::config::filter::network::tcp_proxy::v2::TcpProxy::WeightedCluster::ClusterWeight*
488-
cluster1 = config.mutable_weighted_clusters()->mutable_clusters()->Add();
489-
cluster1->set_name("cluster1");
490-
cluster1->set_weight(10);
491-
492-
envoy::config::filter::network::tcp_proxy::v2::TcpProxy::WeightedCluster::ClusterWeight*
493-
cluster2 = config.mutable_weighted_clusters()->mutable_clusters()->Add();
494-
cluster2->set_name("cluster2");
495-
cluster2->set_weight(90);
496-
497-
configure(config);
498-
499-
EXPECT_EQ(std::string("cluster1"), config_->getWeightedClusterRoute(5));
500-
EXPECT_EQ(std::string("cluster2"), config_->getWeightedClusterRoute(25));
501-
}
502-
503484
TEST_F(TcpProxyTest, DefaultRoutes) {
504485
envoy::config::filter::network::tcp_proxy::v2::TcpProxy config = defaultConfig();
505486

0 commit comments

Comments
 (0)