native-image instrumentation (#4333)

* Ignore .gz resource dependencies when pre-processing Instrumenters as byte-buddy can't process them
* Automatically adjust native-image configuration to support dd-java-agent instrumentation
* Only apply instrumentation, don't start other services, when native-image build is happening
* Turn off weak-map background thread during native-image instrumentation

buildSrc/src/main/groovy/InstrumentPlugin.groovy

@@ -67,8 +67,9 @@ class InstrumentPlugin implements Plugin<Project> {
           compileTask.destinationDirectory = rawClassesDir.asFile
 
           byteBuddyTask.classPath.from((project.configurations.findByName('instrumentationMuzzle') ?: []) +
-            project.configurations.compileClasspath.findAll { it.name != 'previous-compilation-data.bin' } +
-            compileTask.destinationDirectory)
+            project.configurations.compileClasspath.findAll {
+              it.name != 'previous-compilation-data.bin' && !it.name.endsWith(".gz")
+            } + compileTask.destinationDirectory)
 
           byteBuddyTask.transformation {
             it.plugin = InstrumentLoader

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/Agent.java

@@ -127,6 +127,11 @@ public boolean isEnabledByDefault() {
   public static void start(final Instrumentation inst, final URL agentJarURL) {
     createAgentClassloader(agentJarURL);
 
+    if (Platform.isNativeImageBuilder()) {
+      startDatadogAgent(inst);
+      return;
+    }
+
     // Retro-compatibility for the old way to configure CI Visibility
     if ("true".equals(ddGetProperty("dd.integration.junit.enabled"))
         || "true".equals(ddGetProperty("dd.integration.testng.enabled"))) {

dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/WeakMaps.java

@@ -1,6 +1,7 @@
 package datadog.trace.agent.tooling;
 
 import com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap;
+import datadog.trace.api.Platform;
 import datadog.trace.bootstrap.WeakMap;
 import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.AgentTaskScheduler.Task;
@@ -12,12 +13,14 @@ public class WeakMaps {
 
   public static <K, V> WeakMap<K, V> newWeakMap() {
     final WeakConcurrentMap<K, V> map = new WeakConcurrentMap<>(false, true);
-    AgentTaskScheduler.INSTANCE.weakScheduleAtFixedRate(
-        MapCleaningTask.INSTANCE,
-        map,
-        CLEAN_FREQUENCY_SECONDS,
-        CLEAN_FREQUENCY_SECONDS,
-        TimeUnit.SECONDS);
+    if (!Platform.isNativeImageBuilder()) {
+      AgentTaskScheduler.INSTANCE.weakScheduleAtFixedRate(
+          MapCleaningTask.INSTANCE,
+          map,
+          CLEAN_FREQUENCY_SECONDS,
+          CLEAN_FREQUENCY_SECONDS,
+          TimeUnit.SECONDS);
+    }
     return new Adapter<>(map);
   }
 

dd-java-agent/instrumentation/graal/build.gradle

@@ -0,0 +1 @@
+apply from: "$rootDir/gradle/java.gradle"

dd-java-agent/instrumentation/graal/native-image/build.gradle

@@ -0,0 +1,41 @@
+ext {
+  minJavaVersionForTests = JavaVersion.VERSION_11
+}
+
+muzzle {
+  pass {
+    group = "org.graalvm.nativeimage"
+    module = "svm"
+    versions = "[20,)"
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+apply plugin: "idea"
+
+[compileMain_java11Java, compileTestJava].each {
+  it.sourceCompatibility = JavaVersion.VERSION_11
+  it.targetCompatibility = JavaVersion.VERSION_11
+  setJavaVersion(it, 11)
+  it.options.compilerArgs.addAll([
+    '-Xlint:all,-processing,-options,-path',
+    '--add-modules',
+    'jdk.internal.vm.ci',
+    '--add-exports',
+    'jdk.internal.vm.ci/jdk.vm.ci.meta=ALL-UNNAMED'
+  ])
+}
+
+dependencies {
+  main_java11CompileOnly group: 'org.graalvm.nativeimage', name: 'svm', version: '20.0.0'
+}
+
+forbiddenApisMain_java11 {
+  failOnMissingClasses = false
+}
+
+idea {
+  module {
+    jdkName = '11'
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/AbstractNativeImageInstrumentation.java

@@ -0,0 +1,16 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import datadog.trace.api.Platform;
+import java.util.Set;
+
+public abstract class AbstractNativeImageInstrumentation extends Instrumenter.Default {
+  public AbstractNativeImageInstrumentation() {
+    super("native-image");
+  }
+
+  @Override
+  public boolean isApplicable(Set<TargetSystem> enabledSystems) {
+    return Platform.isNativeImageBuilder();
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/AnnotationSubstitutionProcessorInstrumentation.java

@@ -0,0 +1,33 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+
+@AutoService(Instrumenter.class)
+public final class AnnotationSubstitutionProcessorInstrumentation
+    extends AbstractNativeImageInstrumentation implements Instrumenter.ForSingleType {
+
+  @Override
+  public String instrumentedType() {
+    return "com.oracle.svm.hosted.substitute.AnnotationSubstitutionProcessor";
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod()
+            .and(named("lookup"))
+            .and(takesArgument(0, named("jdk.vm.ci.meta.ResolvedJavaField"))),
+        packageName + ".DeleteFieldAdvice");
+  }
+
+  @Override
+  public String[] muzzleIgnoredClassNames() {
+    // ignore JVMCI classes which are part of GraalVM but aren't available in public repositories
+    return new String[] {"jdk.vm.ci.meta.ResolvedJavaType", "jdk.vm.ci.meta.ResolvedJavaField"};
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java/datadog/trace/instrumentation/graal/nativeimage/ClassInitializationFeatureInstrumentation.java

@@ -0,0 +1,30 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+
+import com.google.auto.service.AutoService;
+import datadog.trace.agent.tooling.Instrumenter;
+
+@AutoService(Instrumenter.class)
+public final class ClassInitializationFeatureInstrumentation
+    extends AbstractNativeImageInstrumentation implements Instrumenter.ForSingleType {
+
+  @Override
+  public String instrumentedType() {
+    return "com.oracle.svm.hosted.classinitialization.ClassInitializationFeature";
+  }
+
+  @Override
+  public void adviceTransformations(AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        isMethod()
+            .and(named("processClassInitializationOptions"))
+            .and(
+                takesArgument(
+                    0,
+                    named("com.oracle.svm.hosted.classinitialization.ClassInitializationSupport"))),
+        packageName + ".ClassInitializationAdvice");
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java11/datadog/trace/instrumentation/graal/nativeimage/ClassInitializationAdvice.java

@@ -0,0 +1,47 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import net.bytebuddy.asm.Advice;
+import org.graalvm.nativeimage.impl.RuntimeClassInitializationSupport;
+
+public class ClassInitializationAdvice {
+  @Advice.OnMethodEnter
+  public static void onEnter(
+      @Advice.Argument(0) RuntimeClassInitializationSupport classInitializationSupport) {
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.api.env.CapturedEnvironment", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.api.ConfigDefaults", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.api.Functions", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.api.InstrumenterConfig", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.api.Platform", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.api.GlobalTracer", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.config.provider.ConfigConverter", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.config.provider.ConfigProvider", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.bootstrap.Agent", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.bootstrap.BootstrapProxy", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.CallDepthThreadLocalMap", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.DatadogClassLoader", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.FieldBackedContextStores", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.instrumentation.java.concurrent.ConcurrentState", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.bootstrap.instrumentation.java.concurrent.TPEHelper", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.logging.LoggingSettingsDescription", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "datadog.trace.logging.simplelogger.SLCompatFactory", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.trace.util.CollectionUtils", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.slf4j.impl.StaticLoggerBinder", "");
+    classInitializationSupport.initializeAtBuildTime("datadog.slf4j.LoggerFactory", "");
+    classInitializationSupport.initializeAtBuildTime(
+        "com.blogspot.mydailyjava.weaklockfree.WeakConcurrentMap", "");
+    classInitializationSupport.initializeAtBuildTime("net.bytebuddy.", "");
+    classInitializationSupport.initializeAtBuildTime("com.sun.proxy.", "");
+  }
+}

dd-java-agent/instrumentation/graal/native-image/src/main/java11/datadog/trace/instrumentation/graal/nativeimage/DeleteFieldAdvice.java

@@ -0,0 +1,19 @@
+package datadog.trace.instrumentation.graal.nativeimage;
+
+import com.oracle.svm.core.annotate.Delete;
+import com.oracle.svm.hosted.substitute.AnnotatedField;
+import jdk.vm.ci.meta.ResolvedJavaField;
+import net.bytebuddy.asm.Advice;
+
+public class DeleteFieldAdvice {
+  @Advice.OnMethodExit
+  public static void onExit(
+      @Advice.Argument(0) ResolvedJavaField field,
+      @Advice.Return(readOnly = false) ResolvedJavaField result,
+      @Advice.FieldValue("SUBSTITUTION_DELETE") Delete SUBSTITUTION_DELETE) {
+    if ("datadog.trace.bootstrap.DatadogClassLoader"
+        .equals(field.getDeclaringClass().toClassName())) {
+      result = new AnnotatedField(field, SUBSTITUTION_DELETE);
+    }
+  }
+}

settings.gradle

@@ -174,6 +174,7 @@ include ':dd-java-agent:instrumentation:exception-profiling'
 include ':dd-java-agent:instrumentation:finatra-2.9'
 include ':dd-java-agent:instrumentation:glassfish'
 include ':dd-java-agent:instrumentation:google-http-client'
+include ':dd-java-agent:instrumentation:graal:native-image'
 include ':dd-java-agent:instrumentation:graphql-java-14.0'
 include ':dd-java-agent:instrumentation:grizzly-2'
 include ':dd-java-agent:instrumentation:grizzly-client-1.9'