Address review comments: use requireNonNull consistently, fix race condition in MuzzleMatcher, fix SafeServiceLoader.next()

Signed-off-by: Gregor Zeitlinger <gregor.zeitlinger@grafana.com>

Author: zeitlinger

javaagent-extension-api/src/main/java/io/opentelemetry/javaagent/extension/instrumentation/InstrumentationModule.java

@@ -19,7 +19,6 @@
 import java.util.List;
 import java.util.Set;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
 import net.bytebuddy.matcher.ElementMatcher;
 
 /**
@@ -108,7 +107,7 @@ public boolean defaultEnabled() {
    * @deprecated Use {@link #defaultEnabled()} instead.
    */
   @Deprecated // will be removed in 3.0.0
-  public boolean defaultEnabled(@Nullable ConfigProperties config) {
+  public boolean defaultEnabled(ConfigProperties config) {
     return defaultEnabled();
   }
 

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentExtension.java

@@ -8,7 +8,6 @@
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import io.opentelemetry.sdk.autoconfigure.spi.Ordered;
-import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
 
 /**
@@ -27,7 +26,7 @@ public interface AgentExtension extends Ordered {
    * @return The customized agent. Note that this method MUST return a non-null {@link AgentBuilder}
    *     instance that contains all customizations defined in this extension.
    */
-  AgentBuilder extend(AgentBuilder agentBuilder, @Nullable ConfigProperties config);
+  AgentBuilder extend(AgentBuilder agentBuilder, ConfigProperties config);
 
   /**
    * Returns the name of the extension. It does not have to be unique, but it should be

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/AgentInstaller.java

@@ -9,6 +9,7 @@
 import static io.opentelemetry.javaagent.tooling.SafeServiceLoader.load;
 import static io.opentelemetry.javaagent.tooling.SafeServiceLoader.loadOrdered;
 import static io.opentelemetry.javaagent.tooling.Utils.getResourceName;
+import static java.util.Objects.requireNonNull;
 import static java.util.logging.Level.FINE;
 import static java.util.logging.Level.SEVERE;
 import static net.bytebuddy.matcher.ElementMatchers.any;
@@ -161,7 +162,9 @@ private static void installBytebuddyAgent(
     AutoConfiguredOpenTelemetrySdk autoConfiguredSdk =
         installOpenTelemetrySdk(extensionClassLoader);
 
-    ConfigProperties sdkConfig = AutoConfigureUtil.getConfig(autoConfiguredSdk);
+    ConfigProperties sdkConfig =
+        requireNonNull(
+            AutoConfigureUtil.getConfig(autoConfiguredSdk), "SDK config must not be null");
 
     setBootstrapPackages(sdkConfig, extensionClassLoader);
     ConfiguredResourceAttributesHolder.initialize(
@@ -261,24 +264,24 @@ private static void installEarlyInstrumentation(
 
     VirtualFieldImplementationInstallerFactory virtualFieldInstallerFactory =
         VirtualFieldImplementationInstallerFactory.getInstance();
-    ClassLoader extensionsClassLoader = Utils.getExtensionsClassLoader();
-    if (extensionsClassLoader != null) {
-      for (EarlyInstrumentationModule earlyInstrumentationModule :
-          loadOrdered(EarlyInstrumentationModule.class, extensionsClassLoader)) {
-
-        VirtualFieldImplementationInstaller contextProvider =
-            virtualFieldInstallerFactory.create(
-                earlyInstrumentationModule.getInstrumentationModule());
-        extendableAgentBuilder = contextProvider.injectFields(extendableAgentBuilder);
-      }
+    ClassLoader extensionsClassLoader =
+        requireNonNull(
+            Utils.getExtensionsClassLoader(), "Extensions class loader must not be null");
+    for (EarlyInstrumentationModule earlyInstrumentationModule :
+        loadOrdered(EarlyInstrumentationModule.class, extensionsClassLoader)) {
+
+      VirtualFieldImplementationInstaller contextProvider =
+          virtualFieldInstallerFactory.create(
+              earlyInstrumentationModule.getInstrumentationModule());
+      extendableAgentBuilder = contextProvider.injectFields(extendableAgentBuilder);
     }
 
     agentBuilder = AgentBuilderUtil.optimize(extendableAgentBuilder);
     agentBuilder.installOn(instrumentation);
   }
 
   private static void setBootstrapPackages(
-      @Nullable ConfigProperties config, ClassLoader extensionClassLoader) {
+      ConfigProperties config, ClassLoader extensionClassLoader) {
     BootstrapPackagesBuilderImpl builder = new BootstrapPackagesBuilderImpl();
     for (BootstrapPackagesConfigurer configurer :
         load(BootstrapPackagesConfigurer.class, extensionClassLoader)) {
@@ -294,13 +297,11 @@ private static void setDefineClassHandler() {
   // Need to call deprecated API for backward compatibility with extensions that haven't migrated
   @SuppressWarnings("deprecation")
   private static AgentBuilder configureIgnoredTypes(
-      @Nullable ConfigProperties config,
-      ClassLoader extensionClassLoader,
-      AgentBuilder agentBuilder) {
+      ConfigProperties config, ClassLoader extensionClassLoader, AgentBuilder agentBuilder) {
     IgnoredTypesBuilderImpl builder = new IgnoredTypesBuilderImpl();
     for (IgnoredTypesConfigurer configurer :
         loadOrdered(IgnoredTypesConfigurer.class, extensionClassLoader)) {
-      configurer.configure(builder, config != null ? config : EmptyConfigProperties.INSTANCE);
+      configurer.configure(builder, config);
     }
 
     Trie<Boolean> ignoredTasksTrie = builder.buildIgnoredTasksTrie();

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/DefineClassHandler.java

@@ -86,9 +86,7 @@ private static void addSuperNames(Set<String> superNames, Class<?> clazz) {
 
   @Override
   public void afterDefineClass(DefineClassContext context) {
-    if (context != null) {
-      context.exit();
-    }
+    context.exit();
   }
 
   /**

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/SafeServiceLoader.java

@@ -12,9 +12,9 @@
 import java.util.Comparator;
 import java.util.Iterator;
 import java.util.List;
+import java.util.NoSuchElementException;
 import java.util.ServiceLoader;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
 
 public final class SafeServiceLoader {
 
@@ -35,9 +35,10 @@ public static <T> List<T> load(Class<T> serviceClass, ClassLoader classLoader) {
     List<T> result = new ArrayList<>();
     ServiceLoader<T> services = ServiceLoader.load(serviceClass, classLoader);
     for (Iterator<T> iterator = new SafeIterator<>(services.iterator()); iterator.hasNext(); ) {
-      T service = iterator.next();
-      if (service != null) {
-        result.add(service);
+      try {
+        result.add(iterator.next());
+      } catch (NoSuchElementException ignored) {
+        // UnsupportedClassVersionError was thrown and handled by SafeIterator
       }
     }
     return result;
@@ -83,15 +84,14 @@ public boolean hasNext() {
       }
     }
 
-    @Nullable
     @Override
     public T next() {
       // jdk8 throws UnsupportedClassVersionError in next()
       try {
         return delegate.next();
       } catch (UnsupportedClassVersionError unsupportedClassVersionError) {
         handleUnsupportedClassVersionError(unsupportedClassVersionError);
-        return null;
+        throw new NoSuchElementException(unsupportedClassVersionError.getMessage());
       }
     }
   }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/bootstrap/BootstrapPackagesConfigurer.java

@@ -6,7 +6,6 @@
 package io.opentelemetry.javaagent.tooling.bootstrap;
 
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
-import javax.annotation.Nullable;
 
 /**
  * This SPI can be used to define which packages/classes belong to the bootstrap class loader: all
@@ -25,5 +24,5 @@ public interface BootstrapPackagesConfigurer {
    * Configure the passed {@code builder} and define which classes should always be loaded by the
    * bootstrap class loader.
    */
-  void configure(BootstrapPackagesBuilder builder, @Nullable ConfigProperties config);
+  void configure(BootstrapPackagesBuilder builder, ConfigProperties config);
 }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationLoader.java

@@ -17,18 +17,19 @@
 import io.opentelemetry.javaagent.tooling.Utils;
 import io.opentelemetry.sdk.autoconfigure.spi.ConfigProperties;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
 
 @AutoService(AgentExtension.class)
 public class InstrumentationLoader implements AgentExtension {
   private static final Logger logger = Logger.getLogger(InstrumentationLoader.class.getName());
 
   private final InstrumentationModuleInstaller instrumentationModuleInstaller =
-      new InstrumentationModuleInstaller(InstrumentationHolder.getInstrumentation());
+      new InstrumentationModuleInstaller(
+          requireNonNull(
+              InstrumentationHolder.getInstrumentation(), "Instrumentation must not be null"));
 
   @Override
-  public AgentBuilder extend(AgentBuilder agentBuilder, @Nullable ConfigProperties config) {
+  public AgentBuilder extend(AgentBuilder agentBuilder, ConfigProperties config) {
     int numberOfLoadedModules = 0;
     ClassLoader extensionsClassLoader =
         requireNonNull(

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/InstrumentationModuleInstaller.java

@@ -41,7 +41,6 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Function;
-import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.annotation.AnnotationSource;
 import net.bytebuddy.description.type.TypeDescription;
@@ -58,20 +57,20 @@ public final class InstrumentationModuleInstaller {
   public static final ElementMatcher.Junction<AnnotationSource> NOT_DECORATOR_MATCHER =
       not(isAnnotatedWith(named("javax.decorator.Decorator")));
 
-  @Nullable private final Instrumentation instrumentation;
+  private final Instrumentation instrumentation;
   private final VirtualFieldImplementationInstallerFactory virtualFieldInstallerFactory =
       VirtualFieldImplementationInstallerFactory.getInstance();
 
-  public InstrumentationModuleInstaller(@Nullable Instrumentation instrumentation) {
-    this.instrumentation = instrumentation;
+  public InstrumentationModuleInstaller(Instrumentation instrumentation) {
+    this.instrumentation = requireNonNull(instrumentation, "Instrumentation must not be null");
   }
 
   // Need to call deprecated API for backward compatibility with modules that haven't migrated
   @SuppressWarnings("deprecation")
   AgentBuilder install(
       InstrumentationModule instrumentationModule,
       AgentBuilder parentAgentBuilder,
-      @Nullable ConfigProperties config) {
+      ConfigProperties config) {
     if (!isInstrumentationEnabled(
         instrumentationModule.instrumentationNames(),
         instrumentationModule.defaultEnabled(config))) {
@@ -142,7 +141,7 @@ private AgentBuilder installIndyModule(
             helperGenerator,
             helperResourceBuilder.getResources(),
             instrumentationModule.getClass().getClassLoader(),
-            requireNonNull(instrumentation, "Instrumentation must not be null"));
+            instrumentation);
 
     VirtualFieldImplementationInstaller contextProvider =
         virtualFieldInstallerFactory.create(instrumentationModule);
@@ -193,14 +192,13 @@ private AgentBuilder installInjectingModule(
     ClassLoader extensionsClassLoader =
         requireNonNull(
             Utils.getExtensionsClassLoader(), "Extensions class loader must not be null");
-    Instrumentation inst = requireNonNull(instrumentation, "Instrumentation must not be null");
     AgentBuilder.Transformer helperInjector =
         new HelperInjector(
             instrumentationModule.instrumentationName(),
             helperClassNames,
             helperResourceBuilder.getResources(),
             extensionsClassLoader,
-            inst);
+            instrumentation);
     VirtualFieldImplementationInstaller contextProvider =
         virtualFieldInstallerFactory.create(instrumentationModule);
 
@@ -222,7 +220,8 @@ private AgentBuilder installInjectingModule(
                           .jpmsModulesToOpen()
                           .forEach(
                               (javaModule, packages) -> {
-                                ModuleOpener.open(inst, javaModule, classLoader, packages);
+                                ModuleOpener.open(
+                                    instrumentation, javaModule, classLoader, packages);
                               });
                       openerRun.set(true);
                     }

javaagent-tooling/src/main/java/io/opentelemetry/javaagent/tooling/instrumentation/MuzzleMatcher.java

@@ -21,10 +21,9 @@
 import io.opentelemetry.javaagent.tooling.muzzle.ReferenceMatcher;
 import java.security.ProtectionDomain;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.logging.Level;
 import java.util.logging.Logger;
-import javax.annotation.Nullable;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.utility.JavaModule;
@@ -41,9 +40,8 @@ class MuzzleMatcher implements AgentBuilder.RawMatcher {
   private final TransformSafeLogger instrumentationLogger;
   private final InstrumentationModule instrumentationModule;
   private final Level muzzleLogLevel;
-  private final AtomicBoolean initialized = new AtomicBoolean(false);
+  private final AtomicReference<ReferenceMatcher> referenceMatcher = new AtomicReference<>();
   private final Cache<ClassLoader, Boolean> matchCache = Cache.weak();
-  @Nullable private volatile ReferenceMatcher referenceMatcher;
 
   MuzzleMatcher(
       TransformSafeLogger instrumentationLogger, InstrumentationModule instrumentationModule) {
@@ -115,9 +113,11 @@ private boolean doesMatch(ClassLoader classLoader) {
   // ReferenceMatcher is lazily created to avoid unnecessarily loading the muzzle references from
   // the module during the agent setup
   private ReferenceMatcher getReferenceMatcher() {
-    if (initialized.compareAndSet(false, true)) {
-      referenceMatcher = ReferenceMatcher.of(instrumentationModule);
+    ReferenceMatcher result = referenceMatcher.get();
+    if (result == null) {
+      referenceMatcher.compareAndSet(null, ReferenceMatcher.of(instrumentationModule));
+      result = requireNonNull(referenceMatcher.get());
     }
-    return requireNonNull(referenceMatcher);
+    return result;
   }
 }

muzzle/src/main/java/io/opentelemetry/javaagent/tooling/HelperInjector.java

@@ -222,13 +222,12 @@ public DynamicType.Builder<?> transform(
 
   private void injectHelperResources(
       ClassLoader classLoader, Map<String, URL> additionalResources) {
-    if (helpersSource == null) {
-      return;
-    }
+    ClassLoader source =
+        requireNonNull(helpersSource, "helpersSource must not be null when injecting resources");
     for (HelperResource helperResource : helperResources) {
       List<URL> resources;
       try {
-        resources = Collections.list(helpersSource.getResources(helperResource.getAgentPath()));
+        resources = Collections.list(source.getResources(helperResource.getAgentPath()));
       } catch (IOException e) {
         logger.log(
             SEVERE,