Merge pull request #13 from grafana/fix/extension-sdk-interference

Fix interference between extension and Pyroscope SDK

build.gradle

@@ -10,8 +10,6 @@ plugins {
 ext {
     versions = [
             opentelemetry              : "1.17.0",
-
-
             opentelemetryJavaagentAlpha: "1.17.0-alpha",
     ]
 
@@ -48,10 +46,6 @@ dependencies {
     compileOnly("io.opentelemetry:opentelemetry-sdk-extension-autoconfigure-spi:${versions.opentelemetry}")
     compileOnly("io.opentelemetry.instrumentation:opentelemetry-instrumentation-api:${versions.opentelemetryJavaagentAlpha}")
     compileOnly("io.opentelemetry.javaagent:opentelemetry-javaagent-extension-api:${versions.opentelemetryJavaagentAlpha}")
-
-
-
-
 }
 
 java {
@@ -60,12 +54,14 @@ java {
 }
 
 jar {
-    
+
 }
 
 shadowJar {
     archiveFileName = "pyroscope-otel.jar"
-    relocate 'io.pyroscope', 'io.otel.pyroscope.shadow'
+    relocate("io.pyroscope", "io.otel.pyroscope.shadow") {
+        exclude 'io.pyroscope.javaagent.ProfilerSdk'
+    }
     archiveClassifier.set('')
 }
 

gradle.properties

@@ -1,2 +1,2 @@
-pyroscope_version=0.13.0
+pyroscope_version=0.16.0
 otel_profiling_version=0.10.3

src/main/java/io/otel/pyroscope/OtelProfilerSdkBridge.java

@@ -0,0 +1,71 @@
+package io.otel.pyroscope;
+
+import io.pyroscope.javaagent.api.ProfilerApi;
+import io.pyroscope.javaagent.api.ProfilerScopedContext;
+
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.function.BiConsumer;
+
+public class OtelProfilerSdkBridge implements ProfilerApi {
+
+    private final Object sdkInstance;
+
+    public OtelProfilerSdkBridge(Object sdkInstance) {
+        this.sdkInstance = sdkInstance;
+    }
+
+    @Override
+    public void startProfiling() {
+        try {
+            Method method = sdkInstance.getClass().getDeclaredMethod("startProfiling");
+            method.invoke(sdkInstance);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public boolean isProfilingStarted() {
+        try {
+            Method method = sdkInstance.getClass().getDeclaredMethod("isProfilingStarted");
+            return (boolean) method.invoke(sdkInstance);
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    @Override
+    public ProfilerScopedContext createScopedContext(Map<String, String> labels) {
+        try {
+            Method method = sdkInstance.getClass().getDeclaredMethod("createScopedContext", Map.class);
+            Object ctx = method.invoke(sdkInstance, labels);
+
+            return new ProfilerScopedContext() {
+
+                @Override
+                public void forEachLabel(BiConsumer<String, String> biConsumer) {
+                    try {
+                        Method forEachLabelMethod = ctx.getClass().getDeclaredMethod("forEachLabel", BiConsumer.class);
+                        forEachLabelMethod.invoke(ctx, biConsumer);
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+
+                @Override
+                public void close() {
+                    try {
+                        Method closeMethod = ctx.getClass().getDeclaredMethod("close");
+                        closeMethod.invoke(ctx);
+                    } catch (Exception e) {
+                        throw new RuntimeException(e);
+                    }
+                }
+            };
+
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+}

src/main/java/io/otel/pyroscope/PyroscopeOtelAutoConfigurationCustomizerProvider.java

@@ -14,29 +14,58 @@ public class PyroscopeOtelAutoConfigurationCustomizerProvider
 
     public static final String CONFIG_APP_NAME = "otel.pyroscope.app.name";
     public static final String CONFIG_ENDPOINT = "otel.pyroscope.endpoint";
-
-    private static final String PYROSCOPE_APPLICATION_NAME_CONFIG = "pyroscope.application.name";
-    private static final String PYROSCOPE_SERVER_ADDRESS_CONFIG = "pyroscope.server.address";
     public static final String CONFIG_BASELINE_LABELS = "otel.pyroscope.baseline.labels";
 
     @Override
     public void customize(AutoConfigurationCustomizer autoConfiguration) {
         autoConfiguration.addTracerProviderCustomizer((tpBuilder, cfg) -> {
+            OtelProfilerSdkBridge profilerSdk = null;
+            try {
+                profilerSdk = loadProfilerSdk();
+            } catch (Exception e) {
+                // This usually means we are running without the Pyroscope SDK.
+                // We'll instead use the Profiler bundled with the extension.
+                System.out.println("Could not load the profiler SDK, will continue with the built-in one!");
+            }
+
             boolean startProfiling = getBoolean(cfg, "otel.pyroscope.start.profiling", true);
             if (startProfiling) {
-                Config pyroConfig = Config.build();
-                PyroscopeAgent.start(pyroConfig);
+                if (profilerSdk == null) {
+                    PyroscopeAgent.start(Config.build());
+                } else if (!profilerSdk.isProfilingStarted()) {
+                    profilerSdk.startProfiling();
+                }
             }
+
             PyroscopeOtelConfiguration pyroOtelConfig = new PyroscopeOtelConfiguration.Builder()
                     .setRootSpanOnly(getBoolean(cfg, "otel.pyroscope.root.span.only", true))
                     .setAddSpanName(getBoolean(cfg, "otel.pyroscope.add.span.name", true))
                     .build();
+
             return tpBuilder.addSpanProcessor(
                     new PyroscopeOtelSpanProcessor(
-                            pyroOtelConfig
+                            pyroOtelConfig,
+                            profilerSdk
                     ));
         });
+    }
 
+    /**
+     * Open Telemetry extension classes are loaded by an isolated class loader.
+     * As such, they can't communicate with other parts of the application (e.g., the Pyroscope SDK).
+     *
+     * If the Pyroscope SDK is loaded as a java agent, we'll access it via the system class loader and interact with it
+     * via a bridge.
+     */
+    private static OtelProfilerSdkBridge loadProfilerSdk() {
+        try {
+            ClassLoader systemClassLoader = ClassLoader.getSystemClassLoader();
+            Class<?> sdkClass = systemClassLoader.loadClass("io.pyroscope.javaagent.ProfilerSdk");
+            Object sdk = sdkClass.getDeclaredConstructor().newInstance();
+            return new OtelProfilerSdkBridge(sdk);
+        } catch (Exception e) {
+            throw new RuntimeException("Error loading the profiler SDK", e);
+        }
     }
 
 }

src/main/java/io/otel/pyroscope/PyroscopeOtelSpanProcessor.java

@@ -6,6 +6,9 @@
 import io.opentelemetry.sdk.trace.ReadWriteSpan;
 import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.sdk.trace.SpanProcessor;
+
+import io.pyroscope.javaagent.api.ProfilerScopedContext;
+import io.pyroscope.javaagent.impl.ProfilerScopedContextWrapper;
 import io.pyroscope.labels.LabelsSet;
 import io.pyroscope.labels.ScopedContext;
 
@@ -32,9 +35,11 @@ public class PyroscopeOtelSpanProcessor implements SpanProcessor {
     private final Map<String, PyroscopeContextHolder> pyroscopeContexts = new ConcurrentHashMap<>();
 
     private final PyroscopeOtelConfiguration configuration;
+    private final OtelProfilerSdkBridge profilerSdk;
 
-    public PyroscopeOtelSpanProcessor(PyroscopeOtelConfiguration configuration) {
+    public PyroscopeOtelSpanProcessor(PyroscopeOtelConfiguration configuration, OtelProfilerSdkBridge profilerSdk) {
         this.configuration = configuration;
+        this.profilerSdk = profilerSdk;
     }
 
     @Override
@@ -59,10 +64,10 @@ public void onStart(Context parentContext, ReadWriteSpan span) {
             labels.put(LABEL_SPAN_NAME, span.getName());
         }
 
-        ScopedContext pyroscopeContext = new ScopedContext(new LabelsSet(labels));
+        ProfilerScopedContext scopedContext = createScopedContext(labels);
         span.setAttribute(ATTRIBUTE_KEY_PROFILE_ID, profileId);
         long now = now();
-        PyroscopeContextHolder pyroscopeContextHolder = new PyroscopeContextHolder(profileId, pyroscopeContext, now);
+        PyroscopeContextHolder pyroscopeContextHolder = new PyroscopeContextHolder(profileId, scopedContext, now);
         pyroscopeContexts.put(profileId, pyroscopeContextHolder);
         if (configuration.optimisticTimestamps) {
             long optimisticEnd = now + TimeUnit.HOURS.toMillis(1);
@@ -77,6 +82,13 @@ public void onStart(Context parentContext, ReadWriteSpan span) {
         }
     }
 
+    private ProfilerScopedContext createScopedContext(Map<String, String> labels) {
+        if (profilerSdk == null) {
+            return new ProfilerScopedContextWrapper(new ScopedContext(new LabelsSet(labels)));
+        }
+        return profilerSdk.createScopedContext(labels);
+    }
+
     @Override
     public void onEnd(ReadableSpan span) {
         String profileId = span.getAttribute(ATTRIBUTE_KEY_PROFILE_ID);
@@ -110,13 +122,11 @@ public void onEnd(ReadableSpan span) {
         } finally {
             pyroscopeContext.ctx.close();
         }
-
     }
 
-
     private String buildComparisonQuery(PyroscopeContextHolder pyroscopeContext, long untilMilis) {
         StringBuilder qb = new StringBuilder();
-        pyroscopeContext.ctx.forEach((k, v) -> {
+        pyroscopeContext.ctx.forEachLabel((k, v) -> {
             if (k.equals(LABEL_PROFILE_ID)) {
                 return;
             }
@@ -176,10 +186,10 @@ private static String urlEncode(String query) {
 
     private static class PyroscopeContextHolder {
         final String profileId;
-        final ScopedContext ctx;
+        final ProfilerScopedContext ctx;
         final long startTimeMillis;
 
-        PyroscopeContextHolder(String profileId, ScopedContext ctx, long startTimeMillis) {
+        PyroscopeContextHolder(String profileId, ProfilerScopedContext ctx, long startTimeMillis) {
             this.profileId = profileId;
             this.ctx = ctx;
             this.startTimeMillis = startTimeMillis;