Full thread stack dump on timeout with custom handler

src/main/java/org/junit/internal/runners/statements/FailOnTimeout.java

@@ -2,13 +2,16 @@
 
 import java.lang.management.ManagementFactory;
 import java.lang.management.ThreadMXBean;
-import java.util.Arrays;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.concurrent.Callable;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+import org.junit.rules.Timeout;
+import org.junit.rules.TimeoutHandler;
 import org.junit.runners.model.MultipleFailureException;
 import org.junit.runners.model.Statement;
 
@@ -17,21 +20,23 @@ public class FailOnTimeout extends Statement {
     private final TimeUnit fTimeUnit;
     private final long fTimeout;
     private final boolean fLookForStuckThread;
+    private final TimeoutHandler fTimeoutHandler;
     private ThreadGroup fThreadGroup = null;
 
     public FailOnTimeout(Statement originalStatement, long millis) {
         this(originalStatement, millis, TimeUnit.MILLISECONDS);
     }
 
     public FailOnTimeout(Statement originalStatement, long timeout, TimeUnit unit) {
-        this(originalStatement, timeout, unit, false);
+        this(originalStatement, timeout, unit, false, null);
     }
 
-    public FailOnTimeout(Statement originalStatement, long timeout, TimeUnit unit, boolean lookForStuckThread) {
+    public FailOnTimeout(Statement originalStatement, long timeout, TimeUnit unit, boolean lookForStuckThread, TimeoutHandler timeoutHandler) {
         fOriginalStatement = originalStatement;
         fTimeout = timeout;
         fTimeUnit = unit;
         fLookForStuckThread = lookForStuckThread;
+        fTimeoutHandler = timeoutHandler;
     }
 
     @Override
@@ -72,24 +77,78 @@ private Exception createTimeoutException(Thread thread) {
                 "test timed out after %d %s", fTimeout, fTimeUnit.name().toLowerCase()));
         if (stackTrace != null) {
             currThreadException.setStackTrace(stackTrace);
-            thread.interrupt();
         }
+
+        List<Throwable> exceptions = new ArrayList<Throwable>();
+        exceptions.add(currThreadException);
+
         if (stuckThread != null) {
-            Exception stuckThreadException = 
+            Exception stuckThreadException =
                 new Exception ("Appears to be stuck in thread " +
                                stuckThread.getName());
             stuckThreadException.setStackTrace(getStackTrace(stuckThread));
-            return new MultipleFailureException    
-                (Arrays.<Throwable>asList(currThreadException, stuckThreadException));
+            exceptions.add(stuckThreadException);
+        }
+
+        // Specifically set timeout handler takes precedence
+        TimeoutHandler timeoutHandler = fTimeoutHandler;
+        if (timeoutHandler == null) {
+            timeoutHandler = lookForGlobalTimeoutHandler();
+        }
+        if (timeoutHandler != null) {
+            // For the sake of convenience just allow adding another exception by the custom timeout handler
+            Exception timeoutHandlerException = timeoutHandler.handleTimeout(thread);
+            if (timeoutHandlerException != null) {
+                exceptions.add(timeoutHandlerException);
+            }
+        }
+
+        thread.interrupt();
+
+        if (exceptions.size() > 1) {
+            return new MultipleFailureException(exceptions);
         } else {
             return currThreadException;
         }
     }
 
