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..6807efc --- /dev/null +++ b/src/main/java/io/otel/pyroscope/ProfilingContextStorage.java @@ -0,0 +1,99 @@ +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; + +/** + * EXPERIMENTAL: This API is experimental and may change or be removed in future versions. + * + * 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; + private final AsyncProfiler asprof; + + public ProfilingContextStorage(ContextStorage delegate) { + this.delegate = delegate; + this.asprof = PyroscopeAsyncProfiler.getAsyncProfiler(); + } + + @Override + public Scope attach(Context toAttach) { + long spanId = 0L; + long spanNameId = 0L; + + if (toAttach != null) { + Span span = Span.fromContext(toAttach); + SpanContext spanContext = span.getSpanContext(); + + if (spanContext.isValid()) { + spanId = PyroscopeOtelSpanProcessor.parseSpanId(spanContext.getSpanId()); + + if (span instanceof ReadableSpan) { + String spanName = ((ReadableSpan) span).getName(); + spanNameId = Pyroscope.LabelsWrapper.registerConstant(spanName); + } + } + } + + if (asprof != null) { + asprof.setTracingContext(spanId, spanNameId); + } + + Scope originalScope = delegate.attach(toAttach); + return new ProfilingScope(originalScope, delegate, asprof); + } + + @Override + public Context current() { + return delegate.current(); + } + + 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() { + delegate.close(); + + 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 = PyroscopeOtelSpanProcessor.parseSpanId(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