Skip to content

Commit

Permalink
Include engine type in sampler proto
Browse files Browse the repository at this point in the history
  • Loading branch information
lucko committed Sep 3, 2024
1 parent 684cc5e commit ed33fd5
Show file tree
Hide file tree
Showing 7 changed files with 79 additions and 18 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -179,6 +179,7 @@ protected void sendStatisticsToSocket() {

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)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,13 @@ public interface Sampler {
*/
boolean isRunningInBackground();

/**
* Gets the sampler type.
*
* @return the sampler type
*/
SamplerType getType();

/**
* Gets the sampler mode.
*
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
/*
* This file is part of spark.
*
* Copyright (c) lucko (Luck) <[email protected]>
* 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;
}

}
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
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;
Expand Down Expand Up @@ -209,6 +210,11 @@ public void attachSocket(ViewerSocket socket) {
}
}

@Override
public SamplerType getType() {
return SamplerType.ASYNC;
}

@Override
public SamplerMode getMode() {
return this.sampleCollector.getMode();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
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;
Expand Down Expand Up @@ -202,6 +203,11 @@ public SamplerData toProto(SparkPlatform platform, ExportProps exportProps) {
return proto.build();
}

@Override
public SamplerType getType() {
return SamplerType.JAVA;
}

@Override
public SamplerMode getMode() {
return SamplerMode.EXECUTION;
Expand Down
6 changes: 6 additions & 0 deletions spark-common/src/main/proto/spark/spark_sampler.proto
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -69,6 +70,11 @@ message SamplerMetadata {
EXECUTION = 0;
ALLOCATION = 1;
}

enum SamplerEngine {
JAVA = 0;
ASYNC = 1;
}
}

message ThreadNode {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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");
}
Expand All @@ -62,11 +60,12 @@ public void testSampler(ProfilerType profilerType, @TempDir Path directory) {
.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();
Expand All @@ -75,7 +74,7 @@ public void testSampler(ProfilerType profilerType, @TempDir Path directory) {
.creator(TestCommandSender.INSTANCE.toData())
.classSourceLookup(() -> ClassSourceLookup.create(plugin.platform()));

if (profilerType == ProfilerType.JAVA) {
if (samplerType == SamplerType.JAVA) {
exportProps.mergeStrategy(MergeStrategy.SAME_METHOD);
}

Expand All @@ -93,15 +92,4 @@ public void testSampler(ProfilerType profilerType, @TempDir Path directory) {
}
}

public enum ProfilerType {
JAVA(JavaSampler.class),
ASYNC(AsyncSampler.class);

private final Class<? extends Sampler> expectedClass;

ProfilerType(Class<? extends Sampler> expectedClass) {
this.expectedClass = expectedClass;
}
}

}

0 comments on commit ed33fd5

Please sign in to comment.