Include engine type in sampler proto

5 files changed, 67 insertions, 0 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
index d76b1a1..aecdc71 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/AbstractSampler.java
@@ -179,6 +179,7 @@ public abstract class AbstractSampler implements Sampler {
 
     protected void writeMetadataToProto(SamplerData.Builder proto, SparkPlatform platform, CommandSender.Data creator, String comment, DataAggregator dataAggregator) {
         SamplerMetadata.Builder metadata = SamplerMetadata.newBuilder()
+                .setSamplerEngine(getType().asProto())
                 .setSamplerMode(getMode().asProto())
                 .setStartTime(this.startTime)
                 .setInterval(this.interval)
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
index 5aca704..71ab039 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/Sampler.java
@@ -85,6 +85,13 @@ public interface Sampler {
     boolean isRunningInBackground();
 
     /**
+     * Gets the sampler type.
+     *
+     * @return the sampler type
+     */
+    SamplerType getType();
+
+    /**
      * Gets the sampler mode.
      *
      * @return the sampler mode
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerType.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerType.java
new file mode 100644
index 0000000..aad4b23
--- /dev/null
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/SamplerType.java
@@ -0,0 +1,47 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.common.sampler;
+
+import me.lucko.spark.common.sampler.async.AsyncSampler;
+import me.lucko.spark.common.sampler.java.JavaSampler;
+import me.lucko.spark.proto.SparkSamplerProtos.SamplerMetadata;
+
+public enum SamplerType {
+    JAVA(JavaSampler.class, SamplerMetadata.SamplerEngine.JAVA),
+    ASYNC(AsyncSampler.class, SamplerMetadata.SamplerEngine.ASYNC);
+
+    private final Class<? extends Sampler> expectedClass;
+    private final SamplerMetadata.SamplerEngine proto;
+
+    SamplerType(Class<? extends Sampler> expectedClass, SamplerMetadata.SamplerEngine proto) {
+        this.expectedClass = expectedClass;
+        this.proto = proto;
+    }
+
+    public Class<? extends Sampler> implClass() {
+        return this.expectedClass;
+    }
+
+    public SamplerMetadata.SamplerEngine asProto() {
+        return this.proto;
+    }
+
+}
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
index 62af021..994c03b 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/async/AsyncSampler.java
@@ -25,6 +25,7 @@ import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.sampler.AbstractSampler;
 import me.lucko.spark.common.sampler.SamplerMode;
 import me.lucko.spark.common.sampler.SamplerSettings;
+import me.lucko.spark.common.sampler.SamplerType;
 import me.lucko.spark.common.sampler.window.ProfilingWindowUtils;
 import me.lucko.spark.common.tick.TickHook;
 import me.lucko.spark.common.util.SparkThreadFactory;
@@ -210,6 +211,11 @@ public class AsyncSampler extends AbstractSampler {
     }
 
     @Override
+    public SamplerType getType() {
+        return SamplerType.ASYNC;
+    }
+
+    @Override
     public SamplerMode getMode() {
         return this.sampleCollector.getMode();
     }
diff --git a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
index 20f9383..050c5b4 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/sampler/java/JavaSampler.java
@@ -25,6 +25,7 @@ import me.lucko.spark.common.SparkPlatform;
 import me.lucko.spark.common.sampler.AbstractSampler;
 import me.lucko.spark.common.sampler.SamplerMode;
 import me.lucko.spark.common.sampler.SamplerSettings;
+import me.lucko.spark.common.sampler.SamplerType;
 import me.lucko.spark.common.sampler.window.ProfilingWindowUtils;
 import me.lucko.spark.common.sampler.window.WindowStatisticsCollector;
 import me.lucko.spark.common.tick.TickHook;
@@ -203,6 +204,11 @@ public class JavaSampler extends AbstractSampler implements Runnable {
     }
 
     @Override
+    public SamplerType getType() {
+        return SamplerType.JAVA;
+    }
+
+    @Override
     public SamplerMode getMode() {
         return SamplerMode.EXECUTION;
     }