+    /**
+     * Note: Before this was in {@link Timeout} and called from
+     * {@link Timeout#apply(Statement, org.junit.runner.Description)}, but this would not have
+     * been used for timeouts specified for single test cases.
+     */
+    private TimeoutHandler lookForGlobalTimeoutHandler() {
+        //TODO: due to using reflection, the global timeout handler should be cached once it is
+        // there, and used as long as the system property is still set?
+        TimeoutHandler handler = null;
+        String timeoutHandlerClassName = System.getProperty(Timeout.TIMEOUT_HANDLER_CLASS_NAME_PROPERTY_NAME);
+        if (timeoutHandlerClassName != null) {
+            try {
+System.err.println("FOUND");
+                Class<?> timeoutHandlerClass = Class.forName(timeoutHandlerClassName);
+                Object handlerInstance = timeoutHandlerClass.newInstance();
+                handler = TimeoutHandler.class.cast(handlerInstance);
+            } catch (ClassNotFoundException e) {
+                //TODO: or throw InitializationError?
+                throw new RuntimeException(String.format("Failed to find timeout handler class '%s' specified via system property", timeoutHandlerClassName), e);
+            } catch (InstantiationException e) {
+                //TODO: or throw InitializationError?
+                throw new RuntimeException(String.format("Failed to instantiate timeout handler class '%s'", timeoutHandlerClassName), e);
+            } catch (IllegalAccessException e) {
+                //TODO: or throw InitializationError?
+                throw new RuntimeException(String.format("Failed to access timeout handler class '%s' during instantion", timeoutHandlerClassName), e);
+            } catch (ClassCastException e) {
+                //TODO: or throw InitializationError?
+                throw new RuntimeException(String.format("Failed to cast timeout handler class '%s' to '%s'", timeoutHandlerClassName, TimeoutHandler.class), e);
+            }
+        }
+        return handler;
+    }
+
     /**
      * Retrieves the stack trace for a given thread.
      * @param thread The thread whose stack is to be retrieved.
-     * @return The stack trace; returns a zero-length array if the thread has 
+     * @return The stack trace; returns a zero-length array if the thread has
      * terminated or the stack cannot be retrieved for some other reason.
      */
     private StackTraceElement[] getStackTrace(Thread thread) {
@@ -107,18 +166,18 @@ private StackTraceElement[] getStackTrace(Thread thread) {
      * @param mainThread The main thread created by {@code evaluate()}
      * @return The thread which appears to be causing the problem, if different from
      * {@code mainThread}, or {@code null} if the main thread appears to be the
-     * problem or if the thread cannot be determined.  The return value is never equal 
+     * problem or if the thread cannot be determined.  The return value is never equal
      * to {@code mainThread}.
      */
     private Thread getStuckThread (Thread mainThread) {
-        if (fThreadGroup == null) 
+        if (fThreadGroup == null)
             return null;
         Thread[] threadsInGroup = getThreadArray(fThreadGroup);
-        if (threadsInGroup == null) 
+        if (threadsInGroup == null)
             return null;
-        
+
         // Now that we have all the threads in the test's thread group: Assume that
-        // any thread we're "stuck" in is RUNNABLE.  Look for all RUNNABLE threads. 
+        // any thread we're "stuck" in is RUNNABLE.  Look for all RUNNABLE threads.
         // If just one, we return that (unless it equals threadMain).  If there's more
         // than one, pick the one that's using the most CPU time, if this feature is
         // supported.
@@ -131,13 +190,13 @@ private Thread getStuckThread (Thread mainThread) {
                     stuckThread = thread;
                     maxCpuTime = threadCpuTime;
                 }
-            }               
+            }
         }
         return (stuckThread == mainThread) ? null : stuckThread;
     }
 
     /**
-     * Returns all active threads belonging to a thread group.  
+     * Returns all active threads belonging to a thread group.
      * @param group The thread group.
      * @return The active threads in the thread group.  The result should be a
      * complete list of the active threads at some point in time.  Returns {@code null}
@@ -158,16 +217,16 @@ private Thread[] getThreadArray(ThreadGroup group) {
             // is >= the array's length; therefore we can't trust that it returned all
             // the threads.  Try again.
             enumSize += 100;
-            if (++loopCount >= 5) 
+            if (++loopCount >= 5)
                 return null;
-            // threads are proliferating too fast for us.  Bail before we get into 
+            // threads are proliferating too fast for us.  Bail before we get into
             // trouble.
         }
         return copyThreads(threads, enumCount);
     }
 
     /**
-     * Returns an array of the first {@code count} Threads in {@code threads}. 
+     * Returns an array of the first {@code count} Threads in {@code threads}.
      * (Use instead of Arrays.copyOf to maintain compatibility with Java 1.5.)
      * @param threads The source array.
      * @param count The maximum length of the result array.

src/main/java/org/junit/rules/Timeout.java

@@ -34,9 +34,13 @@
  * @since 4.7
  */
 public class Timeout implements TestRule {
+
+    public static final String TIMEOUT_HANDLER_CLASS_NAME_PROPERTY_NAME = "org.junit.rules.timeout.handler";
+
     private final long fTimeout;
     private final TimeUnit fTimeUnit;
     private boolean fLookForStuckThread;
+    private TimeoutHandler fTimeoutHandler;
 
     /**
      * Create a {@code Timeout} instance with the timeout specified
@@ -68,6 +72,7 @@ public Timeout(long timeout, TimeUnit unit) {
         fTimeout = timeout;
         fTimeUnit = unit;
         fLookForStuckThread = false;
+        fTimeoutHandler = null;
     }
 
     /**
@@ -100,7 +105,12 @@ public Timeout lookForStuckThread(boolean enable) {
         return this;
     }
 
+    public Timeout customTimeoutHandler(TimeoutHandler handler) {
+        fTimeoutHandler = handler;
+        return this;
+    }
+
     public Statement apply(Statement base, Description description) {
-        return new FailOnTimeout(base, fTimeout, fTimeUnit, fLookForStuckThread);
+        return new FailOnTimeout(base, fTimeout, fTimeUnit, fLookForStuckThread, fTimeoutHandler);
     }
 }

src/main/java/org/junit/rules/TimeoutHandler.java

@@ -0,0 +1,14 @@
+package org.junit.rules;
+
+public interface TimeoutHandler {
+
+    /**
+     * If the test runs into test timeout, this callback will be called allowing to gather
+     * additional test failure context (e.g. of other involved systems).
+     *
+     * @param thread the actual test executing main thread
+     * @return the optional exception will be added to the test failure
+     */
+    Exception handleTimeout(Thread thread);
+
+}