Skip to content

Commit 49e9ef6

Browse files
arnabnandy7arnabnandy7
committed
feat: add CPU profiling pull mode
Signed-off-by: Arnab Nandy <arnab_nandy7@yahoo.com>
1 parent 5736a04 commit 49e9ef6

15 files changed

Lines changed: 947 additions & 14 deletions

‎agent/build.gradle‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@ shadowJar {
119119
relocate("okhttp3", "io.pyroscope.vendor.okhttp3")
120120
relocate("kotlin", "io.pyroscope.vendor.kotlin")
121121
relocate("com", "io.pyroscope.vendor.com") {
122+
// The pull listener uses the JDK HTTP server, not a bundled dependency.
123+
exclude "com.sun.net.httpserver.**"
122124
// Keep the public OTel instrumentation scope name stable.
123125
exclude "%regex[com\\.grafana\\.pyroscope/java]"
124126
}

‎agent/src/main/java/io/pyroscope/http/Format.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22

33
public enum Format {
44
JFR ("jfr"),
5+
/** CPU profiles served by the pull-mode HTTP endpoint. */
6+
PPROF ("pprof"),
57
/** Experimental and unstable; the OTLP Profiles protocol may change incompatibly. */
68
OTLP ("otlp");
79

‎agent/src/main/java/io/pyroscope/javaagent/AsyncProfilerDelegate.java‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import io.pyroscope.PyroscopeAsyncProfiler;
44
import io.pyroscope.http.Format;
55
import io.pyroscope.javaagent.config.Config;
6+
import io.pyroscope.javaagent.impl.PprofEncoder;
67
import io.pyroscope.javaagent.util.TmpFileUtil;
78
import io.pyroscope.labels.v2.Pyroscope;
89
import one.profiler.AsyncProfiler;
@@ -15,6 +16,8 @@
1516
import java.nio.charset.StandardCharsets;
1617
import java.time.Duration;
1718
import java.time.Instant;
19+
import java.util.HashMap;
20+
import java.util.Map;
1821

1922
import static io.pyroscope.Preconditions.checkNotNull;
2023
import static java.nio.file.Files.newInputStream;
@@ -55,7 +58,7 @@ public void setConfig(@NotNull final Config config) {
5558
*/
5659
@Override
5760
public synchronized void start() {
58-
if (format == Format.JFR || format == Format.OTLP) {
61+
if (format == Format.JFR || format == Format.OTLP || format == Format.PPROF) {
5962
try {
6063
instance.execute(createStartCommand());
6164
} catch (IOException e) {
@@ -112,7 +115,7 @@ static String createStartCommand(Config config, Format format, File tempJFRFile)
112115
if (format == Format.JFR) {
113116
sb.append(",file=").append(tempJFRFile.toString());
114117
}
115-
if (format == Format.JFR || format == Format.OTLP) {
118+
if (format == Format.JFR || format == Format.OTLP || format == Format.PPROF) {
116119
sb.append(",timeout=").append(asyncProfilerTimeoutSeconds(config.uploadInterval));
117120
}
118121
if (config.APLogLevel != null) {
@@ -138,6 +141,16 @@ private Snapshot dumpImpl(Instant started, Instant ended) {
138141
data = dumpJFR();
139142
} else if (format == Format.OTLP) {
140143
data = instance.dumpOtlp(Counter.TOTAL);
144+
} else if (format == Format.PPROF) {
145+
Map<String, String> labels = new HashMap<>(config.timeseries.getLabels());
146+
labels.putAll(config.labels);
147+
labels.putAll(Pyroscope.getStaticLabels());
148+
try {
149+
data = PprofEncoder.encode(instance.dumpCollapsed(Counter.SAMPLES), started, ended,
150+
config.profilingInterval.toNanos(), labels);
151+
} catch (IOException e) {
152+
throw new IllegalStateException(e);
153+
}
141154
} else {
142155
data = instance.dumpCollapsed(Counter.SAMPLES).getBytes(StandardCharsets.UTF_8);
143156
}

‎agent/src/main/java/io/pyroscope/javaagent/PyroscopeAgent.java‎

Lines changed: 36 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import io.pyroscope.javaagent.api.Logger;
66
import io.pyroscope.javaagent.api.ProfilingScheduler;
77
import io.pyroscope.javaagent.config.Config;
8+
import io.pyroscope.javaagent.config.ProfilingMode;
89
import io.pyroscope.javaagent.impl.*;
910
import io.pyroscope.labels.v2.ScopedContext;
1011
import org.jetbrains.annotations.NotNull;
@@ -67,10 +68,20 @@ public static void start(@NotNull Options options) {
6768
}
6869
try {
6970
options.scheduler.start(options.profiler);
70-
ScopedContext.ENABLED.set(true);
71+
ScopedContext.ENABLED.set(options.config.format != Format.PPROF);
72+
if (options.config.format == Format.PPROF) {
73+
logger.log(Logger.Level.WARN,
74+
"Pull profiles do not include dynamic labels or trace context");
75+
}
7176
logger.log(Logger.Level.INFO, "Profiling started");
7277
ProfilerApiPublisher.publish(logger);
7378
} catch (final Throwable e) {
79+
try {
80+
stopComponents(options);
81+
} catch (Throwable cleanupError) {
82+
e.addSuppressed(cleanupError);
83+
}
84+
ScopedContext.ENABLED.set(false);
7485
logger.log(Logger.Level.ERROR, "Error starting profiler %s", e);
7586
sOptions = null;
7687
}
@@ -85,8 +96,7 @@ public static void stop() {
8596
return;
8697
}
8798
try {
88-
sOptions.scheduler.stop();
89-
sOptions.exporter.stop();
99+
stopComponents(sOptions);
90100
sOptions.logger.log(Logger.Level.INFO, "Profiling stopped");
91101
} catch (Throwable e) {
92102
sOptions.logger.log(Logger.Level.ERROR, "Error stopping profiler %s", e);
@@ -102,6 +112,16 @@ public static boolean isStarted() {
102112
}
103113
}
104114

115+
private static void stopComponents(Options options) {
116+
try {
117+
options.scheduler.stop();
118+
} finally {
119+
if (options.exporter != null) {
120+
options.exporter.stop();
121+
}
122+
}
123+
}
124+
105125
/**
106126
* Options allow to swap pyroscope components:
107127
* - io.pyroscope.javaagent.api.ProfilingScheduler
@@ -164,13 +184,20 @@ public Builder setProfiler(@NotNull ProfilerDelegate profiler) {
164184
logger = new DefaultLogger(config.logLevel, System.err);
165185
}
166186
if (scheduler == null) {
167-
if (exporter == null) {
168-
exporter = new QueuedExporter(config, new PyroscopeExporter(config, logger), logger);
169-
}
170-
if (config.samplingDuration == null) {
171-
scheduler = new ContinuousProfilingScheduler(config, exporter, logger);
187+
if (config.profilingMode == ProfilingMode.PULL) {
188+
if (exporter != null) {
189+
throw new IllegalArgumentException("Pull mode does not use an exporter");
190+
}
191+
scheduler = new PullProfilingScheduler(config, logger);
172192
} else {
173-
scheduler = new SamplingProfilingScheduler(config, exporter, logger);
193+
if (exporter == null) {
194+
exporter = new QueuedExporter(config, new PyroscopeExporter(config, logger), logger);
195+
}
196+
if (config.samplingDuration == null) {
197+
scheduler = new ContinuousProfilingScheduler(config, exporter, logger);
198+
} else {
199+
scheduler = new SamplingProfilingScheduler(config, exporter, logger);
200+
}
174201
}
175202
}
176203
if (profiler == null) {

‎agent/src/main/java/io/pyroscope/javaagent/config/AppName.java‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ public Builder newBuilder() {
2929
return new Builder(name, labels);
3030
}
3131

32+
public Map<String, String> getLabels() {
33+
return labels;
34+
}
35+
3236
public static class Builder {
3337
private String name;
3438
private Map<String, String> labels;

‎agent/src/main/java/io/pyroscope/javaagent/config/Config.java‎

Lines changed: 88 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,9 @@
3232
*/
3333
public final class Config {
3434
private static final String PYROSCOPE_AGENT_ENABLED_CONFIG = "PYROSCOPE_AGENT_ENABLED";
35+
private static final String PYROSCOPE_PROFILING_MODE = "PYROSCOPE_PROFILING_MODE";
36+
private static final String PYROSCOPE_PULL_BIND_ADDRESS = "PYROSCOPE_PULL_BIND_ADDRESS";
37+
private static final String PYROSCOPE_PULL_PORT = "PYROSCOPE_PULL_PORT";
3538
private static final String PYROSCOPE_APPLICATION_NAME_CONFIG = "PYROSCOPE_APPLICATION_NAME";
3639
private static final String PYROSCOPE_PROFILING_INTERVAL_CONFIG = "PYROSCOPE_PROFILING_INTERVAL";
3740
private static final String PYROSCOPE_PROFILER_TYPE_CONFIG = "PYROSCOPE_PROFILER_TYPE";
@@ -82,6 +85,8 @@ public final class Config {
8285
private static final String PYROSCOPE_JFR_PROFILER_SETTINGS = "PYROSCOPE_JFR_PROFILER_SETTINGS";
8386

8487
private static final boolean DEFAULT_AGENT_ENABLED = true;
88+
private static final String DEFAULT_PULL_BIND_ADDRESS = "127.0.0.1";
89+
private static final int DEFAULT_PULL_PORT = 4041;
8590
public static final String DEFAULT_SPY_NAME = "javaspy";
8691
private static final Duration DEFAULT_PROFILING_INTERVAL = Duration.ofMillis(10);
8792
private static final EventType DEFAULT_PROFILER_EVENT = EventType.ITIMER;
@@ -104,6 +109,9 @@ public final class Config {
104109
private static final Duration DEFAULT_PROFILE_EXPORT_TIMEOUT = Duration.ofSeconds(10);
105110

106111
public final boolean agentEnabled;
112+
public final ProfilingMode profilingMode;
113+
public final String pullBindAddress;
114+
public final int pullPort;
107115
public final String applicationName;
108116
public final ProfilerType profilerType;
109117
public final Duration profilingInterval;
@@ -171,8 +179,37 @@ public final class Config {
171179
String basicAuthUser,
172180
String basicAuthPassword,
173181
Duration profileExportTimeout,
174-
@Nullable Path tmpDir) {
182+
@Nullable Path tmpDir,
183+
ProfilingMode profilingMode,
184+
String pullBindAddress,
185+
int pullPort) {
175186
this.agentEnabled = agentEnabled;
187+
this.profilingMode = checkNotNull(profilingMode, "profilingMode");
188+
this.pullBindAddress = checkNotNull(pullBindAddress, "pullBindAddress").trim();
189+
this.pullPort = pullPort;
190+
if (profilingMode == ProfilingMode.PULL) {
191+
if (this.pullBindAddress.isEmpty() || pullPort < 0 || pullPort > 65535) {
192+
throw new IllegalArgumentException("Invalid pull listener address or port");
193+
}
194+
if (format != Format.PPROF || profilerType != ProfilerType.ASYNC) {
195+
throw new IllegalArgumentException("Pull mode requires PPROF format and the ASYNC profiler");
196+
}
197+
if ((profilingEvent != EventType.CPU && profilingEvent != EventType.ITIMER) ||
198+
(profilingAlloc != null && !profilingAlloc.isEmpty()) ||
199+
(profilingLock != null && !profilingLock.isEmpty()) || allocLive ||
200+
samplingDuration != null || samplingEventOrder != null) {
201+
throw new IllegalArgumentException("Pull mode supports only CPU or ITIMER profiling without sampling or additional events");
202+
}
203+
if (APExtraArguments != null && !APExtraArguments.isEmpty()) {
204+
throw new IllegalArgumentException("Pull mode does not support AP extra arguments");
205+
}
206+
if (profilingInterval.isZero() || profilingInterval.isNegative()) {
207+
throw new IllegalArgumentException("Pull mode requires a positive profiling interval");
208+
}
209+
profilingInterval.toNanos();
210+
} else if (format == Format.PPROF) {
211+
throw new IllegalArgumentException("PPROF format requires pull mode");
212+
}
176213
this.applicationName = applicationName;
177214
this.profilerType = profilerType;
178215
this.profilingInterval = profilingInterval;
@@ -240,6 +277,9 @@ public long profilingIntervalInHertz() {
240277
public String toString() {
241278
return "Config{" +
242279
"agentEnabled=" + agentEnabled +
280+
", profilingMode=" + profilingMode +
281+
", pullBindAddress='" + pullBindAddress + '\'' +
282+
", pullPort=" + pullPort +
243283
", applicationName='" + applicationName + '\'' +
244284
", profilerType=" + profilerType +
245285
", profilingInterval=" + profilingInterval +
@@ -322,7 +362,10 @@ private static long durationToHertz(Duration duration) {
322362
cp.get(PYROSCOPE_AP_EXTRA_ARGUMENTS_CONFIG),
323363
cp.get(PYROSCOPE_BASIC_AUTH_USER_CONFIG), cp.get(PYROSCOPE_BASIC_AUTH_PASSWORD_CONFIG),
324364
profileExportTimeout(cp),
325-
tmpDir(cp));
365+
tmpDir(cp),
366+
profilingMode(cp),
367+
pullBindAddress(cp),
368+
pullPort(cp));
326369
}
327370

328371
/**
@@ -567,12 +610,30 @@ private static Format format(ConfigurationProvider configurationProvider) {
567610
return Format.JFR;
568611
case "otlp":
569612
return Format.OTLP;
613+
case "pprof":
614+
return Format.PPROF;
570615
default:
571616
DefaultLogger.PRECONFIG_LOGGER.log(Logger.Level.WARN, "Unknown format %s, using %s", format, DEFAULT_FORMAT);
572617
return DEFAULT_FORMAT;
573618
}
574619
}
575620

621+
private static ProfilingMode profilingMode(ConfigurationProvider cp) {
622+
String value = cp.get(PYROSCOPE_PROFILING_MODE);
623+
return value == null || value.isEmpty() ? ProfilingMode.PUSH :
624+
ProfilingMode.valueOf(value.trim().toUpperCase(Locale.ROOT));
625+
}
626+
627+
private static String pullBindAddress(ConfigurationProvider cp) {
628+
String value = cp.get(PYROSCOPE_PULL_BIND_ADDRESS);
629+
return value == null ? DEFAULT_PULL_BIND_ADDRESS : value;
630+
}
631+
632+
private static int pullPort(ConfigurationProvider cp) {
633+
String value = cp.get(PYROSCOPE_PULL_PORT);
634+
return value == null || value.isEmpty() ? DEFAULT_PULL_PORT : Integer.parseInt(value);
635+
}
636+
576637
private static int pushQueueCapacity(ConfigurationProvider configurationProvider) {
577638
final String strPushQueueCapacity = configurationProvider.get(PYROSCOPE_PUSH_QUEUE_CAPACITY_CONFIG);
578639
if (strPushQueueCapacity == null || strPushQueueCapacity.isEmpty()) {
@@ -732,6 +793,9 @@ private static Duration samplingDuration(ConfigurationProvider configurationProv
732793

733794
public static class Builder {
734795
private boolean agentEnabled = DEFAULT_AGENT_ENABLED;
796+
private ProfilingMode profilingMode = ProfilingMode.PUSH;
797+
private String pullBindAddress = DEFAULT_PULL_BIND_ADDRESS;
798+
private int pullPort = DEFAULT_PULL_PORT;
735799
private String applicationName = null;
736800
private ProfilerType profilerType = ProfilerType.ASYNC;
737801
private Duration profilingInterval = DEFAULT_PROFILING_INTERVAL;
@@ -770,6 +834,9 @@ public Builder() {
770834
public Builder(@NotNull Config buildUpon) {
771835
checkNotNull(buildUpon, "config");
772836
agentEnabled = buildUpon.agentEnabled;
837+
profilingMode = buildUpon.profilingMode;
838+
pullBindAddress = buildUpon.pullBindAddress;
839+
pullPort = buildUpon.pullPort;
773840
applicationName = buildUpon.applicationName;
774841
profilerType = buildUpon.profilerType;
775842
profilingInterval = buildUpon.profilingInterval;
@@ -807,6 +874,21 @@ public Builder setAgentEnabled(boolean agentEnabled) {
807874
return this;
808875
}
809876

877+
public Builder setProfilingMode(ProfilingMode profilingMode) {
878+
this.profilingMode = profilingMode;
879+
return this;
880+
}
881+
882+
public Builder setPullBindAddress(String pullBindAddress) {
883+
this.pullBindAddress = pullBindAddress;
884+
return this;
885+
}
886+
887+
public Builder setPullPort(int pullPort) {
888+
this.pullPort = pullPort;
889+
return this;
890+
}
891+
810892
public Builder setApplicationName(String applicationName) {
811893
this.applicationName = applicationName;
812894
return this;
@@ -1013,7 +1095,10 @@ public Builder setTmpDir(@Nullable String tmpDir) {
10131095
APExtraArguments,
10141096
basicAuthUser, basicAuthPassword,
10151097
profileExportTimeout,
1016-
tmpDir);
1098+
tmpDir,
1099+
profilingMode,
1100+
pullBindAddress,
1101+
pullPort);
10171102
}
10181103
}
10191104
}

‎agent/src/main/java/io/pyroscope/javaagent/config/ProfilingMode.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
package io.pyroscope.javaagent.config;
2+
3+
public enum ProfilingMode {
4+
PUSH,
5+
PULL
6+
}

0 commit comments

Comments
 (0)