Home | History | Annotate | Download | only in junit 
      1 /*
      2  * Copyright (C) 2016 The Android Open Source Project
      3  *
      4  * Licensed under the Apache License, Version 2.0 (the "License");
      5  * you may not use this file except in compliance with the License.
      6  * You may obtain a copy of the License at
      7  *
      8  *      http://www.apache.org/licenses/LICENSE-2.0
      9  *
     10  * Unless required by applicable law or agreed to in writing, software
     11  * distributed under the License is distributed on an "AS IS" BASIS,
     12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     13  * See the License for the specific language governing permissions and
     14  * limitations under the License.
     15  */
     16 
     17 package vogar.target.junit;
     18 
     19 import java.util.concurrent.Callable;
     20 import java.util.concurrent.ExecutionException;
     21 import java.util.concurrent.ExecutorService;
     22 import java.util.concurrent.Executors;
     23 import java.util.concurrent.Future;
     24 import java.util.concurrent.TimeUnit;
     25 import java.util.concurrent.TimeoutException;
     26 import java.util.concurrent.atomic.AtomicReference;
     27 import org.junit.rules.TestRule;
     28 import org.junit.runner.Description;
     29 import org.junit.runners.model.Statement;
     30 import vogar.util.Threads;
     31 
     32 /**
     33  * Times a test out and then aborts the test run.
     34  */
     35 public class TimeoutAndAbortRunRule implements TestRule {
     36 
     37     private final ExecutorService executor = Executors.newCachedThreadPool(
     38             Threads.daemonThreadFactory(getClass().getName()));
     39 
     40     private final int timeoutSeconds;
     41 
     42     /**
     43      * @param timeoutSeconds the timeout in seconds, if 0 then never times out.
     44      */
     45     public TimeoutAndAbortRunRule(int timeoutSeconds) {
     46         this.timeoutSeconds = timeoutSeconds;
     47     }
     48 
     49     @Override
     50     public Statement apply(final Statement base, Description description) {
     51         return new Statement() {
     52             @Override
     53             public void evaluate() throws Throwable {
     54                 runWithTimeout(base);
     55             }
     56         };
     57     }
     58 
     59     /**
     60      * Runs the test on another thread. If the test completes before the
     61      * timeout, this reports the result normally. But if the test times out,
     62      * this reports the timeout stack trace and begins the process of killing
     63      * this no-longer-trustworthy process.
     64      */
     65     private void runWithTimeout(final Statement base) throws Throwable {
     66         // Start the test on a background thread.
     67         final AtomicReference<Thread> executingThreadReference = new AtomicReference<>();
     68         Future<Throwable> result = executor.submit(new Callable<Throwable>() {
     69             public Throwable call() throws Exception {
     70                 executingThreadReference.set(Thread.currentThread());
     71                 try {
     72                     base.evaluate();
     73                     return null;
     74                 } catch (Throwable throwable) {
     75                     return throwable;
     76                 }
     77             }
     78         });
     79 
     80         // Wait until either the result arrives or the test times out.
     81         Throwable thrown;
     82         try {
     83             thrown = getThrowable(result);
     84         } catch (TimeoutException e) {
     85             Thread executingThread = executingThreadReference.get();
     86             if (executingThread != null) {
     87                 executingThread.interrupt();
     88                 e.setStackTrace(executingThread.getStackTrace());
     89             }
     90             // Wrap it in an exception that will cause the current run to be aborted.
     91             thrown = new VmIsUnstableException(e);
     92         }
     93 
     94         if (thrown != null) {
     95             throw thrown;
     96         }
     97     }
     98 
     99     @SuppressWarnings("ThrowableResultOfMethodCallIgnored")
    100     private Throwable getThrowable(Future<Throwable> result)
    101             throws InterruptedException, ExecutionException, TimeoutException {
    102         Throwable thrown;
    103         thrown = timeoutSeconds == 0
    104                 ? result.get()
    105                 : result.get(timeoutSeconds, TimeUnit.SECONDS);
    106         return thrown;
    107     }
    108 }
    109 
    110