Add After callsites support for void methods (#8116)

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

@@ -570,11 +570,20 @@ public AfterSpecification(
 
     @Override
     protected void validateAdvice(@Nonnull final ValidationContext context) {
-      if (advice.isVoidReturn()) {
-        context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID);
-      }
-      if (findReturn() == null) {
-        context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN);
+      if (shouldNotUseReturn(pointcut)) {
+        if (!advice.isVoidReturn()) {
+          context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID);
+        }
+        if (findReturn() != null) {
+          context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN);
+        }
+      } else {
+        if (advice.isVoidReturn()) {
+          context.addError(ErrorCode.ADVICE_AFTER_SHOULD_NOT_RETURN_VOID);
+        }
+        if (findReturn() == null) {
+          context.addError(ErrorCode.ADVICE_AFTER_SHOULD_HAVE_RETURN);
+        }
       }
       if (!pointcut.isConstructor()) {
         if (!isStaticPointcut() && !includeThis()) {
@@ -588,6 +597,10 @@ protected void validateAdvice(@Nonnull final ValidationContext context) {
       super.validateAdvice(context);
     }
 
+    private boolean shouldNotUseReturn(final MethodType type) {
+      return !type.isConstructor() && type.isVoidReturn();
+    }
+
     @Override
     public String toString() {
       return "@CallSite.After(" + signature + ")";

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

@@ -359,6 +359,20 @@ public String apply(final Object[] objects) {
     }
   },
 
+  ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID {
+    @Override
+    public String apply(final Object[] objects) {
+      return "After advice for void method should return void";
+    }
+  },
+
+  ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN {
+    @Override
+    public String apply(final Object[] objects) {
+      return "After advice for void method should not contain @Return annotated parameters";
+    }
+  },
+
   ADVICE_AFTER_SHOULD_HAVE_RETURN {
     @Override
     public String apply(final Object[] objects) {

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

@@ -473,6 +473,36 @@ final class AdviceGeneratorTest extends BaseCsiPluginTest {
     }
   }
 
+
+  @CallSite(spi = CallSites)
+  class AfterAdviceWithVoidReturn {
+    @CallSite.After("void java.lang.StringBuilder.setLength(int)")
+    static void after(@CallSite.This StringBuilder self, @CallSite.Argument(0) int length) {
+    }
+  }
+
+  void 'test after advice with void return'() {
+    setup:
+    final spec = buildClassSpecification(AfterAdviceWithVoidReturn)
+    final generator = buildAdviceGenerator(buildDir)
+
+    when:
+    final result = generator.generate(spec)
+
+    then:
+    assertNoErrors result
+    assertCallSites(result.file) {
+      advices(0) {
+        pointcut('java/lang/StringBuilder', 'setLength', '(I)V')
+        statements(
+          'handler.dupInvoke(owner, descriptor, StackDupMode.COPY);',
+          'handler.method(opcode, owner, name, descriptor, isInterface);',
+          'handler.advice("datadog/trace/plugin/csi/impl/AdviceGeneratorTest$AfterAdviceWithVoidReturn", "after", "(Ljava/lang/StringBuilder;I)V");',
+        )
+      }
+    }
+  }
+
   private static AdviceGenerator buildAdviceGenerator(final File targetFolder) {
     return new AdviceGeneratorImpl(targetFolder, pointcutParser())
   }

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

@@ -542,4 +542,26 @@ class AdviceSpecificationTest extends BaseCsiPluginTest {
     then:
     0 * context.addError(_, _)
   }
+
+
+  @CallSite(spi = CallSites)
+  class AfterWithVoidWrongAdvice {
+    @CallSite.After("void java.lang.String.getChars(int, int, char[], int)")
+    static String after(@CallSite.AllArguments final Object[] args, @CallSite.Return final String result) {
+      return result;
+    }
+  }
+
+  void 'test after advice with void should not use @Return'() {
+    setup:
+    final context = mockValidationContext()
+    final spec = buildClassSpecification(AfterWithVoidWrongAdvice)
+
+    when:
+    spec.advices.each { it.validate(context) }
+
+    then:
+    1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_RETURN_VOID, _)
+    1 * context.addError(ErrorCode.ADVICE_AFTER_VOID_METHOD_SHOULD_NOT_HAVE_RETURN, _)
+  }
 }

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/BaseCallSiteTest.groovy

@@ -83,6 +83,10 @@ class BaseCallSiteTest extends DDSpecification {
     return buildPointcut(String.getDeclaredMethod('concat', String))
   }
 
+  protected static Pointcut stringBuilderSetLengthPointcut() {
+    return buildPointcut(StringBuilder.getDeclaredMethod('setLength', int))
+  }
+
   protected static Pointcut stringReaderPointcut() {
     return buildPointcut(StringReader.getDeclaredConstructor(String))
   }

dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/tooling/csi/CallSiteTransformerTest.groovy

@@ -10,6 +10,7 @@ import net.bytebuddy.dynamic.DynamicType
 import net.bytebuddy.jar.asm.Opcodes
 import net.bytebuddy.jar.asm.Type
 
+import java.util.function.BiConsumer
 import java.util.function.BiFunction
 import java.util.function.Consumer
 
@@ -227,6 +228,38 @@ class CallSiteTransformerTest extends BaseCallSiteTest {
     helperCLass != null
   }
 
+  void 'test after advice with void return'() {
+    setup:
+    final source = Type.getType(StringBuilderSetLengthExample)
+    final target = renameType(source, 'Test')
+    final pointcut = stringBuilderSetLengthPointcut()
+    final advice = Mock(InvokeAdvice)
+    final callSite = Type.getType(StringBuilderSetLengthCallSite)
+    final callSiteTransformer = new CallSiteTransformer(mockAdvices([mockCallSites(advice, pointcut)]))
+    final sb = new StringBuilder("Hello World!")
+    final int length = 5
+    StringBuilderSetLengthCallSite.LAST_CALL = null
+
+    when:
+    final transformedClass = transformType(source, target, callSiteTransformer)
+    final instance = loadType(target, transformedClass) as BiConsumer<StringBuilder, Integer>
+    instance.accept(sb, length)
+
+    then:
+    1 * advice.apply(_ as MethodHandler, Opcodes.INVOKEVIRTUAL, pointcut.type, pointcut.method, pointcut.descriptor, false) >> { params ->
+      final args = params as Object[]
+      final handler = args[0] as MethodHandler
+      final owner = args[2] as String
+      final descriptor = args[4] as String
+      handler.dupInvoke(owner, descriptor, COPY)
+      handler.method(args[1] as int, owner, args[3] as String, descriptor, args[5] as Boolean)
+      handler.advice(callSite.getInternalName(), "after", "(Ljava/lang/StringBuilder;I)V")
+    }
+    sb.toString() == 'Hello'
+    StringBuilderSetLengthCallSite.LAST_CALL[0] == sb
+    StringBuilderSetLengthCallSite.LAST_CALL[1] == length
+  }
+
   static class InstrumentationHelper {
 
     private static Consumer<Object[]> callback = null // codenarc forces the lowercase name

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthCallSite.java

@@ -0,0 +1,10 @@
+package datadog.trace.agent.tooling.csi;
+
+public class StringBuilderSetLengthCallSite {
+
+  public static volatile Object[] LAST_CALL;
+
+  public static void after(final StringBuilder builder, int length) {
+    LAST_CALL = new Object[] {builder, length};
+  }
+}

dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/tooling/csi/StringBuilderSetLengthExample.java

@@ -0,0 +1,16 @@
+package datadog.trace.agent.tooling.csi;
+
+import java.util.function.BiConsumer;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class StringBuilderSetLengthExample implements BiConsumer<StringBuilder, Integer> {
+
+  private static final Logger LOGGER = LoggerFactory.getLogger(StringBuilderSetLengthExample.class);
+
+  public void accept(final StringBuilder builder, final Integer length) {
+    LOGGER.debug("Before apply {} {}", builder, length);
+    builder.setLength(length);
+    LOGGER.debug("After apply {} {}", builder, length);
+  }
+}