Removes CountDownLatch

fork-runner/src/main/java/com/shazam/fork/ForkRunner.java

@@ -36,6 +36,7 @@
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 
+import static com.google.common.util.concurrent.Uninterruptibles.awaitTerminationUninterruptibly;
 import static com.shazam.fork.Utils.namedExecutor;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
@@ -113,28 +114,34 @@ public boolean run() {
     private void executeTests(
             Collection<Pool> pools,
             Collection<TestCaseEvent> testCases
-    ) throws InterruptedException {
+    ) {
         ProgressReporter progressReporter = progressReporterFactory.createProgressReporter();
         progressReporter.start();
 
-        ExecutorService poolExecutor = namedExecutor(pools.size(), "PoolExecutor-%d");
-
-        CountDownLatch poolCountDownLatch = new CountDownLatch(pools.size());
-        for (Pool pool : pools) {
-            Runnable poolTestRunner = poolTestRunnerFactory.createPoolTestRunner(
-                    pool,
-                    testCases,
-                    poolCountDownLatch,
-                    progressReporter
-            );
-            poolExecutor.execute(poolTestRunner);
-        }
-        poolCountDownLatch.await();
+        ExecutorService poolExecutor = null;
+        try {
+            poolExecutor = namedExecutor(pools.size(), "PoolExecutor-%d");
+
+            for (Pool pool : pools) {
+                Runnable poolTestRunner = poolTestRunnerFactory.createPoolTestRunner(
+                        pool,
+                        testCases,
+                        progressReporter
+                );
+                poolExecutor.submit(poolTestRunner);
+            }
 
-        poolExecutor.shutdown();
-        poolExecutor.awaitTermination(10, TimeUnit.MINUTES);
+            poolExecutor.shutdown();
+            awaitTerminationUninterruptibly(poolExecutor);
+        } finally {
+            progressReporter.stop();
+
+            if (poolExecutor != null && !poolExecutor.isTerminated()) {
+                poolExecutor.shutdownNow();
+                awaitTerminationUninterruptibly(poolExecutor);
+            }
+        }
 
-        progressReporter.stop();
     }
 
     private static void reportMissingTests(AggregatedTestResult aggregatedTestResult) {

fork-runner/src/main/java/com/shazam/fork/runner/DeviceTestRunner.java

@@ -23,7 +23,6 @@
 
 import java.io.IOException;
 import java.util.Queue;
-import java.util.concurrent.CountDownLatch;
 
 import static com.shazam.fork.system.io.RemoteFileManager.*;
 
@@ -34,7 +33,6 @@ public class DeviceTestRunner implements Runnable {
     private final Pool pool;
     private final Device device;
     private final Queue<TestCaseEvent> queueOfTestsInPool;
-    private final CountDownLatch deviceCountDownLatch;
     private final ProgressReporter progressReporter;
     private final ScreenRecorder screenRecorder;
     private final TestRunFactory testRunFactory;
@@ -43,15 +41,13 @@ public DeviceTestRunner(Installer installer,
                             Pool pool,
                             Device device,
                             Queue<TestCaseEvent> queueOfTestsInPool,
-                            CountDownLatch deviceCountDownLatch,
                             ProgressReporter progressReporter,
                             ScreenRecorder screenRecorder,
                             TestRunFactory testRunFactory) {
         this.installer = installer;
         this.pool = pool;
         this.device = device;
         this.queueOfTestsInPool = queueOfTestsInPool;
-        this.deviceCountDownLatch = deviceCountDownLatch;
         this.progressReporter = progressReporter;
         this.screenRecorder = screenRecorder;
         this.testRunFactory = testRunFactory;
@@ -81,7 +77,6 @@ public void run() {
             }
         } finally {
             logger.info("Device {} from pool {} finished", device.getSerial(), pool.getName());
-            deviceCountDownLatch.countDown();
         }
     }
 

fork-runner/src/main/java/com/shazam/fork/runner/DeviceTestRunnerFactory.java

@@ -17,7 +17,6 @@
 import com.shazam.fork.system.adb.Installer;
 
 import java.util.Queue;
-import java.util.concurrent.CountDownLatch;
 
 public class DeviceTestRunnerFactory {
 
@@ -31,7 +30,6 @@ public DeviceTestRunnerFactory(Installer installer, TestRunFactory testRunFactor
 
     public Runnable createDeviceTestRunner(Pool pool,
                                            Queue<TestCaseEvent> testClassQueue,
-                                           CountDownLatch deviceInPoolCountDownLatch,
                                            Device device,
                                            ProgressReporter progressReporter
     ) {
@@ -40,7 +38,6 @@ public Runnable createDeviceTestRunner(Pool pool,
                 pool,
                 device,
                 testClassQueue,
-                deviceInPoolCountDownLatch,
                 progressReporter,
                 new ScreenRecorderImpl(device),
                 testRunFactory);

fork-runner/src/main/java/com/shazam/fork/runner/PoolTestRunner.java

@@ -20,9 +20,9 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.Queue;
-import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 
+import static com.google.common.util.concurrent.Uninterruptibles.awaitTerminationUninterruptibly;
 import static com.shazam.fork.Utils.namedExecutor;
 
 public class PoolTestRunner implements Runnable {
@@ -31,17 +31,14 @@ public class PoolTestRunner implements Runnable {
 
     private final Pool pool;
     private final Queue<TestCaseEvent> testCases;
-    private final CountDownLatch poolCountDownLatch;
     private final DeviceTestRunnerFactory deviceTestRunnerFactory;
     private final ProgressReporter progressReporter;
 
     public PoolTestRunner(DeviceTestRunnerFactory deviceTestRunnerFactory, Pool pool,
                           Queue<TestCaseEvent> testCases,
-                          CountDownLatch poolCountDownLatch,
                           ProgressReporter progressReporter) {
         this.pool = pool;
         this.testCases = testCases;
-        this.poolCountDownLatch = poolCountDownLatch;
         this.deviceTestRunnerFactory = deviceTestRunnerFactory;
         this.progressReporter = progressReporter;
     }
@@ -52,23 +49,25 @@ public void run() {
         try {
             int devicesInPool = pool.size();
             concurrentDeviceExecutor = namedExecutor(devicesInPool, "DeviceExecutor-%d");
-            CountDownLatch deviceCountDownLatch = new CountDownLatch(devicesInPool);
             logger.info("Pool {} started", poolName);
             for (Device device : pool.getDevices()) {
-                Runnable deviceTestRunner = deviceTestRunnerFactory.createDeviceTestRunner(pool, testCases,
-                        deviceCountDownLatch, device, progressReporter);
-                concurrentDeviceExecutor.execute(deviceTestRunner);
+                Runnable deviceTestRunner = deviceTestRunnerFactory.createDeviceTestRunner(
+                        pool,
+                        testCases,
+                        device,
+                        progressReporter
+                );
+                concurrentDeviceExecutor.submit(deviceTestRunner);
             }
-            deviceCountDownLatch.await();
-        } catch (InterruptedException e) {
-            logger.warn("Pool {} was interrupted while running", poolName);
+
+            concurrentDeviceExecutor.shutdown();
+            awaitTerminationUninterruptibly(concurrentDeviceExecutor);
         } finally {
-            if (concurrentDeviceExecutor != null) {
-                concurrentDeviceExecutor.shutdown();
+            if (concurrentDeviceExecutor != null && !concurrentDeviceExecutor.isTerminated()) {
+                concurrentDeviceExecutor.shutdownNow();
+                awaitTerminationUninterruptibly(concurrentDeviceExecutor);
             }
             logger.info("Pool {} finished", poolName);
-            poolCountDownLatch.countDown();
-            logger.info("Pools remaining: {}", poolCountDownLatch.getCount());
         }
     }
 }

fork-runner/src/main/java/com/shazam/fork/runner/PoolTestRunnerFactory.java

@@ -15,7 +15,6 @@
 
 import java.util.Collection;
 import java.util.LinkedList;
-import java.util.concurrent.CountDownLatch;
 
 public class PoolTestRunnerFactory {
     private final DeviceTestRunnerFactory deviceTestRunnerFactory;
@@ -26,7 +25,6 @@ public PoolTestRunnerFactory(DeviceTestRunnerFactory deviceTestRunnerFactory) {
 
     public Runnable createPoolTestRunner(Pool pool,
                                          Collection<TestCaseEvent> testCases,
-                                         CountDownLatch poolCountDownLatch,
                                          ProgressReporter progressReporter) {
 
         int totalTests = testCases.size();
@@ -36,7 +34,6 @@ public Runnable createPoolTestRunner(Pool pool,
                 deviceTestRunnerFactory,
                 pool,
                 new LinkedList<>(testCases),
-                poolCountDownLatch,
                 progressReporter);
     }
 }

fork-runner/src/main/java/com/shazam/fork/runner/TestRun.java

@@ -26,6 +26,7 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.concurrent.TimeUnit;
 
 import static java.lang.String.format;
 
@@ -64,7 +65,7 @@ public void execute() {
 		}
 		runner.setRunName(poolName);
 		runner.setMethodName(testClassName, testMethodName);
-		runner.setMaxtimeToOutputResponse(testRunParameters.getTestOutputTimeout());
+		runner.setMaxTimeToOutputResponse(testRunParameters.getTestOutputTimeout(), TimeUnit.MILLISECONDS);
 
         if (testRunParameters.isCoverageEnabled()) {
             runner.setCoverage(true);