Set up virtual field transforms before otel sdk is initialized

Author: laurit

instrumentation/executors/javaagent/build.gradle.kts

@@ -15,6 +15,7 @@ dependencies {
   testImplementation(project(":instrumentation:executors:testing"))
   testImplementation("org.scala-lang:scala-library:2.11.12")
   testCompileOnly(project(":instrumentation:executors:bootstrap"))
+  testCompileOnly(project(":javaagent-bootstrap"))
 }
 
 testing {
@@ -30,6 +31,7 @@ testing {
       dependencies {
         implementation(project(":instrumentation:executors:testing"))
         compileOnly(project(":instrumentation:executors:bootstrap"))
+        compileOnly(project(":javaagent-bootstrap"))
       }
 
       targets {

instrumentation/executors/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/executors/ExecutorsInstrumentationModule.java

@@ -10,10 +10,12 @@
 import com.google.auto.service.AutoService;
 import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
 import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.EarlyInstrumentationModule;
 import java.util.List;
 
-@AutoService(InstrumentationModule.class)
-public class ExecutorsInstrumentationModule extends InstrumentationModule {
+@AutoService({InstrumentationModule.class, EarlyInstrumentationModule.class})
+public class ExecutorsInstrumentationModule extends InstrumentationModule
+    implements EarlyInstrumentationModule {
 
   public ExecutorsInstrumentationModule() {
     super("executors");

instrumentation/executors/javaagent/src/test/java/io/opentelemetry/javaagent/instrumentation/executors/ThreadPoolExecutorTest.java

@@ -10,7 +10,9 @@
 
 import io.opentelemetry.api.baggage.Baggage;
 import io.opentelemetry.context.Scope;
+import io.opentelemetry.javaagent.bootstrap.VirtualFieldInstalledMarker;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.FutureTask;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -19,6 +21,11 @@
 
 class ThreadPoolExecutorTest {
 
+  @Test
+  void virtualFieldsAdded() {
+    assertThat(VirtualFieldInstalledMarker.class).isAssignableFrom(FutureTask.class);
+  }
+
   @Test
   void shouldPassOriginalRunnableToBeforeAfterMethods() throws InterruptedException {
     CountDownLatch latch = new CountDownLatch(1);

instrumentation/internal/internal-reflection/javaagent/src/main/java/io/opentelemetry/javaagent/instrumentation/internal/reflection/ReflectionHelper.java

@@ -11,6 +11,8 @@
 import java.lang.reflect.Field;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 
 public final class ReflectionHelper {
@@ -61,21 +63,23 @@ public static Class<?>[] filterInterfaces(Class<?>[] interfaces, Class<?> contai
       return interfaces;
     }
     List<Class<?>> result = new ArrayList<>(interfaces.length);
+    Collection<String> virtualFieldClassNames = new HashSet<>();
     boolean hasVirtualFieldMarker = false;
     for (Class<?> interfaceClass : interfaces) {
       // filter out virtual field marker and accessor interfaces
-      if (interfaceClass == VirtualFieldInstalledMarker.class
-          || (VirtualFieldAccessorMarker.class.isAssignableFrom(interfaceClass)
-              && interfaceClass.isSynthetic()
-              && interfaceClass.getName().contains("VirtualFieldAccessor$"))) {
-        hasVirtualFieldMarker = true;
+      if (interfaceClass == VirtualFieldInstalledMarker.class) {
+        continue;
+      } else if (VirtualFieldAccessorMarker.class.isAssignableFrom(interfaceClass)
+          && interfaceClass.isSynthetic()
+          && interfaceClass.getName().contains("VirtualFieldAccessor$")) {
+        virtualFieldClassNames.add(interfaceClass.getName());
         continue;
       }
       result.add(interfaceClass);
     }
 
-    if (hasVirtualFieldMarker) {
-      VirtualFieldDetector.markVirtualFieldsPresent(containingClass);
+    if (!virtualFieldClassNames.isEmpty()) {
+      VirtualFieldDetector.markVirtualFields(containingClass, virtualFieldClassNames);
     }
 
     return result.toArray(new Class<?>[0]);

javaagent-bootstrap/src/main/java/io/opentelemetry/javaagent/bootstrap/VirtualFieldDetector.java

@@ -6,36 +6,46 @@
 package io.opentelemetry.javaagent.bootstrap;
 
 import io.opentelemetry.instrumentation.api.internal.cache.Cache;
-import java.util.Arrays;
+import java.util.Collection;
 
 /** Helper class for detecting whether given class has virtual fields. */
 public final class VirtualFieldDetector {
 
-  private static final Cache<Class<?>, Boolean> classesWithVirtualFields = Cache.weak();
+  // class to virtual field interface dot class names (see
+  // GeneratedVirtualFieldNames.getFieldAccessorInterfaceName)
+  private static final Cache<Class<?>, Collection<String>> classesWithVirtualFields = Cache.weak();
 
   private VirtualFieldDetector() {}
 
   /**
-   * Detect whether given class has virtual fields. This method looks for virtual fields only from
-   * the specified class not its super classes.
+   * Detect whether given class has given virtual field. This method looks for virtual fields only
+   * from the specified class not its super classes.
    *
    * @param clazz a class
-   * @return true if given class has virtual fields
+   * @param virtualFieldInterfaceClassName virtual field interface class dot name
+   * @return true if given class has the specified virtual field
    */
-  public static boolean hasVirtualFields(Class<?> clazz) {
+  public static boolean hasVirtualField(Class<?> clazz, String virtualFieldInterfaceClassName) {
+    if (!VirtualFieldInstalledMarker.class.isAssignableFrom(clazz)) {
+      return false;
+    }
     // clazz.getInterfaces() needs to be called before reading from classesWithVirtualFields
     // as the call to clazz.getInterfaces() triggers adding clazz to that map via instrumentation
-    // calling VirtualFieldDetector#markVirtualFieldsPresent() from Class#getInterfaces()
+    // calling VirtualFieldDetector#markVirtualFields() from Class#getInterfaces()
     Class<?>[] interfaces = clazz.getInterfaces();
     // to avoid breaking in case internal-reflection instrumentation is disabled check whether
     // interfaces array contains virtual field marker interface
-    if (Arrays.asList(interfaces).contains(VirtualFieldInstalledMarker.class)) {
-      return true;
+    for (Class<?> interfaceClass : interfaces) {
+      if (virtualFieldInterfaceClassName.equals(interfaceClass.getName())) {
+        return true;
+      }
     }
-    return classesWithVirtualFields.get(clazz) != null;
+
+    Collection<String> virtualFields = classesWithVirtualFields.get(clazz);
+    return virtualFields != null && virtualFields.contains(virtualFieldInterfaceClassName);
   }
 
-  public static void markVirtualFieldsPresent(Class<?> clazz) {
-    classesWithVirtualFields.put(clazz, Boolean.TRUE);
+  public static void markVirtualFields(Class<?> clazz, Collection<String> virtualFieldClassName) {
+    classesWithVirtualFields.put(clazz, virtualFieldClassName);
   }
 }

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

@@ -33,14 +33,6 @@
  */
 public abstract class InstrumentationModule implements Ordered {
   private static final Logger logger = Logger.getLogger(InstrumentationModule.class.getName());
-  private static final boolean indyEnabled;
-
-  static {
-    indyEnabled = ExperimentalConfig.get().indyEnabled();
-    if (indyEnabled) {
-      logger.info("Enabled indy for instrumentation modules");
-    }
-  }
 
   private final Set<String> instrumentationNames;
 
@@ -125,7 +117,7 @@ public boolean isHelperClass(String className) {
    * techniques. The non-inlining of advice will be enforced by muzzle (TODO)
    */
   public boolean isIndyModule() {
-    return indyEnabled;
+    return IndyConfigurationHolder.indyEnabled;
   }
 
   /** Register resource names to inject into the user's class loader. */
@@ -163,4 +155,16 @@ public ElementMatcher.Junction<ClassLoader> classLoaderMatcher() {
   public List<String> getAdditionalHelperClassNames() {
     return Collections.emptyList();
   }
+
+  // InstrumentationModule is loaded before ExperimentalConfig is initialized
+  private static class IndyConfigurationHolder {
+    private static final boolean indyEnabled;
+
+    static {
+      indyEnabled = ExperimentalConfig.get().indyEnabled();
+      if (indyEnabled) {
+        logger.info("Enabled indy for instrumentation modules");
+      }
+    }
+  }
 }

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

@@ -0,0 +1,23 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package io.opentelemetry.javaagent.extension.instrumentation.internal;
+
+import io.opentelemetry.javaagent.extension.instrumentation.InstrumentationModule;
+import io.opentelemetry.sdk.autoconfigure.spi.Ordered;
+
+/**
+ * Marker interface for instrumentation modules whose virtual fields should be set up before
+ * OpenTelemetry SDK is initialized.
+ *
+ * <p>This class is internal and is hence not for public use. Its APIs are unstable and can change
+ * at any time.
+ */
+public interface EarlyInstrumentationModule extends Ordered {
+
+  default InstrumentationModule getInstrumentationModule() {
+    return (InstrumentationModule) this;
+  }
+}

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

@@ -13,6 +13,7 @@
 import static java.util.logging.Level.FINE;
 import static java.util.logging.Level.SEVERE;
 import static net.bytebuddy.matcher.ElementMatchers.any;
+import static net.bytebuddy.matcher.ElementMatchers.none;
 
 import io.opentelemetry.context.Context;
 import io.opentelemetry.context.ContextStorage;
@@ -30,11 +31,15 @@
 import io.opentelemetry.javaagent.bootstrap.internal.ConfiguredResourceAttributesHolder;
 import io.opentelemetry.javaagent.extension.AgentListener;
 import io.opentelemetry.javaagent.extension.ignore.IgnoredTypesConfigurer;
+import io.opentelemetry.javaagent.extension.instrumentation.internal.EarlyInstrumentationModule;
 import io.opentelemetry.javaagent.tooling.asyncannotationsupport.WeakRefAsyncOperationEndStrategies;
 import io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesBuilderImpl;
 import io.opentelemetry.javaagent.tooling.bootstrap.BootstrapPackagesConfigurer;
 import io.opentelemetry.javaagent.tooling.config.ConfigPropertiesBridge;
 import io.opentelemetry.javaagent.tooling.config.EarlyInitAgentConfig;
+import io.opentelemetry.javaagent.tooling.field.FieldBackedImplementationConfiguration;
+import io.opentelemetry.javaagent.tooling.field.VirtualFieldImplementationInstaller;
+import io.opentelemetry.javaagent.tooling.field.VirtualFieldImplementationInstallerFactory;
 import io.opentelemetry.javaagent.tooling.ignore.IgnoredClassLoadersMatcher;
 import io.opentelemetry.javaagent.tooling.ignore.IgnoredTypesBuilderImpl;
 import io.opentelemetry.javaagent.tooling.ignore.IgnoredTypesMatcher;
@@ -50,6 +55,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.logging.LogManager;
 import java.util.logging.Logger;
 import java.util.stream.Stream;
@@ -85,6 +91,8 @@ public class AgentInstaller {
 
   private static final Map<String, List<Runnable>> CLASS_LOAD_CALLBACKS = new HashMap<>();
 
+  private static volatile boolean instrumentationInstalled;
+
   public static void installBytebuddyAgent(
       Instrumentation inst, ClassLoader extensionClassLoader, EarlyInitAgentConfig earlyConfig) {
     addByteBuddyRawSetting();
@@ -99,7 +107,7 @@ public static void installBytebuddyAgent(
     if (earlyConfig.getBoolean(JAVAAGENT_ENABLED_CONFIG, true)) {
       setupUnsafe(inst);
       List<AgentListener> agentListeners = loadOrdered(AgentListener.class, extensionClassLoader);
-      installBytebuddyAgent(inst, extensionClassLoader, agentListeners);
+      installBytebuddyAgent(inst, extensionClassLoader, agentListeners, earlyConfig);
     } else {
       logger.fine("Tracing is disabled, not installing instrumentations.");
     }
@@ -108,31 +116,13 @@ public static void installBytebuddyAgent(
   private static void installBytebuddyAgent(
       Instrumentation inst,
       ClassLoader extensionClassLoader,
-      Iterable<AgentListener> agentListeners) {
+      Iterable<AgentListener> agentListeners,
+      EarlyInitAgentConfig earlyConfig) {
 
     WeakRefAsyncOperationEndStrategies.initialize();
-
     EmbeddedInstrumentationProperties.setPropertiesLoader(extensionClassLoader);
-
     setDefineClassHandler();
-
-    // If noop OpenTelemetry is enabled, autoConfiguredSdk will be null and AgentListeners are not
-    // called
-    AutoConfiguredOpenTelemetrySdk autoConfiguredSdk =
-        installOpenTelemetrySdk(extensionClassLoader);
-
-    ConfigProperties sdkConfig = AutoConfigureUtil.getConfig(autoConfiguredSdk);
-    AgentInstrumentationConfig.internalInitializeConfig(new ConfigPropertiesBridge(sdkConfig));
-    copyNecessaryConfigToSystemProperties(sdkConfig);
-
-    setBootstrapPackages(sdkConfig, extensionClassLoader);
-    ConfiguredResourceAttributesHolder.initialize(
-        SdkAutoconfigureAccess.getResourceAttributes(autoConfiguredSdk));
-
-    for (BeforeAgentListener agentListener :
-        loadOrdered(BeforeAgentListener.class, extensionClassLoader)) {
-      agentListener.beforeAgent(autoConfiguredSdk);
-    }
+    FieldBackedImplementationConfiguration.configure(earlyConfig);
 
     AgentBuilder agentBuilder =
         new AgentBuilder.Default(
@@ -154,9 +144,6 @@ private static void installBytebuddyAgent(
     if (JavaModule.isSupported()) {
       agentBuilder = agentBuilder.with(new ExposeAgentBootstrapListener(inst));
     }
-
-    agentBuilder = configureIgnoredTypes(sdkConfig, extensionClassLoader, agentBuilder);
-
     if (logger.isLoggable(FINE)) {
       agentBuilder =
           agentBuilder
@@ -166,6 +153,28 @@ private static void installBytebuddyAgent(
               .with(new TransformLoggingListener());
     }
 
+    installEarlyInstrumentation(agentBuilder, inst);
+
+    // If noop OpenTelemetry is enabled, autoConfiguredSdk will be null and AgentListeners are not
+    // called
+    AutoConfiguredOpenTelemetrySdk autoConfiguredSdk =
+        installOpenTelemetrySdk(extensionClassLoader);
+
+    ConfigProperties sdkConfig = AutoConfigureUtil.getConfig(autoConfiguredSdk);
+    AgentInstrumentationConfig.internalInitializeConfig(new ConfigPropertiesBridge(sdkConfig));
+    copyNecessaryConfigToSystemProperties(sdkConfig);
+
+    setBootstrapPackages(sdkConfig, extensionClassLoader);
+    ConfiguredResourceAttributesHolder.initialize(
+        SdkAutoconfigureAccess.getResourceAttributes(autoConfiguredSdk));
+
+    for (BeforeAgentListener agentListener :
+        loadOrdered(BeforeAgentListener.class, extensionClassLoader)) {
+      agentListener.beforeAgent(autoConfiguredSdk);
+    }
+
+    agentBuilder = configureIgnoredTypes(sdkConfig, extensionClassLoader, agentBuilder);
+
     int numberOfLoadedExtensions = 0;
     for (AgentExtension agentExtension : loadOrdered(AgentExtension.class, extensionClassLoader)) {
       if (logger.isLoggable(FINE)) {
@@ -192,13 +201,44 @@ private static void installBytebuddyAgent(
 
     agentBuilder = AgentBuilderUtil.optimize(agentBuilder);
     ResettableClassFileTransformer resettableClassFileTransformer = agentBuilder.installOn(inst);
+    instrumentationInstalled = true;
     ClassFileTransformerHolder.setClassFileTransformer(resettableClassFileTransformer);
 
     addHttpServerResponseCustomizers(extensionClassLoader);
 
     runAfterAgentListeners(agentListeners, autoConfiguredSdk);
   }
 
+  private static void installEarlyInstrumentation(
+      AgentBuilder agentBuilder, Instrumentation instrumentation) {
+    // We are only going to install the virtual fields here. Installing virtual field changes class
+    // structure and can not be applied to already loaded classes.
+    agentBuilder = agentBuilder.with(AgentBuilder.RedefinitionStrategy.DISABLED);
+
+    AgentBuilder.Identified.Extendable extendableAgentBuilder =
+        agentBuilder
+            .ignore(
+                target -> instrumentationInstalled, // turn off after instrumentation is installed
+                Objects::nonNull)
+            .type(none())
+            .transform(
+                (builder, typeDescription, classLoader, module, protectionDomain) -> builder);
+
+    VirtualFieldImplementationInstallerFactory virtualFieldInstallerFactory =
+        new VirtualFieldImplementationInstallerFactory();
+    for (EarlyInstrumentationModule earlyInstrumentationModule :
+        loadOrdered(EarlyInstrumentationModule.class, Utils.getExtensionsClassLoader())) {
+
+      VirtualFieldImplementationInstaller contextProvider =
+          virtualFieldInstallerFactory.create(
+              earlyInstrumentationModule.getInstrumentationModule());
+      extendableAgentBuilder = contextProvider.injectFields(extendableAgentBuilder);
+    }
+
+    agentBuilder = AgentBuilderUtil.optimize(extendableAgentBuilder);
+    agentBuilder.installOn(instrumentation);
+  }
+
   private static void copyNecessaryConfigToSystemProperties(ConfigProperties config) {
     for (String property : asList("otel.instrumentation.experimental.span-suppression-strategy")) {
       String value = config.getString(property);