CASSANDRA-16211: Improve metadata exception handling

.gitignore

@@ -0,0 +1,39 @@
+### Intellij ###
+*.iml
+*.ipr
+*.iws
+out/
+.idea/
+.idea_modules/
+
+### Java ###
+# Compiled class file
+*.class
+
+# Log file
+*.log
+
+# Package Files #
+*.jar
+*.war
+*.nar
+*.ear
+*.zip
+*.tar.gz
+*.rar
+
+# virtual machine crash logs, see http://www.java.com/en/download/help/error_hotspot.xml
+hs_err_pid*
+
+### Maven ###
+target/
+pom.xml.tag
+pom.xml.releaseBackup
+pom.xml.versionsBackup
+pom.xml.next
+release.properties
+dependency-reduced-pom.xml
+buildNumber.properties
+.mvn/timing.properties
+.mvn/wrapper/maven-wrapper.jar
+.flattened-pom.xml

common/src/main/java/org/apache/cassandra/diff/ExponentialRetryStrategyProvider.java

@@ -1,3 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 package org.apache.cassandra.diff;
 
 import java.util.concurrent.TimeUnit;

common/src/main/java/org/apache/cassandra/diff/KeyspaceTablePair.java

@@ -1,3 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 package org.apache.cassandra.diff;
 
 import java.io.Serializable;

common/src/main/java/org/apache/cassandra/diff/RetryStrategy.java

@@ -1,6 +1,26 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 package org.apache.cassandra.diff;
 
 import java.util.concurrent.Callable;
