Apply DC filters when building the list of segments

thelastpickle · Oct 19, 2021 · 93c0410 · 93c0410
1 parent f3f80a2
commit 93c0410
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/src/server/src/main/java/io/cassandrareaper/service/RepairRunService.java b/src/server/src/main/java/io/cassandrareaper/service/RepairRunService.java
@@ -202,7 +202,8 @@ List<Segment> generateSegments(
           Segment.builder()
                  .withBaseRange(segment.getBaseRange())
                  .withTokenRanges(segment.getTokenRanges())
-                 .withReplicas(getDCsByNodeForRepairSegment(targetCluster, segment, repairUnit.getKeyspaceName()))
+                 .withReplicas(getDCsByNodeForRepairSegment(
+                    targetCluster, segment, repairUnit.getKeyspaceName(), repairUnit))
                  .build());
     }
 
@@ -212,7 +213,8 @@ List<Segment> generateSegments(
   private Map<String, String> getDCsByNodeForRepairSegment(
       Cluster cluster,
       Segment segment,
-      String keyspace) throws ReaperException {
+      String keyspace,
+      RepairUnit repairUnit) throws ReaperException {
 
     final int maxAttempts = 2;
     for (int attempt = 0; attempt < maxAttempts; attempt++) {
@@ -222,7 +224,13 @@ private Map<String, String> getDCsByNodeForRepairSegment(
         Collection<String> nodes = clusterFacade.tokenRangeToEndpoint(cluster, keyspace, segment);
         Map<String, String> dcByNode = Maps.newHashMap();
         nodes.forEach(node -> dcByNode.put(node, EndpointSnitchInfoProxy.create(jmxConnection).getDataCenter(node)));
-        return dcByNode;
+        if (repairUnit.getDatacenters().isEmpty()) {
+          return dcByNode;
+        } else {
+          return dcByNode.entrySet().stream()
+            .filter(entry -> repairUnit.getDatacenters().contains(entry.getValue()))
+            .collect(Collectors.toMap(entry -> entry.getKey(), entry -> entry.getValue()));
+        }
       } catch (RuntimeException e) {
         if (attempt < maxAttempts - 1) {
           LOG.warn("Failed getting replicas for token range {}. Attempt {} of {}",