Set max stack size in bytebuddy exception handler

Andrew Kent · Andrew Kent · commit 5ffd2142e938 · 2018-03-07T13:14:41.000-05:00
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ExceptionHandlers.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/ExceptionHandlers.java
@@ -3,7 +3,6 @@
 import datadog.trace.agent.bootstrap.ExceptionLogger;
 import net.bytebuddy.implementation.Implementation;
 import net.bytebuddy.implementation.bytecode.StackManipulation;
-import net.bytebuddy.implementation.bytecode.StackSize;
 import net.bytebuddy.jar.asm.Label;
 import net.bytebuddy.jar.asm.MethodVisitor;
 import net.bytebuddy.jar.asm.Opcodes;
@@ -19,7 +18,8 @@ public class ExceptionHandlers {
 
   private static final StackManipulation EXCEPTION_STACK_HANDLER =
       new StackManipulation() {
-        private final Size size = StackSize.SINGLE.toDecreasingSize();
+        // Pops one Throwable off the stack. Maxes the stack to at least 3.
+        private final Size size = new StackManipulation.Size(-1, 3);
 
         @Override
         public boolean isValid() {
diff --git a/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ExceptionHandlerTest.groovy b/dd-java-agent/agent-tooling/src/test/groovy/datadog/trace/agent/test/ExceptionHandlerTest.groovy
@@ -33,6 +33,13 @@ class ExceptionHandlerTest extends Specification {
         .advice(
         isMethod().and(named("isInstrumented")),
         BadAdvice.getName()))
+      .transform(
+      new AgentBuilder.Transformer.ForAdvice()
+        .with(new AgentBuilder.LocationStrategy.Simple(ClassFileLocator.ForClassLoader.of(BadAdvice.getClassLoader())))
+        .withExceptionHandler(ExceptionHandlers.defaultExceptionHandler())
+        .advice(
+        isMethod().and(named("smallStack").or(named("largeStack"))),
+        BadAdvice.NoOpAdvice.getName()))
       .asDecorator()
 
     ByteBuddyAgent.install()
@@ -80,9 +87,31 @@ class ExceptionHandlerTest extends Specification {
     testAppender.list.size() == initLogEvents
   }
 
+  def "exception handler sets the correct stack size"() {
+    when:
+    SomeClass.smallStack()
+    SomeClass.largeStack()
+
+    then:
+    noExceptionThrown()
+  }
+
   static class SomeClass {
     static boolean isInstrumented() {
       return false
     }
+
+    static void smallStack() {
+      // a method with a max stack of 0
+    }
+
+    static void largeStack() {
+      // a method with a max stack of 6
+      long l = 22l
+      int i = 3
+      double d = 32.2d
+      Object o = new Object()
+      println "large stack: $l $i $d $o"
+    }
   }
 }
diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/BadAdvice.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/BadAdvice.java
@@ -8,4 +8,11 @@ public static void throwAnException(@Advice.Return(readOnly = false) boolean ret
     returnVal = true;
     throw new RuntimeException("Test Exception");
   }
+
+  public static class NoOpAdvice {
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void doNothing() {
+      System.currentTimeMillis();
+    }
+  }
 }
diff --git a/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/HelperClass.java b/dd-java-agent/agent-tooling/src/test/java/datadog/trace/agent/test/HelperClass.java
diff --git a/dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanTest.groovy b/dd-trace-ot/src/test/groovy/datadog/opentracing/DDSpanTest.groovy
@@ -7,7 +7,7 @@ import spock.lang.Timeout
 
 import java.util.concurrent.TimeUnit
 
-@Timeout(1)
+@Timeout(10)
 class DDSpanTest extends Specification {
   def writer = new ListWriter()
   def tracer = new DDTracer(writer)

Original file line number	Diff line number	Diff line change
`@@ -8,4 +8,11 @@ public static void throwAnException(@Advice.Return(readOnly = false) boolean ret`
`8`	`8`	`returnVal = true;`
`9`	`9`	`throw new RuntimeException("Test Exception");`
`10`	`10`	`}`
	`11`	`+`
	`12`	`+ public static class NoOpAdvice {`
	`13`	`+ @Advice.OnMethodExit(suppress = Throwable.class)`
	`14`	`+ public static void doNothing() {`
	`15`	`+ System.currentTimeMillis();`
	`16`	`+ }`
	`17`	`+ }`
`11`	`18`	`}`