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

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();
+    }
+}

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<>();