How to dump all thread stacks on junit timeout -- making a Junit 5 extension

user_1366016 0
Viewed 68

I am running tests with junit 5, and occasionally, one will hang. Usually this is in a driver to a cloud service, but we haven't caught it in the act so we don't know exactly where. I'd like to configure junit to output a stack trace of all threads when the timeout happens. But I don't see how that is easily done. Sure, I can build my own timeout mechanism and install it in @BeforeEach and cancel it in @AfterEach, and it is easy to print the thread stack dump, but it would be nice if junit had something built-in like an "execute this lambda on timeout". I don't see anything like that in the assertTimeout*() functions.

UPDATE: A Junit5 extension turned out to be the best solution I cound find, see my answer.

junit5 java
1 Answers

A lot of the comments were hinting at a roll-your-own solution, which ultimately I had to do. It really wasn't hard, because a colleague of mine had written Junit 5 extensions and I was able to use his expertise, which I am happy to share with the community.

This extension involves two things:

  • The extension class
  • An annotation that can be added to tests

Once you have these two things in your classpath for the test, you can annotate your test class (or a base class of your test class) like:

@ExtendWith({TimeoutExtension.class})
public class Test1Timeout {

This enables the extension. Then annotate either a test class or individual test methods with

@Timeout(seconds=100)

which will set the maximum time allowed for a test method to run. If the test method fails to complete in the allotted time, two things will happen:

  • A dump will be printed to System.out of the stack trace of all running threads
  • It will call Thread.interrupt() on the thread which ran the test method, which will tend to dislodge most "waiting for I/O" and "waiting for synchronization" cases.

The extension class:

public class TimeoutExtension implements InvocationInterceptor {
  private static final ScheduledExecutorService exec = Executors.newScheduledThreadPool(1);
  @Override
  public void interceptTestMethod(Invocation<Void> invocation, ReflectiveInvocationContext<Method> invocationContext, ExtensionContext extensionContext) throws Throwable {
    Timeout timeout = invocationContext.getExecutable().getAnnotation(Timeout.class);
    Class<?> clazz = invocationContext.getExecutable().getDeclaringClass();
    while (timeout == null && clazz != Object.class) {
      timeout = clazz.getAnnotation(Timeout.class);
      clazz = clazz.getSuperclass();
    }
    if (timeout == null || timeout.seconds() <= 0) {
      invocation.proceed();
      return;
    }
    int seconds = timeout.seconds();
    Thread caller = Thread.currentThread();
    AtomicBoolean timedOut = new AtomicBoolean();
    Future<Void> future = exec.schedule(() -> {
      System.out.println("**** TIMEOUT ERROR: TEST EXCEEDED " + seconds + " SECONDS ****");
      printThreadDump();
      timedOut.set(true);
      caller.interrupt();
      return null;
    }, seconds, TimeUnit.SECONDS);
    Exception caught = null;
    try {
      invocation.proceed();
    } catch (Exception ex) {
      caught = ex;
    } finally {
      future.cancel(true);
      if (timedOut.get()) {
        if (timeout.expectTimeout()) {
          // awesome!
        } else {
          Exception ex = new TimeoutException("Test exceeded timeout of " + seconds + " seconds");
          if (caught != null) {
            ex.addSuppressed(caught);
          }
          throw ex;
        }
      } else if (caught != null) {
        throw caught;
      } else if (timeout.expectTimeout()) {
        throw new RuntimeException("Test expected to timeout at " + seconds + " but didn't");
      }
    }
  }
}

The Timeout annotation

@Retention(RetentionPolicy.RUNTIME)
@Target( { ElementType.METHOD, ElementType.TYPE })
public @interface Timeout {
  /** Set to zero to disable timeout **/
  int seconds();
  /** Set to true if the test is expected to timeout and this is OK with you. 
      This is mostly for self-testing the extension.
  **/
  boolean expectTimeout() default false;
}

and finally, the thread-dumper utility

public class TestUtils {
  public static String generateThreadDump() {
    final StringBuilder dump = new StringBuilder();
    final ThreadMXBean threadMXBean = ManagementFactory.getThreadMXBean();
    final ThreadInfo[] threadInfos = threadMXBean.getThreadInfo(threadMXBean.getAllThreadIds(), 100);
    for (ThreadInfo threadInfo : threadInfos) {
      dump.append('"');
      dump.append(threadInfo.getThreadName());
      dump.append("\" ");
      final Thread.State state = threadInfo.getThreadState();
      dump.append("\n   java.lang.Thread.State: ");
      dump.append(state);
      final StackTraceElement[] stackTraceElements = threadInfo.getStackTrace();
      for (final StackTraceElement stackTraceElement : stackTraceElements) {
        dump.append("\n        at ");
        dump.append(stackTraceElement);
      }
      dump.append("\n\n");
    }
    return dump.toString();
  }
  public static void printThreadDump() {
    System.out.println(generateThreadDump());
  }
}
user_13660160
Related