CPP-917 Fix infinite loop in token map calculation when using SimpleS…

…trategy
datastax · Apr 13, 2020 · 6875ae3 · 6875ae3
1 parent e918c7d
commit 6875ae3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 3 deletions.
diff --git a/src/token_map_impl.hpp b/src/token_map_impl.hpp
@@ -516,19 +516,20 @@ void ReplicationStrategy<Partitioner>::build_replicas_simple(const TokenHostVec&
   if (it == replication_factors_.end()) {
     return;
   }
-  size_t num_replicas = std::min<size_t>(it->second.count, tokens.size());
+  const size_t num_tokens = tokens.size();
+  const size_t num_replicas = std::min<size_t>(it->second.count, num_tokens);
   for (typename TokenHostVec::const_iterator i = tokens.begin(), end = tokens.end(); i != end;
        ++i) {
     CopyOnWriteHostVec replicas(new HostVec());
     replicas->reserve(num_replicas);
     typename TokenHostVec::const_iterator token_it = i;
-    do {
+    for (size_t j = 0; j < num_tokens && replicas->size() < num_replicas; ++j) {
       add_replica(replicas, Host::Ptr(Host::Ptr(token_it->second)));
       ++token_it;
       if (token_it == tokens.end()) {
         token_it = tokens.begin();
       }
-    } while (replicas->size() < num_replicas);
+    }
     result.push_back(TokenReplicas(i->first, replicas));
   }
 }

diff --git a/tests/src/unit/tests/test_replication_strategy.cpp b/tests/src/unit/tests/test_replication_strategy.cpp
@@ -185,6 +185,27 @@ TEST(ReplicationStrategyUnitTest, Simple) {
   }
 }
 
+TEST(ReplicationStrategyUnitTest, SimpleNumHostsLessThanReplicationFactor) {
+  MockTokenMap<Murmur3Partitioner> token_map;
+
+  token_map.init_simple_strategy(3);
+
+  MockTokenMap<Murmur3Partitioner>::Token t1 = 0;
+
+  // To reproduce the issue the number of tokens needs to be greater than
+  // (or equal to) the RF because the RF is bounded by the number of tokens.
+  token_map.add_token(t1, "1.0.0.1");
+  token_map.add_token(100, "1.0.0.1");
+  token_map.add_token(200, "1.0.0.1");
+  token_map.add_token(300, "1.0.0.1");
+
+  token_map.build_replicas();
+
+  const CopyOnWriteHostVec& hosts = token_map.find_hosts(t1);
+  ASSERT_TRUE(hosts && hosts->size() == 1);
+  check_host((*hosts)[0], "1.0.0.1");
+}
+
 TEST(ReplicationStrategyUnitTest, NetworkTopology) {
   MockTokenMap<Murmur3Partitioner> token_map;