feat: Add thread context propagation for span-profile correlation in thread pools

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}")

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);
+            }
+        }
+    }
+}

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");
+    }
+}

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);
+    }
 }

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);
         }

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) {

src/main/resources/META-INF/services/io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer

@@ -0,0 +1 @@
+io.otel.pyroscope.PyroscopeBootstrapConfig