Add flag to exclude nodes from local DC when running nodetool rebuild

Patch by Saranya Krishnakumar; reviewed by Dinesh Joshi, Marcus Eriksson, Yifan Cai for CASSANDRA-17870

CHANGES.txt

@@ -1,4 +1,5 @@
 4.2
+ * Add flag to exclude nodes from local DC when running nodetool rebuild (CASSANDRA-17870)
  * Adding endpoint verification option to client_encryption_options (CASSANDRA-18034)
  * Replace 'wcwidth.py' with pypi module (CASSANDRA-17287)
  * Add nodetool forcecompact to remove tombstoned or ttl'd data ignoring GC grace for given table and partition keys (CASSANDRA-17711)

src/java/org/apache/cassandra/dht/RangeStreamer.java

@@ -203,6 +203,33 @@ public String message(Replica replica)
         }
     }
 
+    /**
+    * Source filter which excludes nodes from local DC.
+    */
+    public static class ExcludeLocalDatacenterFilter implements SourceFilter
+    {
+        private final IEndpointSnitch snitch;
+        private final String localDc;
+
+        public ExcludeLocalDatacenterFilter(IEndpointSnitch snitch)
+        {
+            this.snitch = snitch;
+            this.localDc = snitch.getLocalDatacenter();
+        }
+
+        @Override
+        public boolean apply(Replica replica)
+        {
+            return !snitch.getDatacenter(replica).equals(localDc);
+        }
+
+        @Override
+        public String message(Replica replica)
+        {
+            return "Filtered " + replica + " out because it belongs to the local datacenter";
+        }
+    }
+
     /**
      * Source filter which excludes the current node from source calculations
      */

src/java/org/apache/cassandra/service/StorageService.java

@@ -1280,17 +1280,28 @@ public boolean isJoined()
 
     public void rebuild(String sourceDc)
     {
-        rebuild(sourceDc, null, null, null);
+        rebuild(sourceDc, null, null, null, false);
     }
 
     public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources)
