From 9d5b50561448cbaa1bf1993c5639d5492b22588a Mon Sep 17 00:00:00 2001 From: Dor Amid Date: Fri, 9 Jan 2026 15:26:40 +0200 Subject: [PATCH 1/2] feat: Add thread context propagation for span-profile correlation - Add ProfilingContextStorage wrapper to synchronize OTel context with async-profiler native pthread TLS on every context switch - Fix onEnd() to restore parent span context instead of clearing to (0,0) - Add otel.pyroscope.context.propagation.enabled config flag (default: true) - Add PyroscopeBootstrapConfig for bootstrap classloader configuration Fixes https://github.com/grafana/otel-profiling-java/issues/44 --- build.gradle | 4 + .../pyroscope/ProfilingContextStorage.java | 124 ++++++++++++++++++ .../pyroscope/PyroscopeBootstrapConfig.java | 15 +++ ...elAutoConfigurationCustomizerProvider.java | 12 +- .../pyroscope/PyroscopeOtelConfiguration.java | 11 +- .../pyroscope/PyroscopeOtelSpanProcessor.java | 11 +- ...ling.bootstrap.BootstrapPackagesConfigurer | 1 + 7 files changed, 175 insertions(+), 3 deletions(-) create mode 100644 src/main/java/io/otel/pyroscope/ProfilingContextStorage.java create mode 100644 src/main/java/io/otel/pyroscope/PyroscopeBootstrapConfig.java create mode 100644 src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer diff --git a/build.gradle b/build.gradle index 158a73b..0650a7a 100644 --- a/build.gradle +++ b/build.gradle @@ -11,6 +11,7 @@ ext { versions = [ opentelemetry : "1.33.6", opentelemetryJavaagentAlpha: "1.33.6-alpha", + opentelemetrySdkAutoconfigure: "1.33.0", ] deps = [ @@ -43,9 +44,12 @@ repositories { dependencies { implementation("io.pyroscope:agent:$pyroscopeVersion") + compileOnly("io.opentelemetry:opentelemetry-sdk:${versions.opentelemetry}") + compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:${versions.opentelemetrySdkAutoconfigure}") compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi:${versions.opentelemetry}") compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:${versions.opentelemetry}") compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api:${versions.opentelemetryJavaagentAlpha}") + compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-tooling:${versions.opentelemetryJavaagentAlpha}") testImplementation("io.opentelemetry:opentelemetry-sdk:${versions.opentelemetry}") testImplementation("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi:${versions.opentelemetry}") diff --git a/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java b/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java new file mode 100644 index 0000000..f25aafe --- /dev/null +++ b/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java @@ -0,0 +1,124 @@ +package io.otel.pyroscope; + +import io.opentelemetry.api.trace.Span; +import io.opentelemetry.api.trace.SpanContext; +import io.opentelemetry.context.Context; +import io.opentelemetry.context.ContextStorage; +import io.opentelemetry.context.Scope; +import io.opentelemetry.sdk.trace.ReadableSpan; + +import io.pyroscope.PyroscopeAsyncProfiler; +import io.pyroscope.labels.v2.Pyroscope; +import io.pyroscope.vendor.one.profiler.AsyncProfiler; + +/** + * ContextStorage wrapper that synchronizes OTel context with async-profiler's + * native pthread thread-local storage on every context switch. + * + * This ensures profiling samples are correctly associated with spans, + * even when work is executed on different threads via executors. + */ +public final class ProfilingContextStorage implements ContextStorage { + private final ContextStorage delegate; + private final AsyncProfiler asprof; + + public ProfilingContextStorage(ContextStorage delegate) { + this.delegate = delegate; + this.asprof = PyroscopeAsyncProfiler.getAsyncProfiler(); + } + + @Override + public Scope attach(Context toAttach) { + // 1. Extract span ID and name from the context being attached + long spanId = 0L; + long spanNameId = 0L; + + if (toAttach != null) { + Span span = Span.fromContext(toAttach); + SpanContext spanContext = span.getSpanContext(); + + if (spanContext.isValid()) { + spanId = parseHexSpanId(spanContext.getSpanId()); + + // Register span name to get an ID for native context + if (span instanceof ReadableSpan) { + String spanName = ((ReadableSpan) span).getName(); + spanNameId = Pyroscope.LabelsWrapper.registerConstant(spanName); + } + } + } + + // 2. Set profiling context in native TLS (calls pthread_setspecific) + if (asprof != null) { + asprof.setTracingContext(spanId, spanNameId); + } + + // 3. Delegate to original storage (sets Java ThreadLocal) + Scope originalScope = delegate.attach(toAttach); + + // 4. Return wrapped scope that restores profiling context on close + return new ProfilingScope(originalScope, delegate, asprof); + } + + @Override + public Context current() { + return delegate.current(); + } + + /** + * Parse 16-character hex span ID to long. + */ + private static long parseHexSpanId(String hexSpanId) { + if (hexSpanId == null || hexSpanId.length() != 16) { + return 0L; + } + try { + return Long.parseUnsignedLong(hexSpanId, 16); + } catch (NumberFormatException e) { + return 0L; + } + } + + /** + * Scope wrapper that restores profiling context when closed. + */ + private static final class ProfilingScope implements Scope { + + private final Scope delegate; + private final ContextStorage storage; + private final AsyncProfiler asprof; + + ProfilingScope(Scope delegate, ContextStorage storage, AsyncProfiler asprof) { + this.delegate = delegate; + this.storage = storage; + this.asprof = asprof; + } + + @Override + public void close() { + // 1. Close original scope (restores Java ThreadLocal) + delegate.close(); + + // 2. Restore profiling context to match the now-current OTel context + Context currentContext = storage.current(); + long spanId = 0L; + long spanNameId = 0L; + + if (currentContext != null) { + Span span = Span.fromContext(currentContext); + SpanContext spanContext = span.getSpanContext(); + if (spanContext.isValid()) { + spanId = parseHexSpanId(spanContext.getSpanId()); + if (span instanceof ReadableSpan) { + String spanName = ((ReadableSpan) span).getName(); + spanNameId = Pyroscope.LabelsWrapper.registerConstant(spanName); + } + } + } + + if (asprof != null) { + asprof.setTracingContext(spanId, spanNameId); + } + } + } +} diff --git a/src/main/java/io/otel/pyroscope/PyroscopeBootstrapConfig.java b/src/main/java/io/otel/pyroscope/PyroscopeBootstrapConfig.java new file mode 100644 index 0000000..8b69831 --- /dev/null +++ b/src/main/java/io/otel/pyroscope/PyroscopeBootstrapConfig.java @@ -0,0 +1,15 @@ +package io.otel.pyroscope; + +import io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesBuilder; +import io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer; +import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties; + +/** + * Configures bootstrap classloader to include Pyroscope packages. + */ +public class PyroscopeBootstrapConfig implements BootstrapPackagesConfigurer { + @Override + public void configure(BootstrapPackagesBuilder bootstrapPackagesBuilder, ConfigProperties configProperties) { + bootstrapPackagesBuilder.add("io.pyroscope"); + } +} diff --git a/src/main/java/io/otel/pyroscope/PyroscopeOtelAutoConfigurationCustomizerProvider.java b/src/main/java/io/otel/pyroscope/PyroscopeOtelAutoConfigurationCustomizerProvider.java index 0ab1225..fefcf99 100644 --- a/src/main/java/io/otel/pyroscope/PyroscopeOtelAutoConfigurationCustomizerProvider.java +++ b/src/main/java/io/otel/pyroscope/PyroscopeOtelAutoConfigurationCustomizerProvider.java @@ -1,6 +1,7 @@ package io.otel.pyroscope; +import io.opentelemetry.context.ContextStorage; import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizer; import io.opentelemetry.sdk.autoconfigure.spi.AutoConfigurationCustomizerProvider; import io.pyroscope.javaagent.PyroscopeAgent; @@ -15,7 +16,6 @@ public class PyroscopeOtelAutoConfigurationCustomizerProvider implements AutoConfigurationCustomizerProvider { - @Override public void customize(AutoConfigurationCustomizer autoConfiguration) { autoConfiguration.addTracerProviderCustomizer((tpBuilder, cfg) -> { @@ -42,8 +42,11 @@ public void customize(AutoConfigurationCustomizer autoConfiguration) { PyroscopeOtelConfiguration pyroOtelConfig = new PyroscopeOtelConfiguration.Builder() .setRootSpanOnly(getBoolean(cfg, "otel.pyroscope.root.span.only", true)) .setAddSpanName(getBoolean(cfg, "otel.pyroscope.add.span.name", true)) + .setContextPropagationEnabled(getBoolean(cfg, "otel.pyroscope.context.propagation.enabled", true)) .build(); + registerContextStorageWrapper(pyroOtelConfig); + return tpBuilder.addSpanProcessor( new PyroscopeOtelSpanProcessor( pyroOtelConfig, @@ -78,4 +81,11 @@ private static String getLabelsWrapperClassName() { // otherwise the relocate plugin renames this string :shrug: return new String(Base64.getDecoder().decode("aW8ucHlyb3Njb3BlLmxhYmVscy52Mi5QeXJvc2NvcGUkTGFiZWxzV3JhcHBlcg==")); } + + private static void registerContextStorageWrapper(PyroscopeOtelConfiguration config) { + if (!config.contextPropagationEnabled) { + return; + } + ContextStorage.addWrapper(ProfilingContextStorage::new); + } } \ No newline at end of file diff --git a/src/main/java/io/otel/pyroscope/PyroscopeOtelConfiguration.java b/src/main/java/io/otel/pyroscope/PyroscopeOtelConfiguration.java index f122df5..e377aba 100644 --- a/src/main/java/io/otel/pyroscope/PyroscopeOtelConfiguration.java +++ b/src/main/java/io/otel/pyroscope/PyroscopeOtelConfiguration.java @@ -5,23 +5,27 @@ public class PyroscopeOtelConfiguration { //todo think about removing both options, so that users don't need to configure or think about anything final boolean rootSpanOnly; final boolean addSpanName; + final boolean contextPropagationEnabled; private PyroscopeOtelConfiguration(Builder builder) { this.rootSpanOnly = builder.rootSpanOnly; this.addSpanName = builder.addSpanName; + this.contextPropagationEnabled = builder.contextPropagationEnabled; } @Override public String toString() { return "PyroscopeOtelConfiguration{" + - ", rootSpanOnly=" + rootSpanOnly + + "rootSpanOnly=" + rootSpanOnly + ", addSpanName=" + addSpanName + + ", contextPropagationEnabled=" + contextPropagationEnabled + '}'; } public static class Builder { boolean rootSpanOnly = true; boolean addSpanName = true; + boolean contextPropagationEnabled = true; public Builder() { } @@ -36,6 +40,11 @@ public Builder setAddSpanName(boolean addSpanName) { return this; } + public Builder setContextPropagationEnabled(boolean contextPropagationEnabled) { + this.contextPropagationEnabled = contextPropagationEnabled; + return this; + } + public PyroscopeOtelConfiguration build() { return new PyroscopeOtelConfiguration(this); } diff --git a/src/main/java/io/otel/pyroscope/PyroscopeOtelSpanProcessor.java b/src/main/java/io/otel/pyroscope/PyroscopeOtelSpanProcessor.java index 05b89c3..3b7d58d 100644 --- a/src/main/java/io/otel/pyroscope/PyroscopeOtelSpanProcessor.java +++ b/src/main/java/io/otel/pyroscope/PyroscopeOtelSpanProcessor.java @@ -61,7 +61,16 @@ public void onStart(Context parentContext, ReadWriteSpan span) { @Override public void onEnd(ReadableSpan span) { - setTracingContextForSpan(0, 0); + // Restore parent span context instead of clearing completely + SpanContext parentContext = span.getParentSpanContext(); + if (parentContext.isValid() && !parentContext.isRemote()) { + long parentSpanId = parseSpanId(parentContext.getSpanId()); + // Note: we don't have parent span name, so use 0 + setTracingContextForSpan(parentSpanId, 0); + } else { + // Root span ended, clear context + setTracingContextForSpan(0, 0); + } } private void setTracingContextForSpan(long spanId, long spanName) { diff --git a/src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer b/src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer new file mode 100644 index 0000000..2399b3b --- /dev/null +++ b/src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer @@ -0,0 +1 @@ +io.otel.pyroscope.PyroscopeBootstrapConfig From 96c5dd4fac863d61cd851b5c407c80a30afe7566 Mon Sep 17 00:00:00 2001 From: dordor12 <50987668+dordor12@users.noreply.github.com> Date: Sun, 8 Feb 2026 09:43:17 +0200 Subject: [PATCH 2/2] refactor: address PR review feedback - Mark API as EXPERIMENTAL - Clarify relationship with SpanProcessor in javadoc - Remove redundant comments - Reuse parseSpanId from PyroscopeOtelSpanProcessor Co-authored-by: Claude --- .../pyroscope/ProfilingContextStorage.java | 39 ++++--------------- 1 file changed, 7 insertions(+), 32 deletions(-) diff --git a/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java b/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java index f25aafe..6807efc 100644 --- a/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java +++ b/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java @@ -12,11 +12,11 @@ import io.pyroscope.vendor.one.profiler.AsyncProfiler; /** - * ContextStorage wrapper that synchronizes OTel context with async-profiler's - * native pthread thread-local storage on every context switch. + * EXPERIMENTAL: This API is experimental and may change or be removed in future versions. * - * This ensures profiling samples are correctly associated with spans, - * even when work is executed on different threads via executors. + * Complements PyroscopeOtelSpanProcessor for wall-clock profiling. SpanProcessor handles + * span start/end on the originating thread, while this ContextStorage wrapper handles + * context propagation to other threads (e.g., via executors), keeping native TLS in sync. */ public final class ProfilingContextStorage implements ContextStorage { private final ContextStorage delegate; @@ -29,7 +29,6 @@ public ProfilingContextStorage(ContextStorage delegate) { @Override public Scope attach(Context toAttach) { - // 1. Extract span ID and name from the context being attached long spanId = 0L; long spanNameId = 0L; @@ -38,9 +37,8 @@ public Scope attach(Context toAttach) { SpanContext spanContext = span.getSpanContext(); if (spanContext.isValid()) { - spanId = parseHexSpanId(spanContext.getSpanId()); + spanId = PyroscopeOtelSpanProcessor.parseSpanId(spanContext.getSpanId()); - // Register span name to get an ID for native context if (span instanceof ReadableSpan) { String spanName = ((ReadableSpan) span).getName(); spanNameId = Pyroscope.LabelsWrapper.registerConstant(spanName); @@ -48,15 +46,11 @@ public Scope attach(Context toAttach) { } } - // 2. Set profiling context in native TLS (calls pthread_setspecific) if (asprof != null) { asprof.setTracingContext(spanId, spanNameId); } - // 3. Delegate to original storage (sets Java ThreadLocal) Scope originalScope = delegate.attach(toAttach); - - // 4. Return wrapped scope that restores profiling context on close return new ProfilingScope(originalScope, delegate, asprof); } @@ -65,25 +59,7 @@ public Context current() { return delegate.current(); } - /** - * Parse 16-character hex span ID to long. - */ - private static long parseHexSpanId(String hexSpanId) { - if (hexSpanId == null || hexSpanId.length() != 16) { - return 0L; - } - try { - return Long.parseUnsignedLong(hexSpanId, 16); - } catch (NumberFormatException e) { - return 0L; - } - } - - /** - * Scope wrapper that restores profiling context when closed. - */ private static final class ProfilingScope implements Scope { - private final Scope delegate; private final ContextStorage storage; private final AsyncProfiler asprof; @@ -96,10 +72,8 @@ private static final class ProfilingScope implements Scope { @Override public void close() { - // 1. Close original scope (restores Java ThreadLocal) delegate.close(); - // 2. Restore profiling context to match the now-current OTel context Context currentContext = storage.current(); long spanId = 0L; long spanNameId = 0L; @@ -108,7 +82,8 @@ public void close() { Span span = Span.fromContext(currentContext); SpanContext spanContext = span.getSpanContext(); if (spanContext.isValid()) { - spanId = parseHexSpanId(spanContext.getSpanId()); + spanId = PyroscopeOtelSpanProcessor.parseSpanId(spanContext.getSpanId()); + if (span instanceof ReadableSpan) { String spanName = ((ReadableSpan) span).getName(); spanNameId = Pyroscope.LabelsWrapper.registerConstant(spanName);