Add opt-in instrumentation support for generated lambda classes

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformer.java

@@ -0,0 +1,14 @@
+package datadog.trace.bootstrap.instrumentation.java.lang.invoke;
+
+/** Transforms a generated lambda class before it is defined. */
+public interface LambdaTransformer {
+  /**
+   * @param slashClassName internal (slash-separated) name of the generated lambda class
+   * @param targetClass the class declaring the lambda
+   * @param classBytes the freshly generated lambda class bytes
+   * @param interfaceClassName the functional interface implemented by the lambda
+   * @return the transformed bytes, or {@code null}/the original bytes if unchanged
+   */
+  byte[] transform(
+      String slashClassName, Class<?> targetClass, byte[] classBytes, String interfaceClassName);
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHelper.java

@@ -0,0 +1,59 @@
+package datadog.trace.bootstrap.instrumentation.java.lang.invoke;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Transforms eligible lambda bytes before definition, falling back to the original on failure. */
+public final class LambdaTransformerHelper {
+  private static final Logger log = LoggerFactory.getLogger(LambdaTransformerHelper.class);
+
+  // Agent transformation may itself create lambdas.
+  private static final ThreadLocal<Boolean> TRANSFORMING = new ThreadLocal<>();
+
+  private LambdaTransformerHelper() {}
+
+  /**
+   * @param classBytes the generated lambda class bytes
+   * @param lambdaClassName internal (slash-separated) name of the generated lambda class
+   * @param targetClass the class declaring the lambda
+   * @param interfaceClass the functional interface implemented by the lambda
+   * @return possibly transformed bytes; the original bytes on any failure
+   */
+  public static byte[] transform(
+      byte[] classBytes, String lambdaClassName, Class<?> targetClass, Class<?> interfaceClass) {
+    try {
+      if (interfaceClass == null) {
+        return classBytes;
+      }
+      String interfaceName = interfaceClass.getName();
+      LambdaTransformer transformer = LambdaTransformerHolder.get();
+      if (transformer == null) {
+        log.debug("Lambda {} skipped: no transformer registered", lambdaClassName);
+        return classBytes;
+      }
+      if (targetClass == null) {
+        log.debug("Lambda {} skipped: no target class", lambdaClassName);
+        return classBytes;
+      }
+      if (Boolean.TRUE.equals(TRANSFORMING.get())) {
+        log.debug("Lambda {} skipped: re-entrant transform", lambdaClassName);
+        return classBytes;
+      }
+      TRANSFORMING.set(Boolean.TRUE);
+      try {
+        byte[] result =
+            transformer.transform(lambdaClassName, targetClass, classBytes, interfaceName);
+        if (result == null) {
+          log.debug("Lambda {} not transformed", lambdaClassName);
+          return classBytes;
+        }
+        return result;
+      } finally {
+        TRANSFORMING.remove();
+      }
+    } catch (Throwable e) {
+      log.debug("Lambda {} skipped: {}", lambdaClassName, e.toString());
+      return classBytes;
+    }
+  }
+}

dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/java/lang/invoke/LambdaTransformerHolder.java

@@ -0,0 +1,19 @@
+package datadog.trace.bootstrap.instrumentation.java.lang.invoke;
+
+/**
+ * Holds the {@link LambdaTransformer} registered by the agent installer. Lives on the bootstrap
+ * class path so it is reachable from instrumented {@code java.lang.invoke} code.
+ */
+public final class LambdaTransformerHolder {
+  private static volatile LambdaTransformer transformer;
+
+  private LambdaTransformerHolder() {}
+
+  public static void set(LambdaTransformer transformer) {
+    LambdaTransformerHolder.transformer = transformer;
+  }
+
+  public static LambdaTransformer get() {
+    return transformer;
+  }
+}

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -2,9 +2,11 @@
 
 import static datadog.trace.agent.tooling.ExtensionFinder.findExtensions;
 import static datadog.trace.agent.tooling.ExtensionLoader.loadExtensions;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnores.isIgnored;
 import static datadog.trace.agent.tooling.bytebuddy.matcher.GlobalIgnoresMatcher.globalIgnoresMatcher;
 import static net.bytebuddy.matcher.ElementMatchers.isDefaultFinalizer;
 
+import datadog.environment.JavaVirtualMachine;
 import datadog.environment.SystemProperties;
 import datadog.instrument.fieldinject.GlobalObjectStore;
 import datadog.trace.agent.tooling.bytebuddy.SharedTypePools;
@@ -20,6 +22,9 @@
 import datadog.trace.api.telemetry.IntegrationsCollector;
 import datadog.trace.bootstrap.FieldBackedContextAccessor;
 import datadog.trace.bootstrap.instrumentation.java.concurrent.ExcludeFilter;
+import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformer;
+import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHelper;
+import datadog.trace.bootstrap.instrumentation.java.lang.invoke.LambdaTransformerHolder;
 import datadog.trace.bootstrap.instrumentation.java.module.JpmsHelper;
 import datadog.trace.util.AgentTaskScheduler;
 import de.thetaphi.forbiddenapis.SuppressForbidden;
@@ -36,6 +41,7 @@
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import java.util.function.BooleanSupplier;
+import java.util.function.Function;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.agent.builder.AgentBuilder;
 import net.bytebuddy.description.type.TypeDescription;
@@ -166,6 +172,15 @@ public static ClassFileTransformer installBytebuddyAgent(
             // .with(AgentBuilder.LambdaInstrumentationStrategy.ENABLED)
             .ignore(globalIgnoresMatcher(skipAdditionalLibraryMatcher));
 
+    boolean lambdaTransformationEnabled =
+        !Platform.isNativeImageBuilder()
+            && InstrumenterConfig.get()
+                .isIntegrationEnabled(Collections.singleton("lambda"), false);
+    if (lambdaTransformationEnabled) {
+      // The injected metafactory call needs java.base to read the bootstrap helper's module.
+      agentBuilder = agentBuilder.assureReadEdgeTo(inst, LambdaTransformerHelper.class);
+    }
+
     if (DEBUG) {
       agentBuilder =
           agentBuilder
@@ -216,7 +231,8 @@ public static ClassFileTransformer installBytebuddyAgent(
     }
 
     CombiningTransformerBuilder transformerBuilder =
-        new CombiningTransformerBuilder(agentBuilder, instrumenterIndex, enabledSystems);
+        new CombiningTransformerBuilder(
+            agentBuilder, instrumenterIndex, enabledSystems, lambdaTransformationEnabled);
 
     int installedCount = 0;
     for (InstrumenterModule module : instrumenterModules) {
@@ -265,12 +281,96 @@ public void applied(Iterable<String> instrumentationNames) {
 
     InstrumenterState.resetDefaultState();
     try {
-      return transformerBuilder.installOn(inst);
+      ClassFileTransformer classFileTransformer = transformerBuilder.installOn(inst);
+      registerLambdaTransformer(
+          lambdaTransformationEnabled, classFileTransformer, transformerBuilder.lambdaInterfaces());
+      return classFileTransformer;
     } finally {
       SharedTypePools.endInstall();
     }
   }
 
+  /** Registers the installed class-file transformer for generated lambdas. */
+  static void registerLambdaTransformer(
+      final boolean enabled,
+      final ClassFileTransformer classFileTransformer,
+      final String[] lambdaInterfaces) {
+    if (!enabled) {
+      // Agent installation can be repeated in tests and embedded environments.
+      LambdaTransformerHolder.set(null);
+      return;
+    }
+    LambdaTransformer transformer =
+        lambdaInterfaces.length == 0 ? null : newLambdaTransformer(classFileTransformer);
+    LambdaTransformerHolder.set(filterLambdaTransformer(transformer, lambdaInterfaces));
+  }
+
+  static LambdaTransformer filterLambdaTransformer(
+      final LambdaTransformer transformer, final String[] lambdaInterfaces) {
+    if (transformer == null) {
+      return null;
+    }
+    return (className, targetClass, classBytes, interfaceName) -> {
+      for (String enabledInterface : lambdaInterfaces) {
+        if (enabledInterface.equals(interfaceName)) {
+          // Apply the system-level name filter before entering the full transformer pipeline.
+          if (isIgnored(targetClass.getName(), true)) {
+            return null;
+          }
+          return transformer.transform(className, targetClass, classBytes, interfaceName);
+        }
+      }
+      return null;
+    };
+  }
+
+  /**
+   * Java 9+ requires the module-aware transformer for injected read edges. Failure must disable
+   * lambda transformation rather than fall back to the module-less overload.
+   */
+  @SuppressWarnings("unchecked")
+  private static LambdaTransformer newLambdaTransformer(
+      final ClassFileTransformer classFileTransformer) {
+    if (JavaVirtualMachine.isJavaVersionAtLeast(9)) {
+      try {
+        Function<ClassFileTransformer, LambdaTransformer> factory =
+            (Function<ClassFileTransformer, LambdaTransformer>)
+                Instrumenter.class
+                    .getClassLoader()
+                    .loadClass("datadog.trace.agent.tooling.bytebuddy.DDJava9LambdaTransformer")
+                    .getField("FACTORY")
+                    .get(null);
+        return factory.apply(classFileTransformer);
+      } catch (Throwable e) {
+        log.debug("Problem loading Java 9 lambda transformer, disabling lambda transformation", e);
+        return null;
+      }
+    }
+    // Avoid invoking the instrumented metafactory while installing its transformer.
+    return new LambdaTransformer() {
+      @Override
+      public byte[] transform(
+          String slashClassName,
+          Class<?> targetClass,
+          byte[] classBytes,
+          String interfaceClassName) {
+        TypePoolFacade.beginLambdaTransform(interfaceClassName);
+        try {
+          return classFileTransformer.transform(
+              targetClass.getClassLoader(),
+              slashClassName,
+              null,
+              targetClass.getProtectionDomain(),
+              classBytes);
+        } catch (Throwable ignored) {
+          return null;
+        } finally {
+          TypePoolFacade.endLambdaTransform();
+        }
+      }
+    };
+  }
+
   /** Returns an iterable that combines the original sequence with any discovered extensions. */
   private static Iterable<InstrumenterModule> withExtensions(Iterable<InstrumenterModule> initial) {
     String extensionsPath = InstrumenterConfig.get().getTraceExtensionsPath();

dd-java-agent/agent-installer/src/main/java/datadog/trace/agent/tooling/CombiningMatcher.java

@@ -2,18 +2,22 @@
 
 import static datadog.trace.api.config.TraceInstrumentationConfig.EXPERIMENTAL_DEFER_INTEGRATIONS_UNTIL;
 import static datadog.trace.util.AgentThreadFactory.AgentThread.RETRANSFORMER;
+import static java.util.Collections.unmodifiableMap;
 
 import datadog.trace.agent.tooling.bytebuddy.matcher.CustomExcludes;
 import datadog.trace.agent.tooling.bytebuddy.matcher.ProxyClassIgnores;
+import datadog.trace.agent.tooling.bytebuddy.outline.TypePoolFacade;
 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.HashMap;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import net.bytebuddy.agent.builder.AgentBuilder;
@@ -45,13 +49,22 @@ final class CombiningMatcher implements AgentBuilder.RawMatcher {
 
   private final BitSet knownTypesMask;
   private final MatchRecorder[] matchers;
+  private final Map<String, LambdaMatchRecorder[]> lambdaMatchers;
 
   private volatile boolean deferring;
 
   CombiningMatcher(
-      Instrumentation instrumentation, BitSet knownTypesMask, List<MatchRecorder> matchers) {
+      Instrumentation instrumentation,
+      BitSet knownTypesMask,
+      List<MatchRecorder> matchers,
+      Map<String, List<LambdaMatchRecorder>> lambdaMatchers) {
     this.knownTypesMask = knownTypesMask;
     this.matchers = matchers.toArray(new MatchRecorder[0]);
+    Map<String, LambdaMatchRecorder[]> lambdaMatchersByInterface = new HashMap<>();
+    lambdaMatchers.forEach(
+        (name, recorders) ->
+            lambdaMatchersByInterface.put(name, recorders.toArray(new LambdaMatchRecorder[0])));
+    this.lambdaMatchers = unmodifiableMap(lambdaMatchersByInterface);
 
     if (DEFER_MATCHING) {
       scheduleResumeMatching(instrumentation, InstrumenterConfig.get().deferIntegrationsUntil());
@@ -66,15 +79,43 @@ public boolean matches(
       Class<?> classBeingRedefined,
       ProtectionDomain pd) {
 
+    String lambdaInterface = lambdaMatchers.isEmpty() ? null : TypePoolFacade.lambdaInterface();
+
     // check initial requests to see if we should defer matching until retransformation
-    if (DEFER_MATCHING && null == classBeingRedefined && deferring && isDeferred(classLoader)) {
+    // Generated lambda classes cannot be retransformed because they are not yet defined here.
+    if (null == lambdaInterface
+        && DEFER_MATCHING
+        && null == classBeingRedefined
+        && deferring
+        && isDeferred(classLoader)) {
       return false;
     }
 
     BitSet ids = recordedMatches.get();
     ids.clear();
 
     long fromTick = InstrumenterMetrics.tick();
+    if (null != lambdaInterface) {
+      LambdaMatchRecorder[] recorders = lambdaMatchers.get(lambdaInterface);
+      if (null != recorders) {
+        for (LambdaMatchRecorder recorder : recorders) {
+          try {
+            recorder.record(target, classLoader, ids);
+          } catch (Throwable e) {
+            if (log.isDebugEnabled()) {
+              log.debug(
+                  "Lambda instrumentation matcher unexpected exception - interface={} matcher={}",
+                  lambdaInterface,
+                  recorder.describe(),
+                  e);
+            }
+          }
+        }
+      }
+      InstrumenterMetrics.matchType(fromTick);
+      return !ids.isEmpty();
+    }
+
     knownTypesIndex.apply(target.getName(), knownTypesMask, ids);
     if (ids.isEmpty()) {
       InstrumenterMetrics.knownTypeMiss(fromTick);