[PDS-243785] Count Attempted Hosts per DC (#5903)

Cassandra host routing now distributes retry attempts between datacenters as evenly as possible.
palantir · Feb 11, 2022 · ee5cee8 · ee5cee8
1 parent 2326544
commit ee5cee8
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 4 deletions.
diff --git a/...assandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraService.java b/...assandra/src/main/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraService.java
@@ -63,6 +63,7 @@
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
@@ -306,19 +307,28 @@ public Optional<CassandraClientPoolingContainer> getRandomGoodHostForPredicate(
 
         Set<InetSocketAddress> hostsMatchingPredicate =
                 pools.keySet().stream().filter(predicate).collect(Collectors.toSet());
-        Set<String> triedDatacenters = triedHosts.stream()
+        Map<String, Long> triedDatacenters = triedHosts.stream()
                 .map(hostToDatacenter::get)
                 .filter(Objects::nonNull)
+                .collect(Collectors.groupingBy(Function.identity(), Collectors.counting()));
+        Optional<Long> maximumAttemptsPerDatacenter =
+                triedDatacenters.values().stream().max(Long::compareTo);
+        Set<String> maximallyAttemptedDatacenters = KeyedStream.stream(triedDatacenters)
+                .filter(attempts -> Objects.equals(
+                        attempts,
+                        maximumAttemptsPerDatacenter.orElseThrow(() -> new SafeIllegalStateException(
+                                "Unexpectedly could not find the max attempts per datacenter"))))
+                .keys()
                 .collect(Collectors.toSet());
 
-        Set<InetSocketAddress> hostsInUntriedDatacenters = hostsMatchingPredicate.stream()
+        Set<InetSocketAddress> hostsInPermittedDatacenters = hostsMatchingPredicate.stream()
                 .filter(pool -> {
                     String datacenter = hostToDatacenter.get(pool);
-                    return datacenter == null || !triedDatacenters.contains(datacenter);
+                    return datacenter == null || !maximallyAttemptedDatacenters.contains(datacenter);
                 })
                 .collect(Collectors.toSet());
         Set<InetSocketAddress> filteredHosts =
-                hostsInUntriedDatacenters.isEmpty() ? hostsMatchingPredicate : hostsInUntriedDatacenters;
+                hostsInPermittedDatacenters.isEmpty() ? hostsMatchingPredicate : hostsInPermittedDatacenters;
 
         if (filteredHosts.isEmpty()) {
             log.info("No hosts match the provided predicate.");

diff --git a/...ndra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraServiceTest.java b/...ndra/src/test/java/com/palantir/atlasdb/keyvalue/cassandra/pool/CassandraServiceTest.java
@@ -47,6 +47,7 @@ public class CassandraServiceTest {
     private static final InetSocketAddress HOST_3 = InetSocketAddress.createUnresolved(HOSTNAME_3, DEFAULT_PORT);
     private static final String DC_1 = "london";
     private static final String DC_2 = "singapore";
+    private static final String DC_3 = "zurich";
 
     private CassandraKeyValueServiceConfig config;
     private Blacklist blacklist;
@@ -165,6 +166,24 @@ public void selectsHostsInAnotherDatacenter() {
                 cassandra.getRandomGoodHostForPredicate(address -> true, ImmutableSet.of(HOST_1)), HOST_2);
     }
 
+    @Test
+    public void choosesTheHostInTheLeastAttemptedDatacenter() {
+        CassandraService cassandra = clientPoolWithServers(ImmutableSet.of(HOST_1, HOST_2, HOST_3));
+        cassandra.overrideHostToDatacenterMapping(ImmutableMap.of(HOST_1, DC_1, HOST_2, DC_2, HOST_3, DC_1));
+        assertContainerHasHost(
+                cassandra.getRandomGoodHostForPredicate(address -> true, ImmutableSet.of(HOST_1, HOST_2, HOST_3)),
+                HOST_2);
+    }
+
+    @Test
+    public void distributesAttemptsWhenMultipleDatacentersAreLeastAttempted() {
+        CassandraService cassandra = clientPoolWithServers(ImmutableSet.of(HOST_1, HOST_2, HOST_3));
+        cassandra.overrideHostToDatacenterMapping(ImmutableMap.of(HOST_1, DC_1, HOST_2, DC_2, HOST_3, DC_3));
+        Set<InetSocketAddress> suggestedHosts =
+                getRecommendedHostsFromAThousandTrials(cassandra, ImmutableSet.of(HOST_1));
+        assertThat(suggestedHosts).containsExactlyInAnyOrder(HOST_2, HOST_3);
+    }
+
     @Test
     public void selectsAnyHostIfAllDatacentersAlreadyTried() {
         ImmutableSet<InetSocketAddress> allHosts = ImmutableSet.of(HOST_1, HOST_2);

diff --git a/changelog/@unreleased/pr-5903.v2.yml b/changelog/@unreleased/pr-5903.v2.yml
@@ -0,0 +1,6 @@
+type: improvement
+improvement:
+  description: Cassandra host routing now distributes retry attempts between datacenters
+    as evenly as possible.
+  links:
+  - https://github.com/palantir/atlasdb/pull/5903