From 6875ae33ca820fbcb2647c0d88a5cd596986667f Mon Sep 17 00:00:00 2001
From: Michael Penick <michael.penick@datastax.com>
Date: Tue, 7 Apr 2020 12:23:12 -0400
Subject: [PATCH] CPP-917 Fix infinite loop in token map calculation when using
 SimpleStrategy

---
 src/token_map_impl.hpp                        |  7 ++++---
 .../unit/tests/test_replication_strategy.cpp  | 21 +++++++++++++++++++
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/src/token_map_impl.hpp b/src/token_map_impl.hpp
index b08b37cbf..8596278d1 100644
--- 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
index e17984137..a8bf9c9e4 100644
--- 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;