Include engine type in sampler proto

7 files changed, 79 insertions, 18 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;
     }
diff --git a/spark-common/src/main/proto/spark/spark_sampler.proto b/spark-common/src/main/proto/spark/spark_sampler.proto
index 10cc6d1..bd48e7d 100644
--- a/spark-common/src/main/proto/spark/spark_sampler.proto
+++ b/spark-common/src/main/proto/spark/spark_sampler.proto
@@ -34,6 +34,7 @@ message SamplerMetadata {
   map<string, PluginOrModMetadata> sources = 13;
   map<string, string> extra_platform_metadata = 14;
   SamplerMode sampler_mode = 15;
+  SamplerEngine sampler_engine = 16;
 
   message ThreadDumper {
     Type type = 1;
@@ -69,6 +70,11 @@ message SamplerMetadata {
     EXECUTION = 0;
     ALLOCATION = 1;
   }
+
+  enum SamplerEngine {
+    JAVA = 0;
+    ASYNC = 1;
+  }
 }
 
 message ThreadNode {
diff --git a/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java b/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
index f6d9f87..53fc7e6 100644
--- a/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
+++ b/spark-common/src/test/java/me/lucko/spark/common/sampler/SamplerTest.java
@@ -20,8 +20,6 @@
 
 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.common.sampler.java.MergeStrategy;
 import me.lucko.spark.common.sampler.source.ClassSourceLookup;
 import me.lucko.spark.proto.SparkSamplerProtos;
@@ -48,8 +46,8 @@ public class SamplerTest {
 
     @ParameterizedTest
     @EnumSource
-    public void testSampler(ProfilerType profilerType, @TempDir Path directory) {
-        if (profilerType == ProfilerType.ASYNC) {
+    public void testSampler(SamplerType samplerType, @TempDir Path directory) {
+        if (samplerType == SamplerType.ASYNC) {
             String os = System.getProperty("os.name").toLowerCase(Locale.ROOT).replace(" ", "");
             assumeTrue(os.equals("linux") || os.equals("macosx"), "async profiler is only supported on Linux and macOS");
         }
@@ -62,11 +60,12 @@ public class SamplerTest {
                     .threadDumper(new ThreadDumper.Specific(thread))
                     .threadGrouper(ThreadGrouper.BY_POOL)
                     .samplingInterval(10)
-                    .forceJavaSampler(profilerType == ProfilerType.JAVA)
+                    .forceJavaSampler(samplerType == SamplerType.JAVA)
                     .completeAfter(2, TimeUnit.SECONDS)
                     .start(plugin.platform());
 
-            assertInstanceOf(profilerType.expectedClass, sampler);
+            assertInstanceOf(samplerType.implClass(), sampler);
+            assertEquals(samplerType, sampler.getType());
 
             assertNotEquals(-1, sampler.getAutoEndTime());
             sampler.getFuture().join();
@@ -75,7 +74,7 @@ public class SamplerTest {
                     .creator(TestCommandSender.INSTANCE.toData())
                     .classSourceLookup(() -> ClassSourceLookup.create(plugin.platform()));
 
-            if (profilerType == ProfilerType.JAVA) {
+            if (samplerType == SamplerType.JAVA) {
                 exportProps.mergeStrategy(MergeStrategy.SAME_METHOD);
             }
 
@@ -93,15 +92,4 @@ public class SamplerTest {
         }
     }
 
-    public enum ProfilerType {
-        JAVA(JavaSampler.class),
-        ASYNC(AsyncSampler.class);
-
-        private final Class<? extends Sampler> expectedClass;
-
-        ProfilerType(Class<? extends Sampler> expectedClass) {
-            this.expectedClass = expectedClass;
-        }
-    }
-
 }