+    {
+        rebuild(sourceDc, keyspace, tokens, specificSources, false);
+    }
+
+    public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources, boolean excludeLocalDatacenterNodes)
     {
         // check ongoing rebuild
         if (!isRebuilding.compareAndSet(false, true))
         {
             throw new IllegalStateException("Node is still rebuilding. Check nodetool netstats.");
         }
 
+        // fail if source DC is local and --exclude-local-dc is set
+        if (sourceDc != null && sourceDc.equals(DatabaseDescriptor.getLocalDataCenter()) && excludeLocalDatacenterNodes)
+        {
+            throw new IllegalArgumentException("Cannot set source data center to be local data center, when excludeLocalDataCenter flag is set");
+        }
+
         try
         {
             // check the arguments
@@ -1317,6 +1328,9 @@ public void rebuild(String sourceDc, String keyspace, String tokens, String spec
             if (sourceDc != null)
                 streamer.addSourceFilter(new RangeStreamer.SingleDatacenterFilter(DatabaseDescriptor.getEndpointSnitch(), sourceDc));
 
+            if (excludeLocalDatacenterNodes)
+                streamer.addSourceFilter(new RangeStreamer.ExcludeLocalDatacenterFilter(DatabaseDescriptor.getEndpointSnitch()));
+
             if (keyspace == null)
             {
                 for (String keyspaceName : Schema.instance.getNonLocalStrategyKeyspaces().names())

src/java/org/apache/cassandra/service/StorageServiceMBean.java

@@ -782,9 +782,23 @@ default int upgradeSSTables(String keyspaceName, boolean excludeCurrentVersion,
      * @param keyspace Name of the keyspace which to rebuild or null to rebuild all keyspaces.
      * @param tokens Range of tokens to rebuild or null to rebuild all token ranges. In the format of:
      *               "(start_token_1,end_token_1],(start_token_2,end_token_2],...(start_token_n,end_token_n]"
+     * @param specificSources list of sources that can be used for rebuilding. Must be other nodes in the cluster.
+     *                        The format of the string is comma separated values.
      */
     public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources);
 
+    /**
+    * Same as {@link #rebuild(String)}, but only for specified keyspace and ranges. It excludes local data center nodes
+    *
+    * @param sourceDc Name of DC from which to select sources for streaming or null to pick any node
+    * @param keyspace Name of the keyspace which to rebuild or null to rebuild all keyspaces.
+    * @param tokens Range of tokens to rebuild or null to rebuild all token ranges. In the format of:
+    *               "(start_token_1,end_token_1],(start_token_2,end_token_2],...(start_token_n,end_token_n]"
+    * @param specificSources list of sources that can be used for rebuilding. Mostly other nodes in the cluster.
+    * @param excludeLocalDatacenterNodes Flag to indicate whether local data center nodes should be excluded as sources for streaming.
+    */
+    public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources, boolean excludeLocalDatacenterNodes);
+
     /** Starts a bulk load and blocks until it completes. */
     public void bulkLoad(String directory);
 

src/java/org/apache/cassandra/tools/NodeProbe.java

@@ -1589,9 +1589,9 @@ public List<String> describeRing(String keyspaceName, boolean withPort) throws I
         return withPort ? ssProxy.describeRingWithPortJMX(keyspaceName) : ssProxy.describeRingJMX(keyspaceName);
     }
 
-    public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources)
+    public void rebuild(String sourceDc, String keyspace, String tokens, String specificSources, boolean excludeLocalDatacenterNodes)
     {
-        ssProxy.rebuild(sourceDc, keyspace, tokens, specificSources);
+        ssProxy.rebuild(sourceDc, keyspace, tokens, specificSources, excludeLocalDatacenterNodes);
     }
 
     public List<String> sampleKeyRange()

src/java/org/apache/cassandra/tools/nodetool/Rebuild.java

@@ -28,7 +28,7 @@
 public class Rebuild extends NodeToolCmd
 {
     @Arguments(usage = "<src-dc-name>",
-               description = "Name of DC from which to select sources for streaming. By default, pick any DC")
+               description = "Name of DC from which to select sources for streaming. By default, pick any DC (except local DC when --exclude-local-dc is set)")
     private String sourceDataCenterName = null;
 
     @Option(title = "specific_keyspace",
@@ -46,6 +46,11 @@ public class Rebuild extends NodeToolCmd
             description = "Use -s to specify hosts that this node should stream from when -ts is used. Multiple hosts should be separated using commas (e.g. 127.0.0.1,127.0.0.2,...)")
     private String specificSources = null;
 
+    @Option(title = "exclude_local_dc",
+            name = {"--exclude-local-dc"},
+            description = "Use --exclude-local-dc to exclude nodes in local data center as source for streaming.")
+    private boolean excludeLocalDatacenterNodes = false;
+
     @Override
     public void execute(NodeProbe probe)
     {
@@ -55,6 +60,6 @@ public void execute(NodeProbe probe)
             throw new IllegalArgumentException("Cannot specify tokens without keyspace.");
         }
 
-        probe.rebuild(sourceDataCenterName, keyspace, tokens, specificSources);
+        probe.rebuild(sourceDataCenterName, keyspace, tokens, specificSources, excludeLocalDatacenterNodes);
     }
 }

test/unit/org/apache/cassandra/service/StorageServiceTest.java

@@ -306,4 +306,19 @@ public void testBatchSizeWarnThresholdInKiB()
             storageService.setBatchSizeWarnThresholdInKiB(previousBatchSizeWarnThreshold);
         }
     }
+
+    @Test
+    public void testLocalDatacenterNodesExcludedDuringRebuild()
+    {
+        StorageService service = StorageService.instance;
+        try
+        {
+            service.rebuild(DatabaseDescriptor.getLocalDataCenter(), "StorageServiceTest", null, null, true);
+            fail();
+        }
+        catch (IllegalArgumentException e)
+        {
+            Assert.assertEquals("Cannot set source data center to be local data center, when excludeLocalDataCenter flag is set", e.getMessage());
+        }
+    }
 }