1 package org.junit.internal.runners.statements; 2 3 import java.util.concurrent.Callable; 4 import java.util.concurrent.CountDownLatch; 5 import java.util.concurrent.ExecutionException; 6 import java.util.concurrent.FutureTask; 7 import java.util.concurrent.TimeUnit; 8 import java.util.concurrent.TimeoutException; 9 10 import org.junit.runners.model.Statement; 11 import org.junit.runners.model.TestTimedOutException; 12 13 public class FailOnTimeout extends Statement { 14 private final Statement originalStatement; 15 private final TimeUnit timeUnit; 16 private final long timeout; 17 18 /** 19 * Returns a new builder for building an instance. 20 * 21 * @since 4.12 22 */ builder()23 public static Builder builder() { 24 return new Builder(); 25 } 26 27 /** 28 * Creates an instance wrapping the given statement with the given timeout in milliseconds. 29 * 30 * @param statement the statement to wrap 31 * @param timeoutMillis the timeout in milliseconds 32 * @deprecated use {@link #builder()} instead. 33 */ 34 @Deprecated FailOnTimeout(Statement statement, long timeoutMillis)35 public FailOnTimeout(Statement statement, long timeoutMillis) { 36 this(builder().withTimeout(timeoutMillis, TimeUnit.MILLISECONDS), statement); 37 } 38 FailOnTimeout(Builder builder, Statement statement)39 private FailOnTimeout(Builder builder, Statement statement) { 40 originalStatement = statement; 41 timeout = builder.timeout; 42 timeUnit = builder.unit; 43 } 44 45 /** 46 * Builder for {@link FailOnTimeout}. 47 * 48 * @since 4.12 49 */ 50 public static class Builder { 51 private long timeout = 0; 52 private TimeUnit unit = TimeUnit.SECONDS; 53 Builder()54 private Builder() { 55 } 56 57 /** 58 * Specifies the time to wait before timing out the test. 59 * 60 * <p>If this is not called, or is called with a {@code timeout} of 61 * {@code 0}, the returned {@code Statement} will wait forever for the 62 * test to complete, however the test will still launch from a separate 63 * thread. This can be useful for disabling timeouts in environments 64 * where they are dynamically set based on some property. 65 * 66 * @param timeout the maximum time to wait 67 * @param unit the time unit of the {@code timeout} argument 68 * @return {@code this} for method chaining. 69 */ withTimeout(long timeout, TimeUnit unit)70 public Builder withTimeout(long timeout, TimeUnit unit) { 71 if (timeout < 0) { 72 throw new IllegalArgumentException("timeout must be non-negative"); 73 } 74 if (unit == null) { 75 throw new NullPointerException("TimeUnit cannot be null"); 76 } 77 this.timeout = timeout; 78 this.unit = unit; 79 return this; 80 } 81 82 /** 83 * Builds a {@link FailOnTimeout} instance using the values in this builder, 84 * wrapping the given statement. 85 * 86 * @param statement 87 */ build(Statement statement)88 public FailOnTimeout build(Statement statement) { 89 if (statement == null) { 90 throw new NullPointerException("statement cannot be null"); 91 } 92 return new FailOnTimeout(this, statement); 93 } 94 } 95 96 @Override evaluate()97 public void evaluate() throws Throwable { 98 CallableStatement callable = new CallableStatement(); 99 FutureTask<Throwable> task = new FutureTask<Throwable>(callable); 100 Thread thread = new Thread(task, "Time-limited test"); 101 thread.setDaemon(true); 102 thread.start(); 103 callable.awaitStarted(); 104 Throwable throwable = getResult(task, thread); 105 if (throwable != null) { 106 throw throwable; 107 } 108 } 109 110 /** 111 * Wait for the test task, returning the exception thrown by the test if the 112 * test failed, an exception indicating a timeout if the test timed out, or 113 * {@code null} if the test passed. 114 */ getResult(FutureTask<Throwable> task, Thread thread)115 private Throwable getResult(FutureTask<Throwable> task, Thread thread) { 116 try { 117 if (timeout > 0) { 118 return task.get(timeout, timeUnit); 119 } else { 120 return task.get(); 121 } 122 } catch (InterruptedException e) { 123 return e; // caller will re-throw; no need to call Thread.interrupt() 124 } catch (ExecutionException e) { 125 // test failed; have caller re-throw the exception thrown by the test 126 return e.getCause(); 127 } catch (TimeoutException e) { 128 return createTimeoutException(thread); 129 } 130 } 131 createTimeoutException(Thread thread)132 private Exception createTimeoutException(Thread thread) { 133 StackTraceElement[] stackTrace = thread.getStackTrace(); 134 Exception currThreadException = new TestTimedOutException(timeout, timeUnit); 135 if (stackTrace != null) { 136 currThreadException.setStackTrace(stackTrace); 137 thread.interrupt(); 138 } 139 return currThreadException; 140 } 141 142 private class CallableStatement implements Callable<Throwable> { 143 private final CountDownLatch startLatch = new CountDownLatch(1); 144 call()145 public Throwable call() throws Exception { 146 try { 147 startLatch.countDown(); 148 originalStatement.evaluate(); 149 } catch (Exception e) { 150 throw e; 151 } catch (Throwable e) { 152 return e; 153 } 154 return null; 155 } 156 awaitStarted()157 public void awaitStarted() throws InterruptedException { 158 startLatch.await(); 159 } 160 } 161 } 162