Make sure WeightedRandomizedLoadBalancer can traverse the whole server list

Author: chenBright

src/brpc/policy/weighted_randomized_load_balancer.cpp

@@ -26,8 +26,9 @@
 namespace brpc {
 namespace policy {
 
-static bool server_compare(const WeightedRandomizedLoadBalancer::Server& lhs, const WeightedRandomizedLoadBalancer::Server& rhs) {
-        return (lhs.current_weight_sum < rhs.current_weight_sum);
+static bool server_compare(const WeightedRandomizedLoadBalancer::Server& lhs,
+                           const WeightedRandomizedLoadBalancer::Server& rhs) {
+        return lhs.current_weight_sum < rhs.current_weight_sum;
 }
 
 bool WeightedRandomizedLoadBalancer::Add(Servers& bg, const ServerId& id) {
@@ -38,15 +39,16 @@ bool WeightedRandomizedLoadBalancer::Add(Servers& bg, const ServerId& id) {
     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.";
+                         << ". 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;
+        bg.server_map.emplace(id.id, bg.server_list.size()).second;
     if (insert_server) {
         uint64_t current_weight_sum = bg.weight_sum + weight;
         bg.server_list.emplace_back(id.id, weight, current_weight_sum);
@@ -124,21 +126,27 @@ int WeightedRandomizedLoadBalancer::SelectServer(const SelectIn& in, SelectOut*
         return ENODATA;
     }
     uint64_t weight_sum = s->weight_sum;
-    for (size_t i = 0; i < n; ++i) {
-        uint64_t random_weight = butil::fast_rand_less_than(weight_sum);
-        const Server random_server(0, 0, random_weight);
-        const auto& server = std::lower_bound(s->server_list.begin(), s->server_list.end(), random_server, server_compare);
-        const SocketId id = server->id;
-        if (((i + 1) == n  // always take last chance
-             || !ExcludedServers::IsExcluded(in.excluded, id))
-            && Socket::Address(id, out->ptr) == 0
-            && (*out->ptr)->IsAvailable()) {
+    uint64_t random_weight = butil::fast_rand_less_than(weight_sum);
+    const Server random_server(0, 0, random_weight);
+    const auto& server = std::lower_bound(
+        s->server_list.begin(), s->server_list.end(), random_server, server_compare);
+    if (IsIdSelected(in, server->id, out->ptr)) {
+        // We found an available server
+        return 0;
+    }
+
+    uint32_t offset = server - s->server_list.begin();
+    uint32_t stride = butil::fast_rand_less_than(n);
+    for (size_t i = 1; i < n; ++i) {
+        offset = (offset + stride) % n;
+        SocketId id = s->server_list[offset].id;
+        if ((i + 1) == n  // always take last chance
+            || IsIdSelected(in, id, out->ptr)) {
             // We found an available server
             return 0;
         }
     }
-    // After we traversed the whole server list, there is still no
-    // available server
+    // After we traversed the whole server list, there is still no available server.
     return EHOSTDOWN;
 }
 

src/brpc/policy/weighted_randomized_load_balancer.h

@@ -31,17 +31,18 @@ namespace policy {
 // Weight is got from tag of ServerId.
 class WeightedRandomizedLoadBalancer : public LoadBalancer {
 public:
-    bool AddServer(const ServerId& id);
-    bool RemoveServer(const ServerId& id);
-    size_t AddServersInBatch(const std::vector<ServerId>& servers);
-    size_t RemoveServersInBatch(const std::vector<ServerId>& servers);
-    int SelectServer(const SelectIn& in, SelectOut* out);
-    LoadBalancer* New(const butil::StringPiece&) const;
-    void Destroy();
-    void Describe(std::ostream& os, const DescribeOptions&);
+    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& os, const DescribeOptions&) override;
 
     struct Server {
-        Server(SocketId s_id = 0, uint32_t s_w = 0, uint64_t s_c_w_s = 0): id(s_id), weight(s_w), current_weight_sum(s_c_w_s) {}
+        Server(SocketId s_id = 0, uint32_t s_w = 0, uint64_t s_c_w_s = 0)
+            : id(s_id), weight(s_w), current_weight_sum(s_c_w_s) {}
         SocketId id;
         uint32_t weight;
         uint64_t current_weight_sum;
@@ -60,6 +61,10 @@ class WeightedRandomizedLoadBalancer : public LoadBalancer {
     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 bool IsIdSelected(const SelectIn& in, SocketId id, SocketUniquePtr* out) {
+        return !ExcludedServers::IsExcluded(in.excluded, id) &&
+               Socket::Address(id, out) == 0 && (*out)->IsAvailable();
+    }
 
     butil::DoublyBufferedData<Servers> _db_servers;
 };