+import java.util.function.Function;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -16,12 +36,28 @@ public abstract class RetryStrategy {
     protected abstract boolean shouldRetry();
 
     public final <T> T retry(Callable<T> retryable) throws Exception {
+        return retryIfNot(retryable);
+    }
+
+    /**
+     * Retry a retryable.
+     * Rethrow the exception from retryable if no more retry is permitted or the thrown exception is in the exclude list.
+     */
+    @SafeVarargs
+    public final <T> T retryIfNot(Callable<T> retryable, Class<? extends Exception>... excludedExceptions) throws Exception {
+        Function<Exception, Boolean> containsException = ex -> {
+            for (Class<? extends Exception> xClass : excludedExceptions) {
+                if (xClass.isInstance(ex))
+                    return true;
+            }
+            return false;
+        };
         while (true) {
             try {
                 return retryable.call();
             }
             catch (Exception exception) {
-                if (!shouldRetry()) {
+                if (containsException.apply(exception) || !shouldRetry()) {
                     throw exception;
                 }
                 logger.warn("Retry with " + toString());

common/src/main/java/org/apache/cassandra/diff/RetryStrategyProvider.java

@@ -1,13 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
 package org.apache.cassandra.diff;
 
+import java.io.Serializable;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /**
  * Provides new RetryStrategy instances.
  * Use abstract class instead of interface in order to retain the referece to retryOptions;
  */
-public abstract class RetryStrategyProvider {
+public abstract class RetryStrategyProvider implements Serializable {
     protected final JobConfiguration.RetryOptions retryOptions;
 
     public RetryStrategyProvider(JobConfiguration.RetryOptions retryOptions) {

common/src/test/java/org/apache/cassandra/diff/ExponentialRetryStrategyTest.java

@@ -91,6 +91,40 @@ public void testOverflowPrevention() {
         }
     }
 
+    @Test
+    public void testNotMatchAndRetryWithRetryIfNot() {
+        AtomicInteger execCount = new AtomicInteger(0);
+        ExponentialRetryStrategy strategy = new ExponentialRetryStrategy(1, 5);
+        // run the code and retry since the thrown exception does not match with the exclude list
+        try {
+            strategy.retryIfNot(() -> {
+                execCount.getAndIncrement();
+                throw new IllegalStateException();
+            }, IllegalArgumentException.class, UnsupportedOperationException.class);
+        }
+        catch (Exception ex) {
+            Assert.assertSame(IllegalStateException.class, ex.getClass());
+            Assert.assertEquals(4, execCount.get());
+        }
+    }
+
+    @Test
+    public void testMatchAndRetryWithRetryIfNot() {
+        AtomicInteger execCount = new AtomicInteger(0);
+        ExponentialRetryStrategy strategy = new ExponentialRetryStrategy(1, 2);
+        // run the code and not retry since the thrown exception matches the exclude list
+        try {
+            strategy.retryIfNot(() -> {
+                execCount.getAndIncrement();
+                throw new IllegalStateException();
+            }, RuntimeException.class);
+        }
+        catch (Exception ex) {
+            Assert.assertSame(IllegalStateException.class, ex.getClass());
+            Assert.assertEquals(1, execCount.get());
+        }
+    }
+
     private JobConfiguration.RetryOptions retryOptions(long baseDelayMs, long totalDelayMs) {
         return new JobConfiguration.RetryOptions() {{
             put(ExponentialRetryStrategy.BASE_DELAY_MS_KEY, String.valueOf(baseDelayMs));

spark-job/src/main/java/org/apache/cassandra/diff/DiffJob.java

@@ -125,12 +125,13 @@ public void run(JobConfiguration configuration, JavaSparkContext sc) {
         try (Cluster metadataCluster = metadataProvider.getCluster();
              Session metadataSession = metadataCluster.connect()) {
 
+            RetryStrategyProvider retryStrategyProvider = RetryStrategyProvider.create(configuration.retryOptions());
             MetadataKeyspaceOptions metadataOptions = configuration.metadataOptions();
-            JobMetadataDb.Schema.maybeInitialize(metadataSession, metadataOptions);
+            JobMetadataDb.Schema.maybeInitialize(metadataSession, metadataOptions, retryStrategyProvider);
 
             // Job params, which once a job is created cannot be modified in subsequent re-runs
             logger.info("Creating or retrieving job parameters");
-            job = new JobMetadataDb.JobLifeCycle(metadataSession, metadataOptions.keyspace);
+            job = new JobMetadataDb.JobLifeCycle(metadataSession, metadataOptions.keyspace, retryStrategyProvider);
             Params params = getJobParams(job, configuration, tablesToCompare);
             logger.info("Job Params: {}", params);
             if (null == params)
@@ -174,7 +175,8 @@ public void run(JobConfiguration configuration, JavaSparkContext sc) {
                                                                                         sourceProvider,
                                                                                         targetProvider,
                                                                                         metadataProvider,
-                                                                                        new TrackerProvider(configuration.metadataOptions().keyspace))
+                                                                                        new TrackerProvider(configuration.metadataOptions().keyspace),
+                                                                                        retryStrategyProvider)
                                                                                  .run())
                                                              .reduce(Differ::accumulate);
             // Publish results. This also removes the job from the currently running list

spark-job/src/main/java/org/apache/cassandra/diff/Differ.java

@@ -90,7 +90,8 @@ public Differ(JobConfiguration config,
                   ClusterProvider sourceProvider,
                   ClusterProvider targetProvider,
                   ClusterProvider metadataProvider,
-                  DiffJob.TrackerProvider trackerProvider)
+                  DiffJob.TrackerProvider trackerProvider,
+                  RetryStrategyProvider retryStrategyProvider)
     {
         logger.info("Creating Differ for {}", split);
         this.jobId = params.jobId;
@@ -101,7 +102,7 @@ public Differ(JobConfiguration config,
         rateLimiter = RateLimiter.create(perExecutorRateLimit);
         this.reverseReadProbability = config.reverseReadProbability();
         this.specificTokens = config.specificTokens();
-        this.retryStrategyProvider = RetryStrategyProvider.create(config.retryOptions());
+        this.retryStrategyProvider = retryStrategyProvider;
         synchronized (Differ.class)
         {
             /*