链接：528. 按权重随机选择 - 力扣（LeetCode）
题解：
cpp
class Solution {
public:
Solution(vector<int>& w) {
if (w.size() > 0) {
_w_prefix_sum.resize(w.size());
_w_prefix_sum[0] = w[0];
_sum = w[0];
for (int i = 1; i < w.size(); ++i) {
_w_prefix_sum[i] = w[i] + _w_prefix_sum[i-1];
_sum += w[i];
}
}
}
int pickIndex() {
int val = random() % _sum;
//if (val < _w_prefix_sum[0]) return 0;
int left = 0;
int right = _w_prefix_sum.size();
while (left + 1 < right) {
int mid = left + (right-left) /2;
// 如果相等的化，应该查找第一个大于 val 的前缀和位置
if (_w_prefix_sum[mid] == val) {
left = mid;
} else if (_w_prefix_sum[mid] > val) {
right = mid;
} else {
left = mid;
}
}
if (val < _w_prefix_sum[left]) {
return left;
}
return right;
}
int _sum;
vector<int> _w_prefix_sum;
};
/**
* Your Solution object will be instantiated and called as such:
* Solution* obj = new Solution(w);
* int param_1 = obj->pickIndex();
*/加权负载均衡算法
cpp
include <algorithm>
#include "butil/fast_rand.h"
#include "brpc/socket.h"
#include "brpc/policy/weighted_round_robin_load_balancer.h"
#include "butil/strings/string_number_conversions.h"
namespace {
const std::vector<uint64_t> prime_stride = {
2,3,5,11,17,29,47,71,107,137,163,251,307,379,569,683,857,1289,1543,1949,2617,
2927,3407,4391,6599,9901,14867,22303,33457,50207,75323,112997,169501,254257,
381389,572087,849083,1273637,1910471,2865727,4298629,6447943,9671923,14507903,
21761863,32642861,48964297,73446469,110169743,165254623,247881989,371822987,
557734537,836601847,1254902827,1882354259,2823531397,4235297173,6352945771,
9529418671};
bool IsCoprime(uint64_t num1, uint64_t num2) {
uint64_t temp;
if (num1 < num2) {
temp = num1;
num1 = num2;
num2 = temp;
}
while (true) {
temp = num1 % num2;
if (temp == 0) {
break;
} else {
num1 = num2;
num2 = temp;
}
}
return num2 == 1;
}
// Get a reasonable stride according to weights configured of servers.
uint64_t GetStride(const uint64_t weight_sum, const size_t num) {
if (weight_sum == 1) {
return 1;
}
uint32_t average_weight = weight_sum / num;
auto iter = std::lower_bound(
prime_stride.begin(), prime_stride.end(), average_weight);
while (iter != prime_stride.end()
&& !IsCoprime(weight_sum, *iter)) {
++iter;
}
CHECK(iter != prime_stride.end()) << "Failed to get stride";
return *iter > weight_sum ? *iter % weight_sum : *iter;
}
} // namespace
namespace brpc {
namespace policy {
bool WeightedRoundRobinLoadBalancer::Add(Servers& bg, const ServerId& id) {
if (bg.server_list.capacity() < 128) {
bg.server_list.reserve(128);
}
uint32_t weight = 0;
if (!butil::StringToUint(id.tag, &weight) || weight <= 0) {
if (FLAGS_default_weight_of_wlb > 0) {
LOG(WARNING) << "Invalid weight is set: " << id.tag
<< ". Now, 'weight' has been set to 'FLAGS_default_weight_of_wlb' by default.";
weight = FLAGS_default_weight_of_wlb;
} else {
LOG(ERROR) << "Invalid weight is set: " << id.tag;
return false;
}
}
bool insert_server =
bg.server_map.emplace(id.id, bg.server_list.size()).second;
if (insert_server) {
bg.server_list.emplace_back(id.id, weight);
bg.weight_sum += weight;
return true;
}
return false;
}
bool WeightedRoundRobinLoadBalancer::Remove(Servers& bg, const ServerId& id) {
auto iter = bg.server_map.find(id.id);
if (iter != bg.server_map.end()) {
const size_t index = iter->second;
bg.weight_sum -= bg.server_list[index].weight;
bg.server_list[index] = bg.server_list.back();
bg.server_map[bg.server_list[index].id] = index;
bg.server_list.pop_back();
bg.server_map.erase(iter);
return true;
}
return false;
}
size_t WeightedRoundRobinLoadBalancer::BatchAdd(
Servers& bg, const std::vector<ServerId>& servers) {
size_t count = 0;
for (size_t i = 0; i < servers.size(); ++i) {
count += !!Add(bg, servers[i]);
}
return count;
}
size_t WeightedRoundRobinLoadBalancer::BatchRemove(
Servers& bg, const std::vector<ServerId>& servers) {
size_t count = 0;
for (size_t i = 0; i < servers.size(); ++i) {
count += !!Remove(bg, servers[i]);
}
return count;
}
bool WeightedRoundRobinLoadBalancer::AddServer(const ServerId& id) {
return _db_servers.Modify(Add, id);
}
bool WeightedRoundRobinLoadBalancer::RemoveServer(const ServerId& id) {
return _db_servers.Modify(Remove, id);
}
size_t WeightedRoundRobinLoadBalancer::AddServersInBatch(
const std::vector<ServerId>& servers) {
const size_t n = _db_servers.Modify(BatchAdd, servers);
LOG_IF(ERROR, n != servers.size())
<< "Fail to AddServersInBatch, expected " << servers.size()
<< " actually " << n;
return n;
}
size_t WeightedRoundRobinLoadBalancer::RemoveServersInBatch(
const std::vector<ServerId>& servers) {
const size_t n = _db_servers.Modify(BatchRemove, servers);
return n;
}
int WeightedRoundRobinLoadBalancer::SelectServer(const SelectIn& in, SelectOut* out) {
butil::DoublyBufferedData<Servers, TLS>::ScopedPtr s;
if (_db_servers.Read(&s) != 0) {
return ENOMEM;
}
if (s->server_list.empty()) {
return ENODATA;
}
TLS& tls = s.tls();
if (tls.IsNeededCalculateNewStride(s->weight_sum, s->server_list.size())) {
if (tls.stride == 0) {
tls.position = butil::fast_rand_less_than(s->server_list.size());
}
tls.stride = GetStride(s->weight_sum, s->server_list.size());
}
// If server list changed, the position may be out of range.
tls.position %= s->server_list.size();
// Check whether remain server was removed from server list.
if (tls.remain_server.weight > 0 &&
tls.remain_server.id != s->server_list[tls.position].id) {
tls.remain_server.weight = 0;
}
// The servers that can not be chosen.
std::unordered_set<SocketId> filter;
TLS tls_temp = tls;
uint64_t remain_weight = s->weight_sum;
size_t remain_servers = s->server_list.size();
while (remain_servers > 0) {
SocketId server_id = GetServerInNextStride(s->server_list, filter, tls_temp);
if ((remain_servers == 1 // always take last chance
|| !ExcludedServers::IsExcluded(in.excluded, server_id))
&& Socket::Address(server_id, out->ptr) == 0
&& (*out->ptr)->IsAvailable()) {
// update tls.
tls.remain_server = tls_temp.remain_server;
tls.position = tls_temp.position;
return 0;
} else {
// Skip this invalid server. We need calculate a new stride for server selection.
if (--remain_servers == 0) {
break;
}
filter.emplace(server_id);
remain_weight -= (s->server_list[s->server_map.at(server_id)]).weight;
// Select from beginning status.
tls_temp.stride = GetStride(remain_weight, remain_servers);
tls_temp.position = tls.position;
tls_temp.remain_server = tls.remain_server;
}
}
return EHOSTDOWN;
}
SocketId WeightedRoundRobinLoadBalancer::GetServerInNextStride(
const std::vector<Server>& server_list,
const std::unordered_set<SocketId>& filter,
TLS& tls) {
SocketId final_server = INVALID_SOCKET_ID;
uint64_t stride = tls.stride;
Server& remain = tls.remain_server;
if (remain.weight > 0) {
if (filter.count(remain.id) == 0) {
final_server = remain.id;
if (remain.weight > stride) {
remain.weight -= stride;
return final_server;
} else {
stride -= remain.weight;
}
}
remain.weight = 0;
++tls.position;
tls.position %= server_list.size();
}
while (stride > 0) {
final_server = server_list[tls.position].id;
if (filter.count(final_server) == 0) {
uint32_t configured_weight = server_list[tls.position].weight;
if (configured_weight > stride) {
remain.id = final_server;
remain.weight = configured_weight - stride;
return final_server;
}
stride -= configured_weight;
}
++tls.position;
tls.position %= server_list.size();
}
return final_server;
}
LoadBalancer* WeightedRoundRobinLoadBalancer::New(
const butil::StringPiece&) const {
return new (std::nothrow) WeightedRoundRobinLoadBalancer;
}
void WeightedRoundRobinLoadBalancer::Destroy() {
delete this;
}
void WeightedRoundRobinLoadBalancer::Describe(
std::ostream &os, const DescribeOptions& options) {
if (!options.verbose) {
os << "wrr";
return;
}
os << "WeightedRoundRobin{";
butil::DoublyBufferedData<Servers, TLS>::ScopedPtr s;
if (_db_servers.Read(&s) != 0) {
os << "fail to read _db_servers";
} else {
os << "n=" << s->server_list.size() << ':';
for (const auto& server : s->server_list) {
os << ' ' << server.id << '(' << server.weight << ')';
}
}
os << '}';
}
} // namespace policy
} // namespace brpc
cpp
#ifndef BRPC_POLICY_WEIGHTED_ROUND_ROBIN_LOAD_BALANCER_H
#define BRPC_POLICY_WEIGHTED_ROUND_ROBIN_LOAD_BALANCER_H
#include <map>
#include <vector>
#include <unordered_set>
#include "butil/containers/doubly_buffered_data.h"
#include "brpc/load_balancer.h"
namespace brpc {
namespace policy {
// This LoadBalancer selects server as the assigned weight.
// Weight is got from tag of ServerId.
class WeightedRoundRobinLoadBalancer : public LoadBalancer {
public:
bool AddServer(const ServerId& id) override;
bool RemoveServer(const ServerId& id) override;
size_t AddServersInBatch(const std::vector<ServerId>& servers) override;
size_t RemoveServersInBatch(const std::vector<ServerId>& servers) override;
int SelectServer(const SelectIn& in, SelectOut* out) override;
LoadBalancer* New(const butil::StringPiece&) const override;
void Destroy() override;
void Describe(std::ostream&, const DescribeOptions& options) override;
private:
struct Server {
Server(SocketId s_id = 0, uint32_t s_w = 0): id(s_id), weight(s_w) {}
SocketId id;
uint32_t weight;
};
struct Servers {
// The value is configured weight for each server.
std::vector<Server> server_list;
// The value is the index of the server in "server_list".
std::map<SocketId, size_t> server_map;
uint64_t weight_sum = 0;
};
struct TLS {
size_t position = 0;
uint64_t stride = 0;
Server remain_server;
// If server list changed, we need calculate a new stride.
bool IsNeededCalculateNewStride(const uint64_t curr_weight_sum,
const size_t curr_servers_num) {
if (curr_weight_sum != weight_sum
|| curr_servers_num != servers_num) {
weight_sum = curr_weight_sum;
servers_num = curr_servers_num;
return true;
}
return false;
}
private:
uint64_t weight_sum = 0;
size_t servers_num = 0;
};
static bool Add(Servers& bg, const ServerId& id);
static bool Remove(Servers& bg, const ServerId& id);
static size_t BatchAdd(Servers& bg, const std::vector<ServerId>& servers);
static size_t BatchRemove(Servers& bg, const std::vector<ServerId>& servers);
static SocketId GetServerInNextStride(const std::vector<Server>& server_list,
const std::unordered_set<SocketId>& filter,
TLS& tls);
butil::DoublyBufferedData<Servers, TLS> _db_servers;
};
} // namespace policy
} // namespace brpc