[IAST] Call Site Instrumentation support

dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentTransformerBuilder.java

@@ -140,6 +140,8 @@ private AgentBuilder.RawMatcher matcher(Instrumenter.Default instrumenter) {
       typeMatcher = ((Instrumenter.ForTypeHierarchy) instrumenter).hierarchyMatcher();
     } else if (instrumenter instanceof Instrumenter.ForConfiguredType) {
       typeMatcher = none(); // handle below, just like when it's combined with other matchers
+    } else if (instrumenter instanceof Instrumenter.ForCallSite) {
+      typeMatcher = ((Instrumenter.ForCallSite) instrumenter).callerType();
     } else {
       return AgentBuilder.RawMatcher.Trivial.NON_MATCHING;
     }

dd-java-agent/agent-tooling/agent-tooling.gradle

@@ -1,3 +1,6 @@
+plugins {
+  id 'me.champeau.jmh'
+}
 apply from: "$rootDir/gradle/java.gradle"
 apply from: "$rootDir/gradle/tries.gradle"
 
@@ -33,4 +36,22 @@ dependencies {
 
   testImplementation project(':dd-java-agent:testing')
   testImplementation group: 'com.google.guava', name: 'guava-testlib', version: '20.0'
+
+  jmhImplementation group: 'org.springframework.boot', name: 'spring-boot-starter-web', version: '2.3.5.RELEASE'
+}
+
+jmh {
+  jmhVersion = '1.32'
+  includeTests = true
 }
+
+forbiddenApisJmh {
+  ignoreFailures = true
+}
+
+final compileTestJava = project.tasks.compileTestJava
+compileTestJava.dependsOn(project.tasks.generateTestClassNameTries)
+
+final jmh = project.tasks.jmh
+jmh.outputs.upToDateWhen { false }
+jmh.dependsOn(compileTestJava)

dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmark.java

@@ -0,0 +1,89 @@
+package datadog.trace.agent.tooling.bytebuddy.csi;
+
+import static java.util.concurrent.TimeUnit.MILLISECONDS;
+
+import datadog.trace.agent.tooling.AgentInstaller;
+import java.lang.instrument.Instrumentation;
+import net.bytebuddy.agent.ByteBuddyAgent;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.springframework.boot.SpringApplication;
+import org.springframework.context.ConfigurableApplicationContext;
+import org.springframework.web.client.RestTemplate;
+
+@State(Scope.Benchmark)
+@BenchmarkMode(Mode.SingleShotTime)
+@OutputTimeUnit(MILLISECONDS)
+@Fork(value = 10)
+public class CallSiteBenchmark {
+
+  private Instrumentation instrumentation;
+
+  @Setup(Level.Trial)
+  public void setUp() {
+    instrumentation = ByteBuddyAgent.install();
+  }
+
+  @Benchmark
+  public void none() throws Exception {
+    runBenchmark(Type.NONE);
+  }
+
+  @Benchmark
+  public void callee() throws Exception {
+    runBenchmark(Type.CALLEE);
+  }
+
+  @Benchmark
+  public void callSite() throws Exception {
+    runBenchmark(Type.CALL_SITE);
+  }
+
+  private void runBenchmark(final Type type) throws Exception {
+    type.apply(instrumentation);
+    final Class<?> server = Class.forName("foo.bar.DummyApplication");
+    try (ConfigurableApplicationContext context = SpringApplication.run(server)) {
+      final RestTemplate template = new RestTemplate();
+      final String url = "http://localhost:8080/benchmark?param=Hello!";
+      final String response = template.getForObject(url, String.class);
+      type.validate(response);
+    }
+  }
+
+  enum Type {
+    NONE(null),
+    CALL_SITE("callSite"),
+    CALLEE("callee");
+
+    private final String instrumenter;
+
+    Type(final String instrumenter) {
+      this.instrumenter = instrumenter;
+    }
+
+    public void apply(final Instrumentation instrumentation) {
+      if (instrumenter != null) {
+        System.setProperty("dd.benchmark.instrumentation", instrumenter);
+        AgentInstaller.installBytebuddyAgent(instrumentation);
+      }
+    }
+
+    public void validate(final String response) {
+      if (response == null) {
+        throw new RuntimeException("Empty response received");
+      }
+      String expected = instrumenter == null ? "Hello!" : "Hello! [Transformed]";
+      if (!expected.equals(response)) {
+        throw new RuntimeException(
+            String.format("Wrong response, expected '%s' but received '%s'", expected, response));
+      }
+    }
+  }
+}

dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkHelper.java

@@ -0,0 +1,17 @@
+package datadog.trace.agent.tooling.bytebuddy.csi;
+
+import javax.servlet.ServletRequest;
+import net.bytebuddy.asm.Advice;
+
+public class CallSiteBenchmarkHelper {
+
+  @Advice.OnMethodExit
+  public static void adviceCallee(@Advice.Return(readOnly = false) String result) {
+    final String currentValue = result;
+    result = currentValue + " [Transformed]";
+  }
+
+  public static String adviceCallSite(final ServletRequest request, final String parameter) {
+    return request.getParameter(parameter) + " [Transformed]";
+  }
+}

dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkInstrumenter.java

@@ -0,0 +1,88 @@
+package datadog.trace.agent.tooling.bytebuddy.csi;
+
+import datadog.trace.agent.tooling.csi.CallSiteAdvice;
+import datadog.trace.agent.tooling.csi.CallSiteAdvice.HasHelpers;
+import datadog.trace.agent.tooling.csi.Pointcut;
+import java.util.Collections;
+import java.util.List;
+import java.util.Set;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.jar.asm.MethodVisitor;
+import net.bytebuddy.jar.asm.Opcodes;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class CallSiteBenchmarkInstrumenter extends CallSiteInstrumenter
+    implements ElementMatcher<TypeDescription> {
+
+  public CallSiteBenchmarkInstrumenter() {
+    super(buildCallSites(), "call-site");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> callerType() {
+    return this;
+  }
+
+  @Override
+  public boolean matches(final TypeDescription target) {
+    return CallSiteTrie.apply(target.getName()) != 1;
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return "callSite".equals(System.getProperty("dd.benchmark.instrumentation", ""));
+  }
+
+  @Override
+  public boolean isApplicable(final Set<TargetSystem> enabledSystems) {
+    return true;
+  }
+
+  private static List<CallSiteAdvice> buildCallSites() {
+    return Collections.<CallSiteAdvice>singletonList(new GetParameterCallSite());
+  }
+
+  private static class GetParameterCallSite implements CallSiteAdvice, HasHelpers, Pointcut {
+
+    @Override
+    public Pointcut pointcut() {
+      return this;
+    }
+
+    @Override
+    public void apply(
+        final MethodVisitor mv,
+        final int opcode,
+        final String owner,
+        final String name,
+        final String descriptor,
+        final boolean isInterface) {
+      mv.visitMethodInsn(
+          Opcodes.INVOKESTATIC,
+          "datadog/trace/agent/tooling/bytebuddy/csi/CallSiteBenchmarkHelper",
+          "adviceCallSite",
+          "(Ljavax/servlet/ServletRequest;Ljava/lang/String;)Ljava/lang/String;",
+          false);
+    }
+
+    @Override
+    public String[] helperClassNames() {
+      return new String[] {CallSiteBenchmarkHelper.class.getName()};
+    }
+
+    @Override
+    public String type() {
+      return "javax/servlet/ServletRequest";
+    }
+
+    @Override
+    public String method() {
+      return "getParameter";
+    }
+
+    @Override
+    public String descriptor() {
+      return "(Ljava/lang/String;)Ljava/lang/String;";
+    }
+  }
+}

dd-java-agent/agent-tooling/src/jmh/java/datadog/trace/agent/tooling/bytebuddy/csi/CalleeBenchmarkInstrumenter.java

@@ -0,0 +1,51 @@
+package datadog.trace.agent.tooling.bytebuddy.csi;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassesNamed;
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.returns;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+import java.util.Set;
+import net.bytebuddy.description.type.TypeDescription;
+import net.bytebuddy.matcher.ElementMatcher;
+
+public class CalleeBenchmarkInstrumenter extends Instrumenter.Default
+    implements Instrumenter.ForTypeHierarchy {
+
+  public CalleeBenchmarkInstrumenter() {
+    super("callee");
+  }
+
+  @Override
+  public ElementMatcher<ClassLoader> classLoaderMatcher() {
+    return hasClassesNamed("javax.servlet.http.HttpServlet");
+  }
+
+  @Override
+  public ElementMatcher<TypeDescription> hierarchyMatcher() {
+    return named("org.apache.catalina.connector.Request");
+  }
+
+  @Override
+  public void adviceTransformations(final AdviceTransformation transformation) {
+    transformation.applyAdvice(
+        named("getParameter").and(takesArguments(String.class)).and(returns(String.class)),
+        CallSiteBenchmarkHelper.class.getName());
+  }
+
+  @Override
+  public String[] helperClassNames() {
+    return new String[] {CallSiteBenchmarkHelper.class.getName()};
+  }
+
+  @Override
+  public boolean isEnabled() {
+    return "callee".equals(System.getProperty("dd.benchmark.instrumentation", ""));
+  }
+
+  @Override
+  public boolean isApplicable(final Set<TargetSystem> enabledSystems) {
+    return true;
+  }
+}

dd-java-agent/agent-tooling/src/jmh/java/foo/bar/DummyApplication.java

@@ -0,0 +1,6 @@
+package foo.bar;
+
+import org.springframework.boot.autoconfigure.SpringBootApplication;
+
+@SpringBootApplication(scanBasePackages = "foo.bar")
+public class DummyApplication {}

dd-java-agent/agent-tooling/src/jmh/java/foo/bar/DummyController.java

@@ -0,0 +1,14 @@
+package foo.bar;
+
+import javax.servlet.ServletRequest;
+import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.RestController;
+
+@RestController
+public class DummyController {
+
+  @GetMapping("/benchmark")
+  public String index(final ServletRequest request) {
+    return request.getParameter("param");
+  }
+}

dd-java-agent/agent-tooling/src/jmh/resources/META-INF/services/datadog.trace.agent.tooling.Instrumenter

@@ -0,0 +1,2 @@
+datadog.trace.agent.tooling.bytebuddy.csi.CalleeBenchmarkInstrumenter
+datadog.trace.agent.tooling.bytebuddy.csi.CallSiteBenchmarkInstrumenter

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

@@ -72,6 +72,11 @@ interface ForTypeHierarchy {
     ElementMatcher<TypeDescription> hierarchyMatcher();
   }
 
+  /** Instrumentation that matches based on the caller of an instruction. */
+  interface ForCallSite {
+    ElementMatcher<TypeDescription> callerType();
+  }
+
   /** Instrumentation that can optionally widen matching to consider the type hierarchy. */
   interface CanShortcutTypeMatching extends ForKnownTypes, ForTypeHierarchy {
     boolean onlyMatchKnownTypes();