shutdown improvements (#841)

* register sub tasks in ParalellImportChainSegmentTask
* rememger and shutdown eth services in EthScheduler

Fixes #768
Signed-off-by: Adrian Sutton <adrian.sutton@consensys.net>

ethereum/eth/src/main/java/tech/pegasys/pantheon/ethereum/eth/manager/AbstractEthTask.java

@@ -117,6 +117,31 @@ public abstract class AbstractEthTask<T> implements EthTask<T> {
     }
   }
 
+  /**
+   * Utility for registring completable futures for cleanup if this EthTask is cancelled.
+   *
+   * @param subTaskFuture the future to be reigstered.
+   * @param <S> the type of data returned from the CompletableFuture
+   * @return The completableFuture that was executed
+   */
+  protected final <S> CompletableFuture<S> registerSubTask(
+      final CompletableFuture<S> subTaskFuture) {
+    synchronized (result) {
+      if (!isCancelled()) {
+        subTaskFutures.add(subTaskFuture);
+        subTaskFuture.whenComplete(
+            (r, t) -> {
+              subTaskFutures.remove(subTaskFuture);
+            });
+        return subTaskFuture;
+      } else {
+        final CompletableFuture<S> future = new CompletableFuture<>();
+        future.completeExceptionally(new CancellationException());
+        return future;
+      }
+    }
+  }
+
   /**
    * Helper method for sending subTask to worker that will clean up if this EthTask is cancelled.
    *

ethereum/eth/src/main/java/tech/pegasys/pantheon/ethereum/eth/manager/EthScheduler.java

@@ -15,8 +15,10 @@ package tech.pegasys.pantheon.ethereum.eth.manager;
 import tech.pegasys.pantheon.util.ExceptionUtils;
 
 import java.time.Duration;
+import java.util.Collection;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedDeque;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -47,6 +49,8 @@ public class EthScheduler {
   private final ExecutorService servicesExecutor;
   private final ExecutorService computationExecutor;
 
+  private Collection<CompletableFuture<?>> serviceFutures = new ConcurrentLinkedDeque<>();
+
   public EthScheduler(
       final int syncWorkerCount, final int txWorkerCount, final int computationWorkerCount) {
     this(
@@ -126,7 +130,10 @@ public class EthScheduler {
   }
 
   public <T> CompletableFuture<T> scheduleServiceTask(final EthTask<T> task) {
-    return task.runAsync(servicesExecutor);
+    final CompletableFuture<T> serviceFuture = task.runAsync(servicesExecutor);
+    serviceFutures.add(serviceFuture);
+    serviceFuture.whenComplete((r, t) -> serviceFutures.remove(serviceFuture));
+    return serviceFuture;
   }
 
   public <T> CompletableFuture<T> scheduleComputationTask(final Supplier<T> computation) {
@@ -230,6 +237,7 @@ public class EthScheduler {
 
   void awaitStop() throws InterruptedException {
     shutdown.await();
+    serviceFutures.forEach(future -> future.cancel(true));
     if (!syncWorkerExecutor.awaitTermination(2L, TimeUnit.MINUTES)) {
       LOG.error("{} worker executor did not shutdown cleanly.", this.getClass().getSimpleName());
       syncWorkerExecutor.shutdownNow();

ethereum/eth/src/main/java/tech/pegasys/pantheon/ethereum/eth/sync/tasks/ParallelImportChainSegmentTask.java

@@ -141,12 +141,16 @@ public class ParallelImportChainSegmentTask<C, B> extends AbstractEthTask<List<B
       final EthScheduler scheduler = ethContext.getScheduler();
       final CompletableFuture<?> downloadHeaderFuture =
           scheduler.scheduleServiceTask(downloadHeadersTask);
+      registerSubTask(downloadHeaderFuture);
       final CompletableFuture<?> validateHeaderFuture =
           scheduler.scheduleServiceTask(validateHeadersTask);
+      registerSubTask(validateHeaderFuture);
       final CompletableFuture<?> downloadBodiesFuture =
           scheduler.scheduleServiceTask(downloadBodiesTask);
+      registerSubTask(downloadBodiesFuture);
       final CompletableFuture<AbstractPeerTask.PeerTaskResult<List<List<B>>>> validateBodiesFuture =
           scheduler.scheduleServiceTask(validateAndImportBodiesTask);
+      registerSubTask(validateBodiesFuture);
 
       // Hook in pipeline completion signaling.
       downloadHeadersTask.shutdown();