Refactor changes #1 (in prep for cluster-metrics change) (#325)

* Made Partition into its own class & refactored stuff to make that work * Made CounterUnit its own class & refactored JobCounter to work with it.
datastax · Nov 1, 2024 · e9bb7e0 · e9bb7e0
1 parent d36afdc
commit e9bb7e0
Show file tree

Hide file tree

Showing 16 changed files with 127 additions and 88 deletions.
diff --git a/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java b/src/main/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatement.java
@@ -26,9 +26,8 @@
 
 import com.datastax.cdm.feature.TrackRun;
 import com.datastax.cdm.feature.TrackRun.RUN_TYPE;
+import com.datastax.cdm.job.Partition;
 import com.datastax.cdm.job.RunNotStartedException;
-import com.datastax.cdm.job.SplitPartitions;
-import com.datastax.cdm.job.SplitPartitions.Partition;
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
 import com.datastax.oss.driver.api.core.cql.ResultSet;
@@ -88,7 +87,7 @@ public TargetUpsertRunDetailsStatement(CqlSession session, String keyspaceTable)
                 + " WHERE table_name = ? AND run_id = ? AND status = ? ALLOW FILTERING");
     }
 
-    public Collection<SplitPartitions.Partition> getPendingPartitions(long prevRunId) throws RunNotStartedException {
+    public Collection<Partition> getPendingPartitions(long prevRunId) throws RunNotStartedException {
         if (prevRunId == 0) {
             return Collections.emptyList();
         }
@@ -105,7 +104,7 @@ public Collection<SplitPartitions.Partition> getPendingPartitions(long prevRunId
             }
         }
 
-        final Collection<SplitPartitions.Partition> pendingParts = new ArrayList<SplitPartitions.Partition>();
+        final Collection<Partition> pendingParts = new ArrayList<Partition>();
         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()));
@@ -114,11 +113,11 @@ public Collection<SplitPartitions.Partition> getPendingPartitions(long prevRunId
         return pendingParts;
     }
 
-    protected Collection<SplitPartitions.Partition> getPartitionsByStatus(long prevRunId, String status) {
+    protected Collection<Partition> getPartitionsByStatus(long prevRunId, String status) {
         ResultSet rs = session.execute(boundSelectStatement.setString("table_name", tableName)
                 .setLong("run_id", prevRunId).setString("status", status));
 
-        final Collection<SplitPartitions.Partition> pendingParts = new ArrayList<SplitPartitions.Partition>();
+        final Collection<Partition> pendingParts = new ArrayList<Partition>();
         rs.forEach(row -> {
             Partition part = new Partition(BigInteger.valueOf(row.getLong("token_min")),
                     BigInteger.valueOf(row.getLong("token_max")));
@@ -127,7 +126,7 @@ protected Collection<SplitPartitions.Partition> getPartitionsByStatus(long prevR
         return pendingParts;
     }
 
-    public void initCdmRun(long runId, long prevRunId, Collection<SplitPartitions.Partition> parts, RUN_TYPE runType) {
+    public void initCdmRun(long runId, long prevRunId, Collection<Partition> parts, RUN_TYPE runType) {
         ResultSet rsInfo = session
                 .execute(boundSelectInfoStatement.setString("table_name", tableName).setLong("run_id", runId));
         if (null != rsInfo.one()) {

diff --git a/src/main/java/com/datastax/cdm/feature/TrackRun.java b/src/main/java/com/datastax/cdm/feature/TrackRun.java
@@ -22,8 +22,8 @@
 import org.slf4j.LoggerFactory;
 
 import com.datastax.cdm.cql.statement.TargetUpsertRunDetailsStatement;
+import com.datastax.cdm.job.Partition;
 import com.datastax.cdm.job.RunNotStartedException;
-import com.datastax.cdm.job.SplitPartitions;
 import com.datastax.oss.driver.api.core.CqlSession;
 
 public class TrackRun {
@@ -42,14 +42,14 @@ public TrackRun(CqlSession session, String keyspaceTable) {
         this.runStatement = new TargetUpsertRunDetailsStatement(session, keyspaceTable);
     }
 
-    public Collection<SplitPartitions.Partition> getPendingPartitions(long prevRunId) throws RunNotStartedException {
-        Collection<SplitPartitions.Partition> pendingParts = runStatement.getPendingPartitions(prevRunId);
+    public Collection<Partition> getPendingPartitions(long prevRunId) throws RunNotStartedException {
+        Collection<Partition> pendingParts = runStatement.getPendingPartitions(prevRunId);
         logger.info("###################### {} partitions pending from previous run id {} ######################",
                 pendingParts.size(), prevRunId);
         return pendingParts;
     }
 
-    public void initCdmRun(long runId, long prevRunId, Collection<SplitPartitions.Partition> parts, RUN_TYPE runType) {
+    public void initCdmRun(long runId, long prevRunId, Collection<Partition> parts, RUN_TYPE runType) {
         runStatement.initCdmRun(runId, prevRunId, parts, runType);
         logger.info("###################### Run Id for this job is: {} ######################", runId);
     }

diff --git a/src/main/java/com/datastax/cdm/job/AbstractJobSession.java b/src/main/java/com/datastax/cdm/job/AbstractJobSession.java
@@ -110,15 +110,15 @@ protected AbstractJobSession(CqlSession originSession, CqlSession targetSession,
         }
     }
 
-    public void processSlice(SplitPartitions.Partition slice, TrackRun trackRunFeature, long runId) {
+    public void processSlice(Partition slice, TrackRun trackRunFeature, long runId) {
         this.trackRunFeature = trackRunFeature;
         this.runId = runId;
         this.processSlice(slice.getMin(), slice.getMax());
     }
 
     protected abstract void processSlice(BigInteger min, BigInteger max);
 
-    public synchronized void initCdmRun(long runId, long prevRunId, Collection<SplitPartitions.Partition> parts,
+    public synchronized void initCdmRun(long runId, long prevRunId, Collection<Partition> parts,
             TrackRun trackRunFeature, TrackRun.RUN_TYPE runType) {
         this.runId = runId;
         this.trackRunFeature = trackRunFeature;

diff --git a/src/main/java/com/datastax/cdm/job/CopyJobSession.java b/src/main/java/com/datastax/cdm/job/CopyJobSession.java
@@ -39,7 +39,7 @@
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 
-public class CopyJobSession extends AbstractJobSession<SplitPartitions.Partition> {
+public class CopyJobSession extends AbstractJobSession<Partition> {
 
     private final PKFactory pkFactory;
     private final boolean isCounterTable;

diff --git a/src/main/java/com/datastax/cdm/job/CopyJobSessionFactory.java b/src/main/java/com/datastax/cdm/job/CopyJobSessionFactory.java
@@ -20,11 +20,11 @@
 import com.datastax.cdm.properties.PropertyHelper;
 import com.datastax.oss.driver.api.core.CqlSession;
 
-public class CopyJobSessionFactory implements IJobSessionFactory<SplitPartitions.Partition>, Serializable {
+public class CopyJobSessionFactory implements IJobSessionFactory<Partition>, Serializable {
     private static final long serialVersionUID = 5255029377029801421L;
     private static CopyJobSession jobSession = null;
 
-    public AbstractJobSession<SplitPartitions.Partition> getInstance(CqlSession originSession, CqlSession targetSession,
+    public AbstractJobSession<Partition> getInstance(CqlSession originSession, CqlSession targetSession,
             PropertyHelper propHelper) {
         if (jobSession == null) {
             synchronized (CopyJobSession.class) {

diff --git a/src/main/java/com/datastax/cdm/job/CounterUnit.java b/src/main/java/com/datastax/cdm/job/CounterUnit.java
@@ -0,0 +1,50 @@
+/*
+ * 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 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);
+
+    public void incrementThreadCounter(long incrementBy) {
+        threadLocalCounter.set(threadLocalCounter.get() + incrementBy);
+    }
+
+    public long getThreadCounter() {
+        return threadLocalCounter.get();
+    }
+
+    public void resetThreadCounter() {
+        threadLocalCounter.set(0L);
+    }
+
+    public void setGlobalCounter(long value) {
+        globalCounter.set(value);
+    }
+
+    public void addThreadToGlobalCounter() {
+        globalCounter.addAndGet(threadLocalCounter.get());
+    }
+
+    public long getGlobalCounter() {
+        return globalCounter.get();
+    }
+}
diff --git a/src/main/java/com/datastax/cdm/job/DiffJobSessionFactory.java b/src/main/java/com/datastax/cdm/job/DiffJobSessionFactory.java
@@ -20,11 +20,11 @@
 import com.datastax.cdm.properties.PropertyHelper;
 import com.datastax.oss.driver.api.core.CqlSession;
 
-public class DiffJobSessionFactory implements IJobSessionFactory<SplitPartitions.Partition>, Serializable {
+public class DiffJobSessionFactory implements IJobSessionFactory<Partition>, Serializable {
     private static final long serialVersionUID = -3543616512495020278L;
     private static DiffJobSession jobSession = null;
 
-    public AbstractJobSession<SplitPartitions.Partition> getInstance(CqlSession originSession, CqlSession targetSession,
+    public AbstractJobSession<Partition> getInstance(CqlSession originSession, CqlSession targetSession,
             PropertyHelper propHelper) {
         if (jobSession == null) {
             synchronized (DiffJobSession.class) {

diff --git a/src/main/java/com/datastax/cdm/job/GuardrailCheckJobSession.java b/src/main/java/com/datastax/cdm/job/GuardrailCheckJobSession.java
@@ -27,7 +27,7 @@
 import com.datastax.oss.driver.api.core.cql.ResultSet;
 import com.datastax.oss.driver.api.core.cql.Row;
 
-public class GuardrailCheckJobSession extends AbstractJobSession<SplitPartitions.Partition> {
+public class GuardrailCheckJobSession extends AbstractJobSession<Partition> {
 
     public Logger logger = LoggerFactory.getLogger(this.getClass().getName());
 

diff --git a/src/main/java/com/datastax/cdm/job/GuardrailCheckJobSessionFactory.java b/src/main/java/com/datastax/cdm/job/GuardrailCheckJobSessionFactory.java
@@ -20,11 +20,11 @@
 import com.datastax.cdm.properties.PropertyHelper;
 import com.datastax.oss.driver.api.core.CqlSession;
 
-public class GuardrailCheckJobSessionFactory implements IJobSessionFactory<SplitPartitions.Partition>, Serializable {
+public class GuardrailCheckJobSessionFactory implements IJobSessionFactory<Partition>, Serializable {
     private static final long serialVersionUID = -4673384128807660843L;
     private static GuardrailCheckJobSession jobSession = null;
 
-    public AbstractJobSession<SplitPartitions.Partition> getInstance(CqlSession originSession, CqlSession targetSession,
+    public AbstractJobSession<Partition> getInstance(CqlSession originSession, CqlSession targetSession,
             PropertyHelper propHelper) {
         if (jobSession == null) {
             synchronized (GuardrailCheckJobSession.class) {

diff --git a/src/main/java/com/datastax/cdm/job/JobCounter.java b/src/main/java/com/datastax/cdm/job/JobCounter.java
@@ -15,15 +15,17 @@
  */
 package com.datastax.cdm.job;
 
+import java.io.Serializable;
 import java.util.HashMap;
-import java.util.concurrent.atomic.AtomicLong;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.datastax.cdm.feature.TrackRun;
 
-public class JobCounter {
+public class JobCounter implements Serializable {
+
+    private static final long serialVersionUID = 7016816604237020549L;
 
     // Enumeration for counter types
     public enum CounterType {
@@ -33,36 +35,6 @@ public enum CounterType {
     // Logger instance
     private final Logger logger = LoggerFactory.getLogger(this.getClass().getName());
 
-    // Internal class to handle atomic counting operations
-    private static class CounterUnit {
-        private final AtomicLong globalCounter = new AtomicLong(0);
-        private final ThreadLocal<Long> threadLocalCounter = ThreadLocal.withInitial(() -> 0L);
-
-        public void incrementThreadCounter(long incrementBy) {
-            threadLocalCounter.set(threadLocalCounter.get() + incrementBy);
-        }
-
-        public long getThreadCounter() {
-            return threadLocalCounter.get();
-        }
-
-        public void resetThreadCounter() {
-            threadLocalCounter.set(0L);
-        }
-
-        public void setGlobalCounter(long value) {
-            globalCounter.set(value);
-        }
-
-        public void addThreadToGlobalCounter() {
-            globalCounter.addAndGet(threadLocalCounter.get());
-        }
-
-        public long getGlobalCounter() {
-            return globalCounter.get();
-        }
-    }
-
     // Declare individual counters for different operations
     private final HashMap<CounterType, CounterUnit> counterMap = new HashMap<>();
 

diff --git a/src/main/java/com/datastax/cdm/job/Partition.java b/src/main/java/com/datastax/cdm/job/Partition.java
@@ -0,0 +1,43 @@
+/*
+ * 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 java.io.Serializable;
+import java.math.BigInteger;
+
+public class Partition implements Serializable {
+    private static final long serialVersionUID = 1L;
+
+    private final BigInteger min;
+    private final BigInteger max;
+
+    public Partition(BigInteger min, BigInteger max) {
+        this.min = min;
+        this.max = max;
+    }
+
+    public BigInteger getMin() {
+        return min;
+    }
+
+    public BigInteger getMax() {
+        return max;
+    }
+
+    public String toString() {
+        return "Processing partition for token range " + min + " to " + max;
+    }
+}
diff --git a/src/main/java/com/datastax/cdm/job/SplitPartitions.java b/src/main/java/com/datastax/cdm/job/SplitPartitions.java
@@ -15,7 +15,6 @@
  */
 package com.datastax.cdm.job;
 
-import java.io.Serializable;
 import java.math.BigInteger;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -76,27 +75,4 @@ private static List<Partition> getSubPartitions(int numSplits, BigInteger min, B
         return partitions;
     }
 
-    public static class Partition implements Serializable {
-        private static final long serialVersionUID = 1L;
-
-        private final BigInteger min;
-        private final BigInteger max;
-
-        public Partition(BigInteger min, BigInteger max) {
-            this.min = min;
-            this.max = max;
-        }
-
-        public BigInteger getMin() {
-            return min;
-        }
-
-        public BigInteger getMax() {
-            return max;
-        }
-
-        public String toString() {
-            return "Processing partition for token range " + min + " to " + max;
-        }
-    }
 }
diff --git a/src/main/scala/com/datastax/cdm/job/BasePartitionJob.scala b/src/main/scala/com/datastax/cdm/job/BasePartitionJob.scala
@@ -20,11 +20,11 @@ import scala.reflect.io.File
 import com.datastax.cdm.feature.TrackRun
 import com.datastax.cdm.properties.KnownProperties
 
-abstract class BasePartitionJob extends BaseJob[SplitPartitions.Partition] {
+abstract class BasePartitionJob extends BaseJob[Partition] {
   var trackRunFeature: TrackRun = _
   var keyspaceTableValue: String = _
 
-  override def getParts(pieces: Int): util.Collection[SplitPartitions.Partition] = {
+  override def getParts(pieces: Int): util.Collection[Partition] = {
     var keyspaceTable: Option[String] = Option(propertyHelper.getString(KnownProperties.TARGET_KEYSPACE_TABLE))
       .filter(_.nonEmpty)
       .orElse(Option(propertyHelper.getString(KnownProperties.ORIGIN_KEYSPACE_TABLE)))

diff --git a/src/test/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatementTest.java b/src/test/java/com/datastax/cdm/cql/statement/TargetUpsertRunDetailsStatementTest.java
@@ -32,8 +32,8 @@
 import org.mockito.Mockito;
 
 import com.datastax.cdm.cql.CommonMocks;
+import com.datastax.cdm.job.Partition;
 import com.datastax.cdm.job.RunNotStartedException;
-import com.datastax.cdm.job.SplitPartitions;
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.cql.BoundStatement;
 import com.datastax.oss.driver.api.core.cql.PreparedStatement;
@@ -98,8 +98,7 @@ public void getPartitionsByStatus() {
         when(mockIterator.next()).thenReturn(row3);
 
         targetUpsertRunDetailsStatement = new TargetUpsertRunDetailsStatement(cqlSession, "ks.table1");
-        Collection<SplitPartitions.Partition> parts = targetUpsertRunDetailsStatement.getPartitionsByStatus(123l,
-                "RUNNING");
+        Collection<Partition> parts = targetUpsertRunDetailsStatement.getPartitionsByStatus(123l, "RUNNING");
 
         // This test is incorrect, but needs to be troubleshot & fixed. The actual code works, but the test does not
         assertEquals(0, parts.size());