DiffData performance improvements (& fix JVM heap space errors)

Author: pravinbhat

pom.xml

@@ -3,7 +3,7 @@
 
   <groupId>datastax.astra.migrate</groupId>
   <artifactId>cassandra-data-migrator</artifactId>
-  <version>1.9</version>
+  <version>2.0</version>
   <packaging>jar</packaging>
 
   <properties>

src/main/java/datastax/astra/migrate/DiffJobSession.java

@@ -2,14 +2,17 @@
 
 import com.datastax.oss.driver.api.core.ConsistencyLevel;
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.cql.AsyncResultSet;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 import org.apache.spark.SparkConf;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import java.math.BigInteger;
-import java.util.concurrent.ForkJoinPool;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.IntStream;
 import java.util.stream.StreamSupport;
@@ -51,7 +54,6 @@ public static DiffJobSession getInstance(CqlSession sourceSession, CqlSession as
     }
 
     public void getDataAndDiff(BigInteger min, BigInteger max) {
-        ForkJoinPool customThreadPool = new ForkJoinPool();
         logger.info("TreadID: " + Thread.currentThread().getId() + " Processing min: " + min + " max:" + max);
         int maxAttempts = maxRetries;
         for (int retryCount = 1; retryCount <= maxAttempts; retryCount++) {
@@ -61,27 +63,30 @@ public void getDataAndDiff(BigInteger min, BigInteger max) {
                 ResultSet resultSet = sourceSession.execute(
                         sourceSelectStatement.bind(hasRandomPartitioner ? min : min.longValueExact(), hasRandomPartitioner ? max : max.longValueExact()).setConsistencyLevel(ConsistencyLevel.LOCAL_QUORUM));
 
-                customThreadPool.submit(() -> {
-                    StreamSupport.stream(resultSet.spliterator(), true).forEach(sRow -> {
-                        readLimiter.acquire(1);
-                        // do not process rows less than writeTimeStampFilter
-                        if (!(writeTimeStampFilter && (getLargestWriteTimeStamp(sRow) < minWriteTimeStampFilter
-                                || getLargestWriteTimeStamp(sRow) > maxWriteTimeStampFilter))) {
-                            if (readCounter.incrementAndGet() % printStatsAfter == 0) {
-                                printCounts("Current");
-                            }
-
-                            Row astraRow = astraSession
-                                    .execute(selectFromAstra(astraSelectStatement, sRow)).one();
-                            diff(sRow, astraRow);
-                        } else {
-                            readCounter.incrementAndGet();
-                            skippedCounter.incrementAndGet();
+                Map<Row, CompletionStage<AsyncResultSet>> writeResults = new HashMap<Row, CompletionStage<AsyncResultSet>>();
+                StreamSupport.stream(resultSet.spliterator(), false).forEach(sRow -> {
+                    readLimiter.acquire(1);
+                    // do not process rows less than writeTimeStampFilter
+                    if (!(writeTimeStampFilter && (getLargestWriteTimeStamp(sRow) < minWriteTimeStampFilter
+                            || getLargestWriteTimeStamp(sRow) > maxWriteTimeStampFilter))) {
+                        if (readCounter.incrementAndGet() % printStatsAfter == 0) {
+                            printCounts("Current");
                         }
-                    });
 
-                    printCounts("Final");
-                }).get();
+                        CompletionStage<AsyncResultSet> writeResultSet = astraSession
+                                .executeAsync(selectFromAstra(astraSelectStatement, sRow));
+                        writeResults.put(sRow, writeResultSet);
+                        if (writeResults.size() > 1000) {
+                            iterateAndClearWriteResults(writeResults);
+                        }
+                    } else {
+                        readCounter.incrementAndGet();
+                        skippedCounter.incrementAndGet();
+                    }
+                });
+                iterateAndClearWriteResults(writeResults);
+
+                printCounts("Final");
 
                 retryCount = maxAttempts;
             } catch (Exception e) {
@@ -91,7 +96,19 @@ public void getDataAndDiff(BigInteger min, BigInteger max) {
             }
         }
 
-        customThreadPool.shutdownNow();
+    }
+
+    private void iterateAndClearWriteResults(Map<Row, CompletionStage<AsyncResultSet>> writeResults) {
+        for (Row sr : writeResults.keySet()) {
+            Row ar = null;
+            try {
+                ar = writeResults.get(sr).toCompletableFuture().get().one();
+            } catch (Exception e) {
+                logger.error("Could not perform diff for Key: " + getKey(sr), e);
+            }
+            diff(sr, ar);
+        }
+        writeResults.clear();
     }
 
     public void printCounts(String finalStr) {
@@ -150,15 +167,13 @@ private void diff(Row sourceRow, Row astraRow) {
     private String isDifferent(Row sourceRow, Row astraRow) {
         StringBuffer diffData = new StringBuffer();
         IntStream.range(0, selectColTypes.size()).parallel().forEach(index -> {
-            if (!writeTimeStampCols.contains(index)) {
-                MigrateDataType dataType = selectColTypes.get(index);
-                Object source = getData(dataType, index, sourceRow);
-                Object astra = getData(dataType, index, astraRow);
-
-                boolean isDiff = dataType.diff(source, astra);
-                if (isDiff) {
-                    diffData.append(" (Index: " + index + " Source: " + source + " Astra: " + astra + " ) ");
-                }
+            MigrateDataType dataType = selectColTypes.get(index);
+            Object source = getData(dataType, index, sourceRow);
+            Object astra = getData(dataType, index, astraRow);
+
+            boolean isDiff = dataType.diff(source, astra);
+            if (isDiff) {
+                diffData.append(" (Index: " + index + " Source: " + source + " Astra: " + astra + " ) ");
             }
         });