// Copyright 2015 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

package org.pantsbuild.tools.junit.impl;

import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.CompletionService;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.ThreadFactory;
import org.junit.runners.model.RunnerScheduler;
import org.pantsbuild.junit.annotations.TestParallel;
import org.pantsbuild.junit.annotations.TestSerial;

public class ConcurrentRunnerScheduler implements RunnerScheduler {
  private final CompletionService<Void> completionService;
  private final Queue<Future<Void>> concurrentTasks;
  private final Queue<Runnable> serialTasks;
  private final Concurrency defaultConcurrency;

  /**
   * A concurrent scheduler to run junit tests in parallel if possible, followed by tests that can
   * only be run in serial.
   *
   * Test classes annotated with {@link TestSerial} will be run in serial.
   * Test classes annotated with {@link TestParallel} will be run in parallel.
   * Test classes without neither annotation will be run in parallel if defaultParallel is set.
   *
   * Call {@link org.junit.runners.ParentRunner#setScheduler} to use this scheduler.
   *
   * @param defaultConcurrency  Describes how to parallelize unannotated classes.
   * @param numThreads Number of parallel threads to use, must be positive.
   */
  public ConcurrentRunnerScheduler(Concurrency defaultConcurrency,
      int numThreads) {
    Preconditions.checkNotNull(defaultConcurrency);
    this.defaultConcurrency = defaultConcurrency;
    ThreadFactory threadFactory = new ThreadFactoryBuilder()
        .setDaemon(true)
        .setNameFormat("concurrent-junit-runner-%d")
        .build();
    completionService = new ExecutorCompletionService<Void>(
        Executors.newFixedThreadPool(numThreads, threadFactory));
    concurrentTasks = new LinkedList<Future<Void>>();
    serialTasks = new LinkedList<Runnable>();
  }

  @Override
  public void schedule(Runnable childStatement) {
    if (shouldMethodsRunParallel()) {
      concurrentTasks.offer(completionService.submit(childStatement, null));
    } else {
      serialTasks.offer(childStatement);
    }
  }

  /**
   * Schedule a test childStatement associated with clazz, using clazz's policy to decide running
   * in serial or parallel.
   */
  public void schedule(Runnable childStatement, Class<?> clazz) {
    if (shouldClassRunParallel(clazz)) {
      concurrentTasks.offer(completionService.submit(childStatement, null));
    } else {
      serialTasks.offer(childStatement);
    }
  }

  private boolean shouldMethodsRunParallel() {
    // TODO(zundel): Add support for an annotation like TestParallelMethods.
    return defaultConcurrency.shouldRunMethodsParallel();
  }

  private boolean shouldClassRunParallel(Class<?> clazz) {
    return !clazz.isAnnotationPresent(TestSerial.class)
        && (clazz.isAnnotationPresent(TestParallel.class) ||
        defaultConcurrency.shouldRunClassesParallel());
  }

  @Override
  public void finished() {
    try {
      // Wait for all concurrent tasks to finish.
      while (!concurrentTasks.isEmpty()) {
        concurrentTasks.poll().get();
      }
      // Then run all serial tasks in serial.
      for (Runnable task : serialTasks) {
        task.run();
      }
    } catch (InterruptedException e) {
      throw new RuntimeException(e);
    } catch (ExecutionException e) {
      // This should not normally happen since junit statements trap and record errors and failures.
      throw Throwables.propagate(e.getCause());
    } finally {
      // In case of error, cancel all in-flight concurrent tasks
      while (!concurrentTasks.isEmpty()) {
        concurrentTasks.poll().cancel(true);
      }
    }
  }
}

@@ -180,345 +180,90 @@
 ort 
-com.google.common.base.Preconditions;%0Aimport com.google.common.base.Throwables;%0Aimport com.google.common.util.concurrent.ThreadFactoryBuilder;%0Aimport java.util.LinkedList;%0Aimport java.util.Queue;%0Aimport java.util.concurrent.CompletionService;%0Aimport java.util.concurrent.ExecutionException;%0Aimport java.util.concurrent.ExecutorCompletion
+java.util.LinkedList;%0Aimport java.util.Queue;%0Aimport java.util.concurrent.Executor
 Serv
@@ -310,44 +310,8 @@
 rs;%0A
-import java.util.concurrent.Future;%0A
 impo
@@ -349,16 +349,164 @@
 actory;%0A
+import java.util.concurrent.TimeUnit;%0A%0Aimport com.google.common.base.Preconditions;%0Aimport com.google.common.util.concurrent.ThreadFactoryBuilder;%0A%0A
 import o
@@ -740,102 +740,32 @@
 nal 
-CompletionService%3CVoid%3E completionService;%0A  private final Queue%3CFuture%3CVoid%3E%3E concurrentTasks
+ExecutorService executor
 ;%0A  
@@ -1861,182 +1861,73 @@
     
-completionService = new ExecutorCompletionService%3CVoid%3E(%0A        Executors.newFixedThreadPool(numThreads, threadFactory));%0A    concurrentTasks = new LinkedList%3CFuture%3CVoid%3E%3E(
+executor = Executors.newFixedThreadPool(numThreads, threadFactory
 );%0A 
@@ -2086,47 +2086,16 @@
     
-concurrentTasks.offer(completionService
+executor
 .sub
@@ -2104,39 +2104,32 @@
 t(childStatement
-, null)
 );%0A    %7D else %7B%0A
@@ -2435,47 +2435,16 @@
     
-concurrentTasks.offer(completionService
+executor
 .sub
@@ -2465,15 +2465,8 @@
 ment
-, null)
 );%0A 
@@ -2988,24 +2988,49 @@
 inished() %7B%0A
+    executor.shutdown();%0A
     try %7B%0A  
@@ -3087,90 +3087,66 @@
     
-while (!concurrentTasks.isEmpty()) %7B%0A        concurrentTasks.poll().get();%0A      %7D
+executor.awaitTermination(Long.MAX_VALUE, TimeUnit.DAYS);%0A
 %0A   
@@ -3344,364 +3344,112 @@
   %7D 
-catch (ExecutionException e) %7B%0A      // This should not normally happen since junit statements trap and record errors and failures.%0A      throw Throwables.propagate(e.getCause());%0A    %7D finally %7B%0A      // In case of error, cancel all in-flight concurrent tasks%0A      while (!concurrentTasks.isEmpty()) %7B%0A        concurrentTasks.poll().cancel(true);%0A      %7D
+finally %7B%0A      // In case of error, cancel all in-flight concurrent tasks%0A      executor.shutdownNow();
 %0A   

Simplify `ConcurrentRunnerScheduler` & cleanup.