Major refactor to make CDM Spark Native (#328)

* Made Partition into its own class & refactored stuff to make that work

* Made CounterUnit its own class & refactored JobCounter to work with it.

* Made JobType (Migrate, Validate & Guardrail) independent of track-run feature and renamed slices/partitions to PartitionRanges. Also provided actual jobs access to PartitionRange class.

* Refactored code to be Spark Native, fixed metrics reporting & also improved trackRun feature.

* Updated readme

* Fixed metrics issue when trackRun was disabled.

Author: pravinbhat

README.md

@@ -160,7 +160,7 @@ spark-submit --properties-file cdm.properties \
 - If a table has only collection and/or UDT non-key columns, the `writetime` used on target will be time the job was run. If you want to avoid this, we recommend setting `spark.cdm.schema.ttlwritetime.calc.useCollections` param to `true` in such scenarios. 
 - When CDM migration (or validation with autocorrect) is run multiple times on the same table (for whatever reasons), it could lead to duplicate entries in `list` type columns. Note this is [due to a Cassandra/DSE bug](https://issues.apache.org/jira/browse/CASSANDRA-11368) and not a CDM issue. This issue can be addressed by enabling and setting a positive value for `spark.cdm.transform.custom.writetime.incrementBy` param. This param was specifically added to address this issue.
 - When you rerun job to resume from a previous run, the run metrics (read, write, skipped, etc.) captured in table `cdm_run_info` will be only for the current run. If the previous run was killed for some reasons, its run metrics may not have been saved. If the previous run did complete (not killed) but with errors, then you will have all run metrics from previous run as well.
-- The Spark Cluster based deployment currently has a bug. It reports '0' for all count metrics, while doing underlying tasks (Migration, Validation, etc.). We are working to address this in the upcoming releases. Also note that this issue is only with the Spark cluster deployment and not with the single VM run.  
+- When running on a Spark Cluster (and not a single VM), the rate-limit values (`spark.cdm.perfops.ratelimit.origin` & `spark.cdm.perfops.ratelimit.target`) applies to individual Spark worker nodes. Hence this value should be set to `effective-rate-limit-you-need`/`number-of-spark-worker-nodes` . E.g. If you need an effective rate-limit of 10000, and the number of Spark worker nodes are 4, then you should set the above rate-limit params to a value of 2500.
 
 # Performance recommendations 
 Below recommendations may only be useful when migrating large tables where the default performance is not good enough

RELEASE.md

@@ -1,6 +1,11 @@
 # Release Notes
+## [5.0.0] - 2024-11-08
+- CDM refactored to be fully Spark Native and more performant when deployed on a multi-node Spark Cluster
+- `trackRun` feature has been expanded to record `run-info` for each part in the `CDM_RUN_DETAILS` table. Along with granular metrics, this information can be used to troubleshoot any unbalanced problematic partitions.
+- This release has feature parity with 4.x release and is also backword compatible while adding the above mentioned improvements. However, we are upgrading it to 5.x as its a major rewrite of the code to make it Spark native.
+
 ## [4.7.0] - 2024-10-25
-- CDM refractored to work when deployed on a Spark Cluster
+- CDM refactored to work when deployed on a Spark Cluster
 - More performant for large migration efforts (multi-terabytes clusters with several billions of rows) using Spark Cluster (instead of individual VMs)
 - No functional changes and fully backward compatible, just refactor to support Spark cluster deployment
 

src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java

@@ -60,25 +60,26 @@ public TargetUpsertRunDetailsStatement(CqlSession session, String keyspaceTable)
 
         this.session.execute("CREATE TABLE IF NOT EXISTS " + cdmKsTabInfo
                 + " (table_name TEXT, run_id BIGINT, run_type TEXT, prev_run_id BIGINT, start_time TIMESTAMP, end_time TIMESTAMP, run_info TEXT, status TEXT, PRIMARY KEY (table_name, run_id))");
+        this.session.execute("CREATE TABLE IF NOT EXISTS " + cdmKsTabDetails
+                + " (table_name TEXT, run_id BIGINT, start_time TIMESTAMP, token_min BIGINT, token_max BIGINT, status TEXT, run_info TEXT, PRIMARY KEY ((table_name, run_id), token_min))");
 
         // TODO: Remove this code block after a few releases, its only added for backward compatibility
         try {
             this.session.execute("ALTER TABLE " + cdmKsTabInfo + " ADD status TEXT");
+            this.session.execute("ALTER TABLE " + cdmKsTabDetails + " ADD run_info TEXT");
         } catch (Exception e) {
             // ignore if column already exists
             logger.trace("Column 'status' already exists in table {}", cdmKsTabInfo);
         }
-        this.session.execute("CREATE TABLE IF NOT EXISTS " + cdmKsTabDetails
-                + " (table_name TEXT, run_id BIGINT, start_time TIMESTAMP, token_min BIGINT, token_max BIGINT, status TEXT, PRIMARY KEY ((table_name, run_id), token_min))");
 
         boundInitInfoStatement = bindStatement("INSERT INTO " + cdmKsTabInfo
                 + " (table_name, run_id, run_type, prev_run_id, start_time, status) VALUES (?, ?, ?, ?, dateof(now()), ?)");
         boundInitStatement = bindStatement("INSERT INTO " + cdmKsTabDetails
                 + " (table_name, run_id, token_min, token_max, status) VALUES (?, ?, ?, ?, ?)");
         boundEndInfoStatement = bindStatement("UPDATE " + cdmKsTabInfo
                 + " SET end_time = dateof(now()), run_info = ?, status = ? WHERE table_name = ? AND run_id = ?");
-        boundUpdateStatement = bindStatement(
-                "UPDATE " + cdmKsTabDetails + " SET status = ? WHERE table_name = ? AND run_id = ? AND token_min = ?");
+        boundUpdateStatement = bindStatement("UPDATE " + cdmKsTabDetails
+                + " SET status = ?, run_info = ? WHERE table_name = ? AND run_id = ? AND token_min = ?");
         boundUpdateStartStatement = bindStatement("UPDATE " + cdmKsTabDetails
                 + " SET start_time = dateof(now()), status = ? WHERE table_name = ? AND run_id = ? AND token_min = ?");
         boundSelectInfoStatement = bindStatement(
@@ -87,7 +88,8 @@ public TargetUpsertRunDetailsStatement(CqlSession session, String keyspaceTable)
                 + " WHERE table_name = ? AND run_id = ? AND status = ? ALLOW FILTERING");
     }
 
-    public Collection<PartitionRange> getPendingPartitions(long prevRunId) throws RunNotStartedException {
+    public Collection<PartitionRange> getPendingPartitions(long prevRunId, JobType jobType)
+            throws RunNotStartedException {
         if (prevRunId == 0) {
             return Collections.emptyList();
         }
@@ -105,27 +107,29 @@ public Collection<PartitionRange> getPendingPartitions(long prevRunId) throws Ru
         }
 
         final Collection<PartitionRange> pendingParts = new ArrayList<PartitionRange>();
-        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.NOT_STARTED.toString()));
-        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.STARTED.toString()));
-        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.FAIL.toString()));
-        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.DIFF.toString()));
+        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.NOT_STARTED.toString(), jobType));
+        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.STARTED.toString(), jobType));
+        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.FAIL.toString(), jobType));
+        pendingParts.addAll(getPartitionsByStatus(prevRunId, TrackRun.RUN_STATUS.DIFF.toString(), jobType));
 
         return pendingParts;
     }
 
-    protected Collection<PartitionRange> getPartitionsByStatus(long prevRunId, String status) {
-        ResultSet rs = session.execute(boundSelectStatement.setString("table_name", tableName)
-                .setLong("run_id", prevRunId).setString("status", status));
-
+    protected Collection<PartitionRange> getPartitionsByStatus(long runId, String status, JobType jobType) {
         final Collection<PartitionRange> pendingParts = new ArrayList<PartitionRange>();
-        rs.forEach(row -> {
+        getResultSetByStatus(runId, status).forEach(row -> {
             PartitionRange part = new PartitionRange(BigInteger.valueOf(row.getLong("token_min")),
-                    BigInteger.valueOf(row.getLong("token_max")));
+                    BigInteger.valueOf(row.getLong("token_max")), jobType);
             pendingParts.add(part);
         });
         return pendingParts;
     }
 
+    protected ResultSet getResultSetByStatus(long runId, String status) {
+        return session.execute(boundSelectStatement.setString("table_name", tableName).setLong("run_id", runId)
+                .setString("status", status));
+    }
+
     public void initCdmRun(long runId, long prevRunId, Collection<PartitionRange> parts, JobType jobType) {
         ResultSet rsInfo = session
                 .execute(boundSelectInfoStatement.setString("table_name", tableName).setLong("run_id", runId));
@@ -153,13 +157,14 @@ public void endCdmRun(long runId, String runInfo) {
                 .setString("run_info", runInfo).setString("status", TrackRun.RUN_STATUS.ENDED.toString()));
     }
 
-    public void updateCdmRun(long runId, BigInteger min, TrackRun.RUN_STATUS status) {
+    public void updateCdmRun(long runId, BigInteger min, TrackRun.RUN_STATUS status, String runInfo) {
         if (TrackRun.RUN_STATUS.STARTED.equals(status)) {
             session.execute(boundUpdateStartStatement.setString("table_name", tableName).setLong("run_id", runId)
                     .setLong("token_min", min.longValue()).setString("status", status.toString()));
         } else {
             session.execute(boundUpdateStatement.setString("table_name", tableName).setLong("run_id", runId)
-                    .setLong("token_min", min.longValue()).setString("status", status.toString()));
+                    .setLong("token_min", min.longValue()).setString("status", status.toString())
+                    .setString("run_info", runInfo));
         }
     }
 

src/main/java/com/datastax/cdm/feature/TrackRun.java

@@ -39,8 +39,9 @@ public TrackRun(CqlSession session, String keyspaceTable) {
         this.runStatement = new TargetUpsertRunDetailsStatement(session, keyspaceTable);
     }
 
-    public Collection<PartitionRange> getPendingPartitions(long prevRunId) throws RunNotStartedException {
-        Collection<PartitionRange> pendingParts = runStatement.getPendingPartitions(prevRunId);
+    public Collection<PartitionRange> getPendingPartitions(long prevRunId, JobType jobType)
+            throws RunNotStartedException {
+        Collection<PartitionRange> pendingParts = runStatement.getPendingPartitions(prevRunId, jobType);
         logger.info("###################### {} partitions pending from previous run id {} ######################",
                 pendingParts.size(), prevRunId);
         return pendingParts;
@@ -51,8 +52,8 @@ public void initCdmRun(long runId, long prevRunId, Collection<PartitionRange> pa
         logger.info("###################### Run Id for this job is: {} ######################", runId);
     }
 
-    public void updateCdmRun(long runId, BigInteger min, RUN_STATUS status) {
-        runStatement.updateCdmRun(runId, min, status);
+    public void updateCdmRun(long runId, BigInteger min, RUN_STATUS status, String runInfo) {
+        runStatement.updateCdmRun(runId, min, status, runInfo);
     }
 
     public void endCdmRun(long runId, String runInfo) {

src/main/java/com/datastax/cdm/job/AbstractJobSession.java

@@ -40,7 +40,6 @@ public abstract class AbstractJobSession<T> extends BaseJobSession {
     protected EnhancedSession originSession;
     protected EnhancedSession targetSession;
     protected Guardrail guardrailFeature;
-    protected JobCounter jobCounter;
     protected Long printStatsAfter;
     protected TrackRun trackRunFeature;
     protected long runId;
@@ -65,8 +64,6 @@ protected AbstractJobSession(CqlSession originSession, CqlSession targetSession,
                     KnownProperties.getDefault(KnownProperties.PRINT_STATS_AFTER));
             printStatsAfter = propertyHelper.getLong(KnownProperties.PRINT_STATS_AFTER);
         }
-        this.jobCounter = new JobCounter(printStatsAfter,
-                propertyHelper.getBoolean(KnownProperties.PRINT_STATS_PER_PART));
 
         rateLimiterOrigin = RateLimiter.create(propertyHelper.getInteger(KnownProperties.PERF_RATELIMIT_ORIGIN));
         rateLimiterTarget = RateLimiter.create(propertyHelper.getInteger(KnownProperties.PERF_RATELIMIT_TARGET));
@@ -127,11 +124,4 @@ public synchronized void initCdmRun(long runId, long prevRunId, Collection<Parti
         DataUtility.deleteGeneratedSCB(runId);
     }
 
-    public synchronized void printCounts(boolean isFinal) {
-        if (isFinal) {
-            jobCounter.printFinal(runId, trackRunFeature);
-        } else {
-            jobCounter.printProgress();
-        }
-    }
 }

src/main/java/com/datastax/cdm/job/CDMMetricsAccumulator.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 com.datastax.cdm.job;
+
+import org.apache.spark.util.AccumulatorV2;
+
+import com.datastax.cdm.job.IJobSessionFactory.JobType;
+
+public class CDMMetricsAccumulator extends AccumulatorV2<JobCounter, JobCounter> {
+
+    private static final long serialVersionUID = -4185304101452658315L;
+    private JobCounter jobCounter;
+
+    public CDMMetricsAccumulator(JobType jobType) {
+        jobCounter = new JobCounter(jobType);
+    }
+
+    @Override
+    public void add(JobCounter v) {
+        jobCounter.add(v);
+    }
+
+    @Override
+    public AccumulatorV2<JobCounter, JobCounter> copy() {
+        return this;
+    }
+
+    @Override
+    public boolean isZero() {
+        return jobCounter.isZero();
+    }
+
+    @Override
+    public void merge(AccumulatorV2<JobCounter, JobCounter> other) {
+        jobCounter.add(other.value());
+    }
+
+    @Override
+    public void reset() {
+        jobCounter.reset();
+    }
+
+    @Override
+    public JobCounter value() {
+        return jobCounter;
+    }
+
+}

src/main/java/com/datastax/cdm/job/CopyJobSession.java

@@ -51,9 +51,6 @@ public class CopyJobSession extends AbstractJobSession<PartitionRange> {
 
     protected CopyJobSession(CqlSession originSession, CqlSession targetSession, PropertyHelper propHelper) {
         super(originSession, targetSession, propHelper);
-        this.jobCounter.setRegisteredTypes(JobCounter.CounterType.READ, JobCounter.CounterType.WRITE,
-                JobCounter.CounterType.SKIPPED, JobCounter.CounterType.ERROR, JobCounter.CounterType.UNFLUSHED);
-
         pkFactory = this.originSession.getPKFactory();
         isCounterTable = this.originSession.getCqlTable().isCounterTable();
         fetchSize = this.originSession.getCqlTable().getFetchSizeInRows();
@@ -69,10 +66,10 @@ protected void processPartitionRange(PartitionRange range) {
         ThreadContext.put(THREAD_CONTEXT_LABEL, getThreadLabel(min, max));
         logger.info("ThreadID: {} Processing min: {} max: {}", Thread.currentThread().getId(), min, max);
         if (null != trackRunFeature)
-            trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.STARTED);
+            trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.STARTED, "");
 
         BatchStatement batch = BatchStatement.newInstance(BatchType.UNLOGGED);
-        jobCounter.threadReset();
+        JobCounter jobCounter = range.getJobCounter();
 
         try {
             OriginSelectByPartitionRangeStatement originSelectByPartitionRangeStatement = this.originSession
@@ -117,20 +114,21 @@ protected void processPartitionRange(PartitionRange range) {
             jobCounter.threadIncrement(JobCounter.CounterType.WRITE,
                     jobCounter.getCount(JobCounter.CounterType.UNFLUSHED));
             jobCounter.threadReset(JobCounter.CounterType.UNFLUSHED);
-            if (null != trackRunFeature)
-                trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.PASS);
+            jobCounter.globalIncrement();
+            if (null != trackRunFeature) {
+                trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.PASS, jobCounter.getThreadCounters(true));
+            }
         } catch (Exception e) {
             jobCounter.threadIncrement(JobCounter.CounterType.ERROR,
                     jobCounter.getCount(JobCounter.CounterType.READ) - jobCounter.getCount(JobCounter.CounterType.WRITE)
                             - jobCounter.getCount(JobCounter.CounterType.SKIPPED));
-            if (null != trackRunFeature)
-                trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.FAIL);
             logger.error("Error with PartitionRange -- ThreadID: {} Processing min: {} max: {}",
                     Thread.currentThread().getId(), min, max, e);
             logger.error("Error stats " + jobCounter.getThreadCounters(false));
-        } finally {
             jobCounter.globalIncrement();
-            printCounts(false);
+            if (null != trackRunFeature) {
+                trackRunFeature.updateCdmRun(runId, min, TrackRun.RUN_STATUS.FAIL, jobCounter.getThreadCounters(true));
+            }
         }
     }
 

src/main/java/com/datastax/cdm/job/CopyJobSessionFactory.java

@@ -36,7 +36,4 @@ public AbstractJobSession<PartitionRange> getInstance(CqlSession originSession,
         return jobSession;
     }
 
-    public JobType getJobType() {
-        return JobType.MIGRATE;
-    }
 }

src/main/java/com/datastax/cdm/job/CounterUnit.java

@@ -16,35 +16,34 @@
 package com.datastax.cdm.job;
 
 import java.io.Serializable;
-import java.util.concurrent.atomic.AtomicLong;
 
 public class CounterUnit implements Serializable {
 
     private static final long serialVersionUID = 2194336948011681878L;
-    private final AtomicLong globalCounter = new AtomicLong(0);
-    private final transient ThreadLocal<Long> threadLocalCounter = ThreadLocal.withInitial(() -> 0L);
+    private long globalCounter = 0;
+    private long threadLocalCounter = 0;
 
     public void incrementThreadCounter(long incrementBy) {
-        threadLocalCounter.set(threadLocalCounter.get() + incrementBy);
+        threadLocalCounter += incrementBy;
     }
 
     public long getThreadCounter() {
-        return threadLocalCounter.get();
+        return threadLocalCounter;
     }
 
     public void resetThreadCounter() {
-        threadLocalCounter.set(0L);
+        threadLocalCounter = 0;
     }
 
     public void setGlobalCounter(long value) {
-        globalCounter.set(value);
+        globalCounter = value;
     }
 
     public void addThreadToGlobalCounter() {
-        globalCounter.addAndGet(threadLocalCounter.get());
+        globalCounter += threadLocalCounter;
     }
 
     public long getGlobalCounter() {
-        return globalCounter.get();
+        return globalCounter;
     }
 }