Fix agentextension publishing and LocalRootSpan bug in agentextension (#218)

Author: JonasKunz

agentextension/build.gradle.kts

@@ -14,6 +14,11 @@ dependencies {
   shadowDependencies(project(":custom"))
 }
 
+publishingConventions {
+  artifactTasks.add(tasks.shadowJar)
+  artifactTasks.add(tasks.javadocJar)
+  artifactTasks.add(tasks.sourcesJar)
+}
 
 tasks {
 

common/src/main/java/co/elastic/otel/common/LocalRootSpan.java

@@ -79,6 +79,19 @@ public static void onSpanStart(ReadableSpan startedSpan, Context parentContext)
     }
   }
 
+  /** See {@link LocalRootSpan#getFor(Span)}. */
+  @Nullable
+  public static ReadableSpan getFor(ReadableSpan span) {
+    Object rootSpanVal = localRoot.get(span);
+    if (rootSpanVal == LOCAL_ROOT_MARKER) {
+      return span;
+    }
+    if (rootSpanVal == INFERRED_SPAN_UNKNOWN_ROOT_MARKER) {
+      return null;
+    }
+    return (ReadableSpan) rootSpanVal;
+  }
+
   /**
    * If the provided span is a local root span, itself is returned. Otherwise, returns the
    * (transitive) parent which is a local root.
@@ -93,31 +106,14 @@ public static void onSpanStart(ReadableSpan startedSpan, Context parentContext)
    * </ul>
    */
   @Nullable
-  public static Span getFor(Span span) {
-    return resolveRoot(span, localRoot.get(span));
+  public static ReadableSpan getFor(Span span) {
+    return getFor((ReadableSpan) span);
   }
 
   /** See {@link LocalRootSpan#getFor(Span)}. */
   @Nullable
-  public static Span getFor(ReadableSpan span) {
-    return resolveRoot(span, localRoot.get(span));
-  }
-
-  /** See {@link LocalRootSpan#getFor(Span)}. */
-  @Nullable
-  public static Span getFor(ReadWriteSpan span) {
-    return resolveRoot(span, localRoot.get(span));
-  }
-
-  @Nullable
-  private static Span resolveRoot(Object span, Object rootSpanVal) {
-    if (rootSpanVal == LOCAL_ROOT_MARKER) {
-      return (Span) span;
-    }
-    if (rootSpanVal == INFERRED_SPAN_UNKNOWN_ROOT_MARKER) {
-      return null;
-    }
-    return (Span) rootSpanVal;
+  public static ReadableSpan getFor(ReadWriteSpan span) {
+    return getFor((ReadableSpan) span);
   }
 
   private LocalRootSpan() {}

common/src/main/java/co/elastic/otel/common/SpanValueStorageProvider.java

@@ -21,11 +21,27 @@
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.sdk.trace.FieldBackedSpanValueStorageProvider;
+import java.util.logging.Level;
+import java.util.logging.Logger;
 import javax.annotation.Nullable;
 
 public interface SpanValueStorageProvider {
 
+  Logger logger = Logger.getLogger(SpanValueStorageProvider.class.getName());
+
   static SpanValueStorageProvider get() {
+    try {
+      Class<?> sdkSpan = Class.forName("io.opentelemetry.sdk.trace.SdkSpan");
+      if (sdkSpan.getClassLoader() != SpanValueStorage.class.getClassLoader()) {
+        // If we are running in a different classloader, this means we aren't running in our distro
+        logger.log(
+            Level.FINE,
+            "Using map-backed storage for SpanValues because SdkSpan lives in a different classloader and therefore is inaccessible");
+        return MapBacked.getInstance();
+      }
+    } catch (ClassNotFoundException e) {
+      throw new IllegalStateException("Expected SdkSpan class to exist", e);
+    }
     return FieldBackedSpanValueStorageProvider.INSTANCE != null
         ? FieldBackedSpanValueStorageProvider.INSTANCE
         : MapBacked.getInstance();

universal-profiling-integration/src/main/java/co/elastic/otel/ProfilerSharedMemoryWriter.java

@@ -23,6 +23,7 @@
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanContext;
 import io.opentelemetry.sdk.resources.Resource;
+import io.opentelemetry.sdk.trace.ReadableSpan;
 import io.opentelemetry.semconv.ResourceAttributes;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -99,7 +100,7 @@ static void updateThreadCorrelationStorage(Span newSpan) {
 
       SpanContext spanCtx = newSpan.getSpanContext();
       if (spanCtx.isValid() && !spanCtx.isRemote()) {
-        Span localRoot = LocalRootSpan.getFor(newSpan);
+        ReadableSpan localRoot = LocalRootSpan.getFor(newSpan);
         if (localRoot != null) {
           String localRootSpanId = localRoot.getSpanContext().getSpanId();
           tls.put(TLS_TRACE_PRESENT_OFFSET, (byte) 1);