CPP-713 - Fix load distribution of RR and DC-aware policies (#218)

CPP-701 introduced a bug where round-robin and DC-aware policies would
not equally distribute requests to the remaining nodes during a node
outage or removal.

Author: mpenick

cpp-driver/gtests/src/unit/tests/test_load_balancing.cpp

@@ -83,6 +83,38 @@ void verify_sequence(cass::QueryPlan* qp, const cass::Vector<size_t>& sequence)
   EXPECT_FALSE(qp->compute_next(&received));
 }
 
+typedef cass::Map<cass::Address, int> QueryCounts;
+
+QueryCounts run_policy(cass::LoadBalancingPolicy& policy, int count) {
+  QueryCounts counts;
+  for (int i = 0; i < 12; ++i) {
+    cass::ScopedPtr<cass::QueryPlan> qp(policy.new_query_plan("ks", NULL, NULL));
+    cass::Host::Ptr host(qp->compute_next());
+    if (host) {
+      counts[host->address()] += 1;
+    }
+  }
+  return counts;
+}
+
+void verify_dcs(const QueryCounts& counts,
+                const cass::HostMap& hosts,
+                const cass::String& expected_dc) {
+  for (QueryCounts::const_iterator it = counts.begin(),
+       end = counts.end(); it != end; ++it) {
+    cass::HostMap::const_iterator host_it = hosts.find(it->first);
+    ASSERT_NE(host_it, hosts.end());
+    EXPECT_EQ(expected_dc, host_it->second->dc());
+  }
+}
+
+void verify_query_counts(const QueryCounts& counts, int expected_count) {
+  for (QueryCounts::const_iterator it = counts.begin(),
+       end = counts.end(); it != end; ++it) {
+    EXPECT_EQ(expected_count, it->second);
+  }
+}
+
 struct RunPeriodicTask : public cass::EventLoop {
   RunPeriodicTask(cass::LatencyAwarePolicy* policy)
     : policy(policy) {
@@ -194,7 +226,7 @@ TEST(RoundRobinLoadBalancingUnitTest, OnAdd) {
   cass::Address addr_new = addr_for_sequence(seq_new);
   cass::SharedRefPtr<cass::Host> host = host_for_addr(addr_new);
   policy.on_host_added(host);
-  policy.on_host_up(host->address());
+  policy.on_host_up(host);
 
   cass::ScopedPtr<cass::QueryPlan> qp2(policy.new_query_plan("ks", NULL, NULL));
   const size_t seq2[] = {2, seq_new, 1};
@@ -248,26 +280,64 @@ TEST(RoundRobinLoadBalancingUnitTest, OnUpAndDown) {
   }
 
   // host is added to the list, but not 'up'
-  policy.on_host_up(host->address());
+  policy.on_host_up(host);
 
   cass::ScopedPtr<cass::QueryPlan> qp_after1(policy.new_query_plan("ks", NULL, NULL));
   cass::ScopedPtr<cass::QueryPlan> qp_after2(policy.new_query_plan("ks", NULL, NULL));
 
   policy.on_host_down(host->address());
   // 1 is dynamically excluded from plan
   {
-    const size_t seq[] = {3, 2};
+    const size_t seq[] = {2, 3};
     verify_sequence(qp_after1.get(), VECTOR_FROM(size_t, seq));
   }
 
-  policy.on_host_up(host->address());
+  policy.on_host_up(host);
   // now included
   {
-    const size_t seq[] = {1, 2, 3};
+    const size_t seq[] = {2, 3, 1};
     verify_sequence(qp_after2.get(), VECTOR_FROM(size_t, seq));
   }
 }
 
+TEST(RoundRobinLoadBalancingUnitTest, VerifyEqualDistribution) {
+  cass::HostMap hosts;
+  populate_hosts(3, "rack", "dc", &hosts);
+
+  cass::RoundRobinPolicy policy;
+  policy.init(cass::SharedRefPtr<cass::Host>(), hosts, NULL);
+
+  { // All nodes
+    QueryCounts counts(run_policy(policy, 12));
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_down(hosts.begin()->first);
+
+  { // One node down
+    QueryCounts counts(run_policy(policy, 12));
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+
+  policy.on_host_up(hosts.begin()->second);
+
+  { // All nodes again
+    QueryCounts counts(run_policy(policy, 12));
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_removed(hosts.begin()->second);
+
+  { // One node removed
+    QueryCounts counts(run_policy(policy, 12));
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+}
+
 TEST(DatacenterAwareLoadBalancingUnitTest, SomeDatacenterLocalUnspecified) {
   const size_t total_hosts = 3;
   cass::HostMap hosts;
@@ -302,7 +372,7 @@ TEST(DatacenterAwareLoadBalancingUnitTest, SingleLocalDown) {
     verify_sequence(qp_before.get(), VECTOR_FROM(size_t, seq));
   }
 
-  policy.on_host_up(target_host->address());
+  policy.on_host_up(target_host);
   {
     const size_t seq[] = {2, 3, 1, 4}; // local dc wrapped before remote offered
     verify_sequence(qp_after.get(), VECTOR_FROM(size_t, seq));
@@ -328,7 +398,7 @@ TEST(DatacenterAwareLoadBalancingUnitTest, AllLocalRemovedReturned) {
     verify_sequence(qp_after.get(), VECTOR_FROM(size_t, seq));
   }
 
-  policy.on_host_up(target_host->address());
+  policy.on_host_up(target_host);
 
   // make sure we get the local node first after on_up
   cass::ScopedPtr<cass::QueryPlan> qp(policy.new_query_plan("ks", NULL, NULL));
@@ -358,7 +428,7 @@ TEST(DatacenterAwareLoadBalancingUnitTest, RemoteRemovedReturned) {
     verify_sequence(qp_after.get(), VECTOR_FROM(size_t, seq));
   }
 
-  policy.on_host_up(target_host->address());
+  policy.on_host_up(target_host);
 
   // make sure we get both nodes, correct order after
   cass::ScopedPtr<cass::QueryPlan> qp(policy.new_query_plan("ks", NULL, NULL));
@@ -474,6 +544,102 @@ cass::Vector<cass::String> single_token(int64_t token) {
   return cass::Vector<cass::String>(1, ss.str());
 }
 
+TEST(DatacenterAwareLoadBalancingUnitTest, VerifyEqualDistributionLocalDc) {
+  cass::HostMap hosts;
+  populate_hosts(3, "rack", LOCAL_DC, &hosts);
+  populate_hosts(3, "rack", REMOTE_DC, &hosts);
+
+  cass::DCAwarePolicy policy("", 0, false);
+  policy.init(hosts.begin()->second, hosts, NULL);
+
+  { // All local nodes
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, LOCAL_DC);
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_down(hosts.begin()->first);
+
+  { // One local node down
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, LOCAL_DC);
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+
+  policy.on_host_up(hosts.begin()->second);
+
+  { // All local nodes again
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, LOCAL_DC);
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_removed(hosts.begin()->second);
+
+  { // One local node removed
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, LOCAL_DC);
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+}
+
+TEST(DatacenterAwareLoadBalancingUnitTest, VerifyEqualDistributionRemoteDc) {
+  cass::HostMap hosts;
+  populate_hosts(3, "rack", LOCAL_DC, &hosts);
+  populate_hosts(3, "rack", REMOTE_DC, &hosts);
+
+  cass::DCAwarePolicy policy("", 3, false); // Allow all remote DC nodes
+  policy.init(hosts.begin()->second, hosts, NULL);
+
+  cass::Host::Ptr remote_dc_node1;
+  { // Mark down all local nodes
+    cass::HostMap::iterator it = hosts.begin();
+    for (int i = 0; i < 3; ++i) {
+      policy.on_host_down(it->first);
+      it++;
+    }
+    remote_dc_node1 = it->second;
+  }
+
+  { // All remote nodes
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, REMOTE_DC);
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_down(remote_dc_node1->address());
+
+  { // One remote node down
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, REMOTE_DC);
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+
+  policy.on_host_up(remote_dc_node1);
+
+  { // All remote nodes again
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, REMOTE_DC);
+    ASSERT_EQ(counts.size(), 3u);
+    verify_query_counts(counts, 4);
+  }
+
+  policy.on_host_removed(remote_dc_node1);
+
+  { // One remote node removed
+    QueryCounts counts(run_policy(policy, 12));
+    verify_dcs(counts, hosts, REMOTE_DC);
+    ASSERT_EQ(counts.size(), 2u);
+    verify_query_counts(counts, 6);
+  }
+}
+
 TEST(TokenAwareLoadBalancingUnitTest, Simple) {
   const int64_t num_hosts = 4;
   cass::HostMap hosts;
@@ -530,7 +696,7 @@ TEST(TokenAwareLoadBalancingUnitTest, Simple) {
   }
 
   // Restore the first host and bring down the first token aware replica
-  policy.on_host_up(curr_host_it->second->address());
+  policy.on_host_up(curr_host_it->second);
   ++curr_host_it; // 2.0.0.0
   ++curr_host_it; // 3.0.0.0
   ++curr_host_it; // 4.0.0.0
@@ -601,19 +767,19 @@ TEST(TokenAwareLoadBalancingUnitTest, NetworkTopology) {
 
   {
     cass::ScopedPtr<cass::QueryPlan> qp(policy.new_query_plan("test", request_handler.get(), token_map.get()));
-    const size_t seq[] = { 3, 5, 7, 6, 2, 4 };
+    const size_t seq[] = { 3, 5, 7, 4, 6, 2 };
     verify_sequence(qp.get(), VECTOR_FROM(size_t, seq));
   }
 
   // Restore the first host and bring down the first token aware replica
-  policy.on_host_up(curr_host_it->second->address());
+  policy.on_host_up(curr_host_it->second);
   ++curr_host_it; // 2.0.0.0
   ++curr_host_it; // 3.0.0.0
   policy.on_host_down(curr_host_it->second->address());
 
   {
     cass::ScopedPtr<cass::QueryPlan> qp(policy.new_query_plan("test", request_handler.get(), token_map.get()));
-    const size_t seq[] = { 5, 7, 1, 2, 4, 6 };
+    const size_t seq[] = { 5, 7, 1, 6, 2, 4 };
     verify_sequence(qp.get(), VECTOR_FROM(size_t, seq));
   }
 }

cpp-driver/src/cluster.cpp

@@ -495,7 +495,7 @@ void Cluster::internal_notify_host_up(const Address& address) {
 
   for (LoadBalancingPolicy::Vec::const_iterator it = load_balancing_policies_.begin(),
        end = load_balancing_policies_.end(); it != end; ++it) {
-    (*it)->on_host_up(address);
+    (*it)->on_host_up(host);
   }
 
   if (is_host_ignored(host)) {

cpp-driver/src/dc_aware_policy.cpp

@@ -119,12 +119,20 @@ void DCAwarePolicy::on_host_removed(const Host::Ptr& host) {
   available_.erase(host->address());
 }
 
-void DCAwarePolicy::on_host_up(const Address& address) {
+void DCAwarePolicy::on_host_up(const Host::Ptr& host) {
+  on_host_added(host);
+
   ScopedWriteLock wl(&available_rwlock_);
-  available_.insert(address);
+  available_.insert(host->address());
 }
 
 void DCAwarePolicy::on_host_down(const Address& address) {
+  if (!remove_host(local_dc_live_hosts_, address) &&
+      !per_remote_dc_live_hosts_.remove_host(address)) {
+    LOG_DEBUG("Attempted to mark host %s as DOWN, but it doesn't exist",
+              address.to_string().c_str());
+  }
+
   ScopedWriteLock wl(&available_rwlock_);
   available_.erase(address);
 }
@@ -145,8 +153,19 @@ void DCAwarePolicy::PerDCHostMap::remove_host_from_dc(const String& dc, const Ho
   ScopedWriteLock wl(&rwlock_);
   Map::iterator i = map_.find(dc);
   if (i != map_.end()) {
-    remove_host(i->second, host);
+    cass::remove_host(i->second, host);
+  }
+}
+
+bool DCAwarePolicy::PerDCHostMap::remove_host(const Address& address) {
+  ScopedWriteLock wl(&rwlock_);
+  for (Map::iterator i = map_.begin(),
+       end = map_.end(); i != end; ++i) {
+    if (cass::remove_host(i->second, address)) {
+      return true;
+    }
   }
+  return false;
 }
 
 const CopyOnWriteHostVec& DCAwarePolicy::PerDCHostMap::get_hosts(const String& dc) const {

cpp-driver/src/dc_aware_policy.hpp

@@ -49,7 +49,7 @@ class DCAwarePolicy : public LoadBalancingPolicy {
 
   virtual void on_host_added(const Host::Ptr& host);
   virtual void on_host_removed(const Host::Ptr& host);
-  virtual void on_host_up(const Address& address);
+  virtual void on_host_up(const Host::Ptr& host);
   virtual void on_host_down(const Address& address);
 
   virtual LoadBalancingPolicy* new_instance() {
@@ -71,6 +71,7 @@ class DCAwarePolicy : public LoadBalancingPolicy {
 
     void add_host_to_dc(const String& dc, const Host::Ptr& host);
     void remove_host_from_dc(const String& dc, const Host::Ptr& host);
+    bool remove_host(const Address& address);
     const CopyOnWriteHostVec& get_hosts(const String& dc) const;
     void copy_dcs(KeySet* dcs) const;
 

cpp-driver/src/host.cpp

@@ -36,13 +36,18 @@ void add_host(CopyOnWriteHostVec& hosts, const Host::Ptr& host) {
 }
 
 void remove_host(CopyOnWriteHostVec& hosts, const Host::Ptr& host) {
+  remove_host(hosts, host->address());
+}
+
+bool remove_host(CopyOnWriteHostVec& hosts, const Address& address) {
   HostVec::iterator i;
   for (i = hosts->begin(); i != hosts->end(); ++i) {
-    if ((*i)->address() == host->address()) {
+    if ((*i)->address() == address) {
       hosts->erase(i);
-      break;
+      return true;
     }
   }
+  return false;
 }
 
 void Host::LatencyTracker::update(uint64_t latency_ns) {

cpp-driver/src/host.hpp

@@ -289,6 +289,7 @@ typedef CopyOnWritePtr<HostVec> CopyOnWriteHostVec;
 
 void add_host(CopyOnWriteHostVec& hosts, const Host::Ptr& host);
 void remove_host(CopyOnWriteHostVec& hosts, const Host::Ptr& host);
+bool remove_host(CopyOnWriteHostVec& hosts, const Address& address);
 
 } // namespace cass
 

cpp-driver/src/load_balancing.hpp

@@ -94,7 +94,7 @@ class LoadBalancingPolicy
   virtual bool is_host_up(const Address& address) const = 0;
   virtual void on_host_added(const Host::Ptr& host) = 0;
   virtual void on_host_removed(const Host::Ptr& host) = 0;
-  virtual void on_host_up(const Address& address) = 0;
+  virtual void on_host_up(const Host::Ptr& host) = 0;
   virtual void on_host_down(const Address& address) = 0;
 
   virtual QueryPlan* new_query_plan(const String& keyspace,
@@ -134,7 +134,7 @@ class ChainedLoadBalancingPolicy : public LoadBalancingPolicy {
 
   virtual void on_host_added(const Host::Ptr& host) { child_policy_->on_host_added(host); }
   virtual void on_host_removed(const Host::Ptr& host) { child_policy_->on_host_removed(host); }
-  virtual void on_host_up(const Address& address) { child_policy_->on_host_up(address); }
+  virtual void on_host_up(const Host::Ptr& host) { child_policy_->on_host_up(host); }
   virtual void on_host_down(const Address& address) { child_policy_->on_host_down(address); }
 
 protected: