Merge branch 'master' into andrea.marziali/jetty-websocket

Author: amarziali

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/AdviceGeneratorImpl.java

@@ -2,6 +2,7 @@
 
 import static com.github.javaparser.ast.Modifier.Keyword.PUBLIC;
 import static datadog.trace.plugin.csi.impl.CallSiteFactory.typeResolver;
+import static datadog.trace.plugin.csi.util.CallSiteConstants.ADVICE_TYPE_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.AUTO_SERVICE_FQDN;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.CALL_SITES_CLASS;
 import static datadog.trace.plugin.csi.util.CallSiteConstants.CALL_SITES_FQCN;
@@ -185,20 +186,24 @@ private void addAdviceLambda(
     final MethodType pointCut = spec.getPointcut();
     final BlockStmt adviceBody = new BlockStmt();
     final Expression advice;
+    final String type;
     if (spec.isInvokeDynamic()) {
       advice = invokeDynamicAdviceSignature(adviceBody);
     } else {
       advice = invokeAdviceSignature(adviceBody);
     }
     if (spec instanceof BeforeSpecification) {
+      type = "BEFORE";
       writeStackOperations(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
     } else if (spec instanceof AfterSpecification) {
+      type = "AFTER";
       writeStackOperations(spec, adviceBody);
       writeOriginalMethodCall(spec, adviceBody);
       writeAdviceMethodCall(spec, adviceBody);
     } else {
+      type = "AROUND";
       writeAdviceMethodCall(spec, adviceBody);
     }
     body.addStatement(
@@ -207,6 +212,10 @@ private void addAdviceLambda(
             .setName("addAdvice")
             .setArguments(
                 new NodeList<>(
+                    new FieldAccessExpr()
+                        .setScope(
+                            new TypeExpr(new ClassOrInterfaceType().setName(ADVICE_TYPE_CLASS)))
+                        .setName(type),
                     new StringLiteralExpr(pointCut.getOwner().getInternalName()),
                     new StringLiteralExpr(pointCut.getMethodName()),
                     new StringLiteralExpr(pointCut.getMethodType().getDescriptor()),

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/impl/ext/IastExtension.java

@@ -364,19 +364,19 @@ private static LambdaExpr findAdviceLambda(
     final MethodType pointcut = spec.getPointcut();
     for (final MethodCallExpr add : addAdvices) {
       final NodeList<Expression> arguments = add.getArguments();
-      final String owner = arguments.get(0).asStringLiteralExpr().asString();
+      final String owner = arguments.get(1).asStringLiteralExpr().asString();
       if (!owner.equals(pointcut.getOwner().getInternalName())) {
         continue;
       }
-      final String method = arguments.get(1).asStringLiteralExpr().asString();
+      final String method = arguments.get(2).asStringLiteralExpr().asString();
       if (!method.equals(pointcut.getMethodName())) {
         continue;
       }
-      final String description = arguments.get(2).asStringLiteralExpr().asString();
+      final String description = arguments.get(3).asStringLiteralExpr().asString();
       if (!description.equals(pointcut.getMethodType().getDescriptor())) {
         continue;
       }
-      return arguments.get(3).asLambdaExpr();
+      return arguments.get(4).asLambdaExpr();
     }
     throw new IllegalArgumentException("Cannot find lambda expression for pointcut " + pointcut);
   }

buildSrc/call-site-instrumentation-plugin/src/main/java/datadog/trace/plugin/csi/util/CallSiteConstants.java

@@ -30,6 +30,8 @@ private CallSiteConstants() {}
 
   public static final String HAS_ENABLED_PROPERTY_CLASS = CALL_SITES_CLASS + ".HasEnabledProperty";
 
+  public static final String ADVICE_TYPE_CLASS = "AdviceType";
+
   public static final String STACK_DUP_MODE_CLASS = "StackDupMode";
 
   public static final String METHOD_HANDLER_CLASS = "MethodHandler";

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/AdviceGeneratorTest.groovy

@@ -44,6 +44,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(BeforeAdvice)
       advices(0) {
+        type("BEFORE")
         pointcut('java/security/MessageDigest', 'getInstance', '(Ljava/lang/String;)Ljava/security/MessageDigest;')
         statements(
           'handler.dupParameters(descriptor, StackDupMode.COPY);',
@@ -76,6 +77,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AroundAdvice)
       advices(0) {
+        type("AROUND")
         pointcut('java/lang/String', 'replaceAll', '(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AroundAdvice", "around", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;");'
@@ -106,6 +108,7 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
       interfaces(CallSites)
       helpers(AfterAdvice)
       advices(0) {
+        type("AFTER")
         pointcut('java/lang/String', 'concat', '(Ljava/lang/String;)Ljava/lang/String;')
         statements(
           'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AdviceAssert.groovy

@@ -1,11 +1,16 @@
 package datadog.trace.plugin.csi.impl.assertion
 
 class AdviceAssert {
+  protected String type
   protected String owner
   protected String method
   protected String descriptor
   protected Collection<String> statements
 
+  void type(String type) {
+    assert type == this.type
+  }
+
   void pointcut(String owner, String method, String descriptor) {
     assert owner == this.owner
     assert method == this.method

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/assertion/AssertBuilder.groovy

@@ -83,10 +83,12 @@ class AssertBuilder<C extends CallSiteAssert> {
     return getMethodCalls(acceptMethod).findAll {
       it.nameAsString == 'addAdvice'
     }.collect {
-      def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-      final handlerLambda = it.arguments[3].asLambdaExpr()
+      final adviceType = it.arguments.get(0).asFieldAccessExpr().getName()
+      def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+      final handlerLambda = it.arguments[4].asLambdaExpr()
       final advice = handlerLambda.body.asBlockStmt().statements*.toString()
       return new AdviceAssert([
+        type      : adviceType,
         owner     : owner,
         method    : method,
         descriptor: descriptor,

buildSrc/call-site-instrumentation-plugin/src/test/groovy/datadog/trace/plugin/csi/impl/ext/IastExtensionTest.groovy

@@ -217,8 +217,8 @@ class IastExtensionTest extends BaseCsiPluginTest {
       return getMethodCalls(acceptMethod).findAll {
         it.nameAsString == 'addAdvice'
       }.collect {
-        def (owner, method, descriptor) = it.arguments.subList(0, 3)*.asStringLiteralExpr()*.asString()
-        final handlerLambda = it.arguments[3].asLambdaExpr()
+        def (owner, method, descriptor) = it.arguments.subList(1, 4)*.asStringLiteralExpr()*.asString()
+        final handlerLambda = it.arguments[4].asLambdaExpr()
         final statements = handlerLambda.body.asBlockStmt().statements
         final instrumentedStmt = statements.get(0).asIfStmt()
         final executedStmt = statements.get(1).asIfStmt()

components/context/src/main/java/datadog/context/Context.java

@@ -42,7 +42,7 @@ public interface Context {
    * @return the initial local context that all contexts extend.
    */
   static Context root() {
-    return manager().root();
+    return EmptyContext.INSTANCE;
   }
 
   /**

components/context/src/main/java/datadog/context/ContextManager.java

@@ -2,17 +2,10 @@
 
 /** Manages context across execution units. */
 public interface ContextManager {
-  /**
-   * Returns the root context.
-   *
-   * @return the initial local context that all contexts extend.
-   */
-  Context root();
-
   /**
    * Returns the context attached to the current execution unit.
    *
-   * @return the attached context; {@link #root()} if there is none.
+   * @return the attached context; {@link Context#root()} if there is none.
    */
   Context current();
 
@@ -28,7 +21,7 @@ public interface ContextManager {
    * Swaps the given context with the one attached to current execution unit.
    *
    * @param context the context to swap.
-   * @return the previously attached context; {@link #root()} if there was none.
+   * @return the previously attached context; {@link Context#root()} if there was none.
    */
   Context swap(Context context);
 

components/context/src/main/java/datadog/context/ThreadLocalContextManager.java

@@ -5,11 +5,6 @@ final class ThreadLocalContextManager implements ContextManager {
   private static final ThreadLocal<Context[]> CURRENT_HOLDER =
       ThreadLocal.withInitial(() -> new Context[] {EmptyContext.INSTANCE});
 
-  @Override
-  public Context root() {
-    return EmptyContext.INSTANCE;
-  }
-
   @Override
   public Context current() {
     return CURRENT_HOLDER.get()[0];