Some misc refactoring

5 files changed, 28 insertions, 24 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/ThreadGrouper.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/ThreadGrouper.java
index d70ba1a..e63ebc8 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/ThreadGrouper.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/ThreadGrouper.java
@@ -30,17 +30,22 @@ import java.util.regex.Pattern;
 /**
  * Function for grouping threads together
  */
-public enum ThreadGrouper {
+public interface ThreadGrouper {
 
     /**
      * Implementation of {@link ThreadGrouper} that just groups by thread name.
      */
-    BY_NAME {
+    ThreadGrouper BY_NAME = new ThreadGrouper() {
         @Override
         public String getGroup(long threadId, String threadName) {
             return threadName;
         }
-    },
+
+        @Override
+        public SamplerMetadata.DataAggregator.ThreadGrouper asProto() {
+            return SamplerMetadata.DataAggregator.ThreadGrouper.BY_NAME;
+        }
+    };
 
     /**
      * Implementation of {@link ThreadGrouper} that attempts to group by the name of the pool
@@ -49,7 +54,7 @@ public enum ThreadGrouper {
      * <p>The regex pattern used to match pools expects a digit at the end of the thread name,
      * separated from the pool name with any of one or more of ' ', '-', or '#'.</p>
      */
-    BY_POOL {
+    ThreadGrouper BY_POOL = new ThreadGrouper() {
         private final Map<Long, String> cache = new ConcurrentHashMap<>();
         private final Pattern pattern = Pattern.compile("^(.*?)[-# ]+\\d+$");
 
@@ -69,17 +74,27 @@ public enum ThreadGrouper {
             this.cache.put(threadId, group); // we don't care about race conditions here
             return group;
         }
-    },
+
+        @Override
+        public SamplerMetadata.DataAggregator.ThreadGrouper asProto() {
+            return SamplerMetadata.DataAggregator.ThreadGrouper.BY_POOL;
+        }
+    };
 
     /**
      * Implementation of {@link ThreadGrouper} which groups all threads as one, under
      * the name "All".
      */
-    AS_ONE {
+    ThreadGrouper AS_ONE = new ThreadGrouper() {
         @Override
         public String getGroup(long threadId, String threadName) {
             return "All";
         }
+
+        @Override
+        public SamplerMetadata.DataAggregator.ThreadGrouper asProto() {
+            return SamplerMetadata.DataAggregator.ThreadGrouper.AS_ONE;
+        }
     };
 
     /**
@@ -89,19 +104,8 @@ public enum ThreadGrouper {
      * @param threadName the name of the thread
      * @return the group
      */
-    public abstract String getGroup(long threadId, String threadName);
-
-    public static SamplerMetadata.DataAggregator.ThreadGrouper asProto(ThreadGrouper threadGrouper) {
-        switch (threadGrouper) {
-            case BY_NAME:
-                return SamplerMetadata.DataAggregator.ThreadGrouper.BY_NAME;
-            case BY_POOL:
-                return SamplerMetadata.DataAggregator.ThreadGrouper.BY_POOL;
-            case AS_ONE:
-                return SamplerMetadata.DataAggregator.ThreadGrouper.AS_ONE;
-            default:
-                throw new AssertionError();
-        }
-    }
+    String getGroup(long threadId, String threadName);
+
+    SamplerMetadata.DataAggregator.ThreadGrouper asProto();
 
 }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/SimpleDataAggregator.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/SimpleDataAggregator.java
index db74995..4272ccd 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/SimpleDataAggregator.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/SimpleDataAggregator.java
@@ -41,7 +41,7 @@ public class SimpleDataAggregator extends AbstractDataAggregator {
     public SamplerMetadata.DataAggregator getMetadata() {
         return SamplerMetadata.DataAggregator.newBuilder()
                 .setType(SamplerMetadata.DataAggregator.Type.SIMPLE)
-                .setThreadGrouper(ThreadGrouper.asProto(this.threadGrouper))
+                .setThreadGrouper(this.threadGrouper.asProto())
                 .build();
     }
 
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/TickedDataAggregator.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/TickedDataAggregator.java
index 567c865..b5a8aad 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/TickedDataAggregator.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/aggregator/TickedDataAggregator.java
@@ -66,7 +66,7 @@ public class TickedDataAggregator extends AbstractDataAggregator {
     public SamplerMetadata.DataAggregator getMetadata() {
         return SamplerMetadata.DataAggregator.newBuilder()
                 .setType(SamplerMetadata.DataAggregator.Type.TICKED)
-                .setThreadGrouper(ThreadGrouper.asProto(this.threadGrouper))
+                .setThreadGrouper(this.threadGrouper.asProto())
                 .setTickLengthThreshold(this.tickLengthThreshold)
                 .build();
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/AbstractTickHook.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/AbstractTickHook.java
index b4e5054..72e927f 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/AbstractTickHook.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/AbstractTickHook.java
@@ -30,7 +30,7 @@ public abstract class AbstractTickHook implements TickHook {
 
     protected void onTick() {
         for (Callback r : this.tasks) {
-            r.onTick(this);
+            r.onTick(this.tick);
         }
         this.tick++;
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/TickHook.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/TickHook.java
index a0dd425..8731216 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/TickHook.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/tick/TickHook.java
@@ -58,7 +58,7 @@ public interface TickHook extends AutoCloseable {
     void removeCallback(Callback runnable);
 
     interface Callback {
-        void onTick(TickHook hook);
+        void onTick(int currentTick);
     }
 
 }