Merge pull request DataDog#6887 from DataDog/mcculls/support-deferred…

…-instrumentation [EXPERIMENTAL] Support deferred matching and transformation for particular class-loaders
galadrielcik · May 2, 2024 · 0b24171 · 0b24171
2 parents 5c8d55c + efe7fe0
commit 0b24171
Show file tree

Hide file tree

Showing 10 changed files with 267 additions and 24 deletions.
diff --git a/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java b/dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java
@@ -1,8 +1,21 @@
 package datadog.trace.agent.tooling;
 
+import static datadog.trace.api.config.TraceInstrumentationConfig.EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL;
+import static datadog.trace.util.AgentThreadFactory.AgentThread.RETRANSFORMER;
+
+import datadog.trace.agent.tooling.bytebuddy.matcher.CustomExcludes;
+import datadog.trace.agent.tooling.bytebuddy.matcher.ProxyClassIgnores;
+import datadog.trace.api.InstrumenterConfig;
+import datadog.trace.api.time.TimeUtils;
+import datadog.trace.util.AgentTaskScheduler;
+import java.lang.instrument.Instrumentation;
 import java.security.ProtectionDomain;
+import java.util.ArrayList;
 import java.util.BitSet;
+import java.util.Iterator;
 import java.util.List;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.utility.JavaModule;
@@ -13,6 +26,14 @@
 final class CombiningMatcher implements AgentBuilder.RawMatcher {
   private static final Logger log = LoggerFactory.getLogger(CombiningMatcher.class);
 
+  private static final boolean DEFER_MATCHING =
+      null != InstrumenterConfig.get().deferIntegrationsUntil();
+
+  private static final Set<String> DEFERRED_CLASSLOADER_NAMES =
+      InstrumenterConfig.get().getDeferredClassLoaders();
+
+  private static final boolean DEFER_ALL = DEFERRED_CLASSLOADER_NAMES.isEmpty();
+
   // optimization to avoid repeated allocations inside BitSet as matched ids are set
   static final int MAX_COMBINED_ID_HINT = 512;
 
@@ -25,9 +46,16 @@ final class CombiningMatcher implements AgentBuilder.RawMatcher {
   private final BitSet knownTypesMask;
   private final MatchRecorder[] matchers;
 
-  CombiningMatcher(BitSet knownTypesMask, List<MatchRecorder> matchers) {
+  private volatile boolean deferring;
+
+  CombiningMatcher(
+      Instrumentation instrumentation, BitSet knownTypesMask, List<MatchRecorder> matchers) {
     this.knownTypesMask = knownTypesMask;
     this.matchers = matchers.toArray(new MatchRecorder[0]);
+
+    if (DEFER_MATCHING) {
+      scheduleResumeMatching(instrumentation, InstrumenterConfig.get().deferIntegrationsUntil());
+    }
   }
 
   @Override
@@ -38,6 +66,11 @@ public boolean matches(
       Class<?> classBeingRedefined,
       ProtectionDomain pd) {
 
+    // check initial requests to see if we should defer matching until retransformation
+    if (DEFER_MATCHING && null == classBeingRedefined && deferring && isDeferred(classLoader)) {
+      return false;
+    }
+
     BitSet ids = recordedMatches.get();
     ids.clear();
 
@@ -63,4 +96,103 @@ public boolean matches(
 
     return !ids.isEmpty();
   }
+
+  /** Arranges for any deferred matching to resume at the requested trigger point. */
+  private void scheduleResumeMatching(Instrumentation instrumentation, String untilTrigger) {
+    if (null != untilTrigger && !untilTrigger.isEmpty()) {
+      long delay = TimeUtils.parseSimpleDelay(untilTrigger);
+      if (delay < 0) {
+        log.info(
+            "Unrecognized value for dd.{}: {}",
+            EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL,
+            untilTrigger);
+      } else if (delay >= 5) { // don't bother deferring small delays
+
+        new AgentTaskScheduler(RETRANSFORMER)
+            .schedule(this::resumeMatching, instrumentation, delay, TimeUnit.SECONDS);
+
+        deferring = true;
+      }
+    }
+  }
+
+  /**
+   * Scans loaded classes to find which ones we should retransform to resume matching them.
+   *
+   * <p>We try to only trigger retransformations for classes we know would match. Caching and
+   * memoization means running matching twice is cheaper than unnecessary retransformations.
+   */
+  void resumeMatching(Instrumentation instrumentation) {
+    if (!deferring) {
+      return;
+    }
+
+    deferring = false;
+
+    Iterator<Iterable<Class<?>>> rediscovery =
+        AgentStrategies.rediscoveryStrategy().resolve(instrumentation).iterator();
+
+    List<Class<?>> resuming = new ArrayList<>();
+    while (rediscovery.hasNext()) {
+      for (Class<?> clazz : rediscovery.next()) {
+        ClassLoader classLoader = clazz.getClassLoader();
+        if (isDeferred(classLoader)
+            && !wouldIgnore(clazz.getName())
+            && instrumentation.isModifiableClass(clazz)
+            && wouldMatch(classLoader, clazz)) {
+          resuming.add(clazz);
+        }
+      }
+    }
+
+    try {
+      log.debug("Resuming deferred matching for {}", resuming);
+      instrumentation.retransformClasses(resuming.toArray(new Class[0]));
+    } catch (Throwable e) {
+      log.debug("Problem resuming deferred matching", e);
+    }
+  }
+
+  /**
+   * Tests whether matches involving this class-loader should be deferred until later.
+   *
+   * <p>The bootstrap class-loader is never deferred.
+   */
+  private static boolean isDeferred(ClassLoader classLoader) {
+    return null != classLoader
+        && (DEFER_ALL || DEFERRED_CLASSLOADER_NAMES.contains(classLoader.getClass().getName()));
+  }
+
+  /** Tests whether this class would be ignored on retransformation. */
+  private static boolean wouldIgnore(String name) {
+    return name.indexOf('/') >= 0 // don't retransform lambdas
+        || CustomExcludes.isExcluded(name)
+        || ProxyClassIgnores.isIgnored(name);
+  }
+
+  /** Tests whether this class would be matched at least once on retransformation. */
+  private boolean wouldMatch(ClassLoader classLoader, Class<?> clazz) {
+    BitSet ids = recordedMatches.get();
+    ids.clear();
+
+    knownTypesIndex.apply(clazz.getName(), knownTypesMask, ids);
+    if (!ids.isEmpty()) {
+      return true;
+    }
+
+    TypeDescription target = new TypeDescription.ForLoadedType(clazz);
+
+    for (MatchRecorder matcher : matchers) {
+      try {
+        matcher.record(target, classLoader, clazz, ids);
+        if (!ids.isEmpty()) {
+          return true;
+        }
+      } catch (Throwable ignore) {
+        // skip misbehaving matchers
+      }
+    }
+
+    return false;
+  }
 }
diff --git a/.../agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java b/.../agent-builder/src/main/java/datadog/trace/agent/tooling/CombiningTransformerBuilder.java
@@ -257,7 +257,7 @@ public ClassFileTransformer installOn(Instrumentation instrumentation) {
     }
 
     return agentBuilder
-        .type(new CombiningMatcher(knownTypesMask, matchers))
+        .type(new CombiningMatcher(instrumentation, knownTypesMask, matchers))
         .and(NOT_DECORATOR_MATCHER)
         .transform(defaultTransformers())
         .transform(new SplittingTransformer(transformers))

diff --git a/...ling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/ClassLoaderMatchers.java b/...ling/src/main/java/datadog/trace/agent/tooling/bytebuddy/matcher/ClassLoaderMatchers.java
@@ -14,6 +14,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import net.bytebuddy.matcher.ElementMatcher;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,8 +26,10 @@ public final class ClassLoaderMatchers {
 
   private static final ClassLoader BOOTSTRAP_CLASSLOADER = null;
 
-  private static final boolean HAS_CLASSLOADER_EXCLUDES =
-      !InstrumenterConfig.get().getExcludedClassLoaders().isEmpty();
+  private static final Set<String> EXCLUDED_CLASSLOADER_NAMES =
+      InstrumenterConfig.get().getExcludedClassLoaders();
+
+  private static final boolean CHECK_EXCLUDES = !EXCLUDED_CLASSLOADER_NAMES.isEmpty();
 
   /** A private constructor that must not be invoked. */
   private ClassLoaderMatchers() {
@@ -45,8 +48,8 @@ public static boolean canSkipClassLoaderByName(final ClassLoader loader) {
       case "datadog.trace.bootstrap.DatadogClassLoader":
         return true;
     }
-    if (HAS_CLASSLOADER_EXCLUDES) {
-      return InstrumenterConfig.get().getExcludedClassLoaders().contains(classLoaderName);
+    if (CHECK_EXCLUDES) {
+      return EXCLUDED_CLASSLOADER_NAMES.contains(classLoaderName);
     }
     return false;
   }

diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java
@@ -29,6 +29,10 @@ public final class TraceInstrumentationConfig {
   public static final String TRACE_CLASSES_EXCLUDE_FILE = "trace.classes.exclude.file";
   public static final String TRACE_CLASSLOADERS_EXCLUDE = "trace.classloaders.exclude";
   public static final String TRACE_CODESOURCES_EXCLUDE = "trace.codesources.exclude";
+  public static final String TRACE_CLASSLOADERS_DEFER = "trace.classloaders.defer";
+
+  public static final String EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL =
+      "experimental.defer.integrations.until";
 
   @SuppressWarnings("unused")
   public static final String TRACE_TESTS_ENABLED = "trace.tests.enabled";

diff --git a/dd-trace-core/src/main/java/datadog/trace/core/flare/TracerFlareService.java b/dd-trace-core/src/main/java/datadog/trace/core/flare/TracerFlareService.java
@@ -6,6 +6,7 @@
 import datadog.trace.api.Config;
 import datadog.trace.api.DynamicConfig;
 import datadog.trace.api.flare.TracerFlare;
+import datadog.trace.api.time.TimeUtils;
 import datadog.trace.core.CoreTracer;
 import datadog.trace.core.DDTraceCoreInfo;
 import datadog.trace.logging.GlobalLogLevelSwitcher;
@@ -27,8 +28,6 @@
 import java.time.ZonedDateTime;
 import java.util.Collections;
 import java.util.concurrent.TimeUnit;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 import java.util.zip.ZipOutputStream;
 import okhttp3.HttpUrl;
 import okhttp3.MediaType;
@@ -49,8 +48,6 @@ final class TracerFlareService {
 
   private static final MediaType OCTET_STREAM = MediaType.get("application/octet-stream");
 
-  private static final Pattern DELAY_TRIGGER = Pattern.compile("(\\d+)([HhMmSs]?)");
-
   private final AgentTaskScheduler scheduler = new AgentTaskScheduler(TRACER_FLARE);
 
   private final Config config;
@@ -81,20 +78,11 @@ final class TracerFlareService {
 
   private void applyTriageReportTrigger(String triageTrigger) {
     if (null != triageTrigger && !triageTrigger.isEmpty()) {
-      Matcher delayMatcher = DELAY_TRIGGER.matcher(triageTrigger);
-      if (delayMatcher.matches()) {
-        long delay = Integer.parseInt(delayMatcher.group(1));
-        String unit = delayMatcher.group(2);
-        if ("H".equalsIgnoreCase(unit)) {
-          delay = TimeUnit.HOURS.toSeconds(delay);
-        } else if ("M".equalsIgnoreCase(unit)) {
-          delay = TimeUnit.MINUTES.toSeconds(delay);
-        } else {
-          // already in seconds
-        }
-        scheduleTriageReport(delay);
-      } else {
+      long delay = TimeUtils.parseSimpleDelay(triageTrigger);
+      if (delay < 0) {
         log.info("Unrecognized triage trigger {}", triageTrigger);
+      } else {
+        scheduleTriageReport(delay);
       }
     }
   }

diff --git a/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java b/internal-api/src/main/java/datadog/trace/api/InstrumenterConfig.java
@@ -30,6 +30,7 @@
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_ENABLED;
 import static datadog.trace.api.config.ProfilingConfig.PROFILING_ENABLED_DEFAULT;
 import static datadog.trace.api.config.TraceInstrumentationConfig.AXIS_TRANSPORT_CLASS_NAME;
+import static datadog.trace.api.config.TraceInstrumentationConfig.EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL;
 import static datadog.trace.api.config.TraceInstrumentationConfig.HTTP_URL_CONNECTION_CLASS_NAME;
 import static datadog.trace.api.config.TraceInstrumentationConfig.INTEGRATIONS_ENABLED;
 import static datadog.trace.api.config.TraceInstrumentationConfig.JAX_RS_ADDITIONAL_ANNOTATIONS;
@@ -50,6 +51,7 @@
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ANNOTATION_ASYNC;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSES_EXCLUDE_FILE;
+import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSLOADERS_DEFER;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CLASSLOADERS_EXCLUDE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_CODESOURCES_EXCLUDE;
 import static datadog.trace.api.config.TraceInstrumentationConfig.TRACE_ENABLED;
@@ -120,6 +122,9 @@ public class InstrumenterConfig {
   private final String excludedClassesFile;
   private final Set<String> excludedClassLoaders;
   private final List<String> excludedCodeSources;
+  private final Set<String> deferredClassLoaders;
+
+  private final String deferIntegrationsUntil;
 
   private final ResolverCacheConfig resolverCacheConfig;
   private final String resolverCacheDir;
@@ -206,6 +211,9 @@ private InstrumenterConfig() {
     excludedClassesFile = configProvider.getString(TRACE_CLASSES_EXCLUDE_FILE);
     excludedClassLoaders = tryMakeImmutableSet(configProvider.getList(TRACE_CLASSLOADERS_EXCLUDE));
     excludedCodeSources = tryMakeImmutableList(configProvider.getList(TRACE_CODESOURCES_EXCLUDE));
+    deferredClassLoaders = tryMakeImmutableSet(configProvider.getList(TRACE_CLASSLOADERS_DEFER));
+
+    deferIntegrationsUntil = configProvider.getString(EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL);
 
     resolverCacheConfig =
         configProvider.getEnum(
@@ -353,6 +361,14 @@ public List<String> getExcludedCodeSources() {
     return excludedCodeSources;
   }
 
+  public Set<String> getDeferredClassLoaders() {
+    return deferredClassLoaders;
+  }
+
+  public String deferIntegrationsUntil() {
+    return deferIntegrationsUntil;
+  }
+
   public int getResolverNoMatchesSize() {
     return resolverCacheConfig.noMatchesSize();
   }
@@ -512,6 +528,10 @@ public String toString() {
         + excludedClassLoaders
         + ", excludedCodeSources="
         + excludedCodeSources
+        + ", deferredClassLoaders="
+        + deferredClassLoaders
+        + ", deferIntegrationsUntil="
+        + deferIntegrationsUntil
         + ", resolverCacheConfig="
         + resolverCacheConfig
         + ", resolverCacheDir="

diff --git a/internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java b/internal-api/src/main/java/datadog/trace/util/AgentThreadFactory.java
@@ -51,7 +51,9 @@ public enum AgentThread {
     CI_GIT_DATA_SHUTDOWN_HOOK("dd-ci-git-data-shutdown-hook"),
     CI_TEST_EVENTS_SHUTDOWN_HOOK("dd-ci-test-events-shutdown-hook"),
     CI_PROJECT_CONFIGURATOR("dd-ci-project-configurator"),
-    CI_SIGNAL_SERVER("dd-ci-signal-server");
+    CI_SIGNAL_SERVER("dd-ci-signal-server"),
+
+    RETRANSFORMER("dd-retransformer");
 
     public final String threadName;
 

diff --git a/utils/time-utils/build.gradle b/utils/time-utils/build.gradle
@@ -1,5 +1,12 @@
 apply from: "$rootDir/gradle/java.gradle"
 
+ext {
+  excludedClassesCoverage = [
+    'datadog.trace.api.time.ControllableTimeSource:',
+    'datadog.trace.api.time.SystemTimeSource'
+  ]
+}
+
 dependencies {
   testImplementation project(':utils:test-utils')
 }
diff --git a/utils/time-utils/src/main/java/datadog/trace/api/time/TimeUtils.java b/utils/time-utils/src/main/java/datadog/trace/api/time/TimeUtils.java
@@ -0,0 +1,37 @@
+package datadog.trace.api.time;
+
+import java.util.concurrent.TimeUnit;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public abstract class TimeUtils {
+
+  /** Number followed by an optional time unit of hours (h), minutes (m), or seconds (s). */
+  private static final Pattern SIMPLE_DELAY_PATTERN = Pattern.compile("(\\d+)([HhMmSs]?)");
+
+  /**
+   * Parses the string as a simple delay, such as "30s" or "10m".
+   *
+   * @param delayString number followed by an optional time unit
+   * @return delay in seconds; -1 if the string cannot be parsed
+   */
+  public static long parseSimpleDelay(String delayString) {
+    if (null != delayString) {
+      Matcher delayMatcher = SIMPLE_DELAY_PATTERN.matcher(delayString);
+      if (delayMatcher.matches()) {
+        long delay = Integer.parseInt(delayMatcher.group(1));
+        String unit = delayMatcher.group(2);
+        if ("H".equalsIgnoreCase(unit)) {
+          return TimeUnit.HOURS.toSeconds(delay);
+        } else if ("M".equalsIgnoreCase(unit)) {
+          return TimeUnit.MINUTES.toSeconds(delay);
+        } else {
+          return delay; // already in seconds
+        }
+      }
+    }
+    return -1; // unrecognized
+  }
+
+  private TimeUtils() {}
+}