Allow method instrumentation trace methods in boot loader (#12454)

Author: laurit

instrumentation/methods/javaagent/build.gradle.kts

@@ -15,6 +15,6 @@ dependencies {
 
 tasks.withType<Test>().configureEach {
   jvmArgs(
-    "-Dotel.instrumentation.methods.include=io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCallable[call];io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCompletableFuture[getResult]"
+    "-Dotel.instrumentation.methods.include=io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCallable[call];io.opentelemetry.javaagent.instrumentation.methods.MethodTest\$ConfigTracedCompletableFuture[getResult];javax.naming.directory.InitialDirContext[search]"
   )
 }

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentation.java

@@ -8,6 +8,7 @@
 import static io.opentelemetry.javaagent.bootstrap.Java8BytecodeBridge.currentContext;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasClassesNamed;
 import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
+import static io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons.getBootstrapLoader;
 import static io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons.instrumenter;
 import static net.bytebuddy.matcher.ElementMatchers.named;
 import static net.bytebuddy.matcher.ElementMatchers.namedOneOf;
@@ -36,7 +37,15 @@ public MethodInstrumentation(String className, Set<String> methodNames) {
 
   @Override
   public ElementMatcher<ClassLoader> classLoaderOptimization() {
-    return hasClassesNamed(className);
+    ElementMatcher<ClassLoader> delegate = hasClassesNamed(className);
+    return target -> {
+      // hasClassesNamed does not support null class loader, so we provide a custom loader that
+      // can be used to look up resources in bootstrap loader
+      if (target == null) {
+        target = getBootstrapLoader();
+      }
+      return delegate.matches(target);
+    };
   }
 
   @Override

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodInstrumentationModule.java

@@ -6,13 +6,13 @@
 package io.opentelemetry.javaagent.instrumentation.methods;
 
 import static java.util.Collections.emptyList;
-import static java.util.Collections.singletonList;
 
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.bootstrap.internal.AgentInstrumentationConfig;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
 import io.opentelemetry.javaagent.tooling.config.MethodsConfigurationParser;
+import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -45,7 +45,9 @@ public MethodInstrumentationModule() {
   public List<String> getAdditionalHelperClassNames() {
     return typeInstrumentations.isEmpty()
         ? emptyList()
-        : singletonList("io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons");
+        : Arrays.asList(
+            "io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons",
+            "io.opentelemetry.javaagent.instrumentation.methods.MethodSingletons$BootstrapLoader");
   }
 
   @Override

instrumentation/methods/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/methods/MethodSingletons.java

@@ -17,6 +17,7 @@ public final class MethodSingletons {
   private static final String INSTRUMENTATION_NAME = "io.opentelemetry.methods";
 
   private static final Instrumenter<ClassAndMethod, Void> INSTRUMENTER;
+  private static final ClassLoader bootstrapLoader = new BootstrapLoader();
 
   static {
     CodeAttributesGetter<ClassAndMethod> codeAttributesGetter =
@@ -35,5 +36,19 @@ public static Instrumenter<ClassAndMethod, Void> instrumenter() {
     return INSTRUMENTER;
   }
 
+  public static ClassLoader getBootstrapLoader() {
+    return bootstrapLoader;
+  }
+
   private MethodSingletons() {}
+
+  private static class BootstrapLoader extends ClassLoader {
+    static {
+      ClassLoader.registerAsParallelCapable();
+    }
+
+    BootstrapLoader() {
+      super(null);
+    }
+  }
 }

instrumentation/methods/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/methods/MethodTest.java

@@ -9,6 +9,7 @@
 import static io.opentelemetry.sdk.testing.assertj.OpenTelemetryAssertions.equalTo;
 import static io.opentelemetry.semconv.incubating.CodeIncubatingAttributes.CODE_FUNCTION;
 import static io.opentelemetry.semconv.incubating.CodeIncubatingAttributes.CODE_NAMESPACE;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.api.trace.SpanKind;
@@ -18,6 +19,10 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import javax.naming.NoInitialContextException;
+import javax.naming.directory.InitialDirContext;
+import javax.naming.ldap.InitialLdapContext;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 
@@ -40,6 +45,33 @@ void methodTraced() {
                             equalTo(CODE_FUNCTION, "call"))));
   }
 
+  @Test
+  void bootLoaderMethodTraced() throws Exception {
+    InitialLdapContext context = new InitialLdapContext();
+    AtomicReference<Throwable> throwableReference = new AtomicReference<>();
+    assertThatThrownBy(
+            () -> {
+              try {
+                context.search("foo", null);
+              } catch (Throwable throwable) {
+                throwableReference.set(throwable);
+                throw throwable;
+              }
+            })
+        .isInstanceOf(NoInitialContextException.class);
+
+    testing.waitAndAssertTraces(
+        trace ->
+            trace.hasSpansSatisfyingExactly(
+                span ->
+                    span.hasName("InitialDirContext.search")
+                        .hasKind(SpanKind.INTERNAL)
+                        .hasException(throwableReference.get())
+                        .hasAttributesSatisfyingExactly(
+                            equalTo(CODE_NAMESPACE, InitialDirContext.class.getName()),
+                            equalTo(CODE_FUNCTION, "search"))));
+  }
+
   static class ConfigTracedCallable implements Callable<String> {
 
     @Override