Cancel in-progress async operations when the pipeline is aborted. (#1205

)

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/AsyncOperationProcessor.java

@@ -62,6 +62,11 @@ public void finalize(final WritePipe<O> outputPipe) {
     }
   }
 
+  @Override
+  public void abort() {
+    inProgress.forEach(future -> future.cancel(true));
+  }
+
   private void outputNextCompletedTask(final WritePipe<O> outputPipe) {
     try {
       waitForAnyFutureToComplete();

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/BatchingReadPipe.java

@@ -36,6 +36,11 @@ public boolean hasMore() {
     return input.hasMore();
   }
 
+  @Override
+  public boolean isAborted() {
+    return input.isAborted();
+  }
+
   @Override
   public List<T> get() {
     final T firstItem = input.get();

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/Pipe.java

@@ -55,6 +55,11 @@ public boolean isOpen() {
     return !closed.get() && !aborted.get();
   }
 
+  @Override
+  public boolean isAborted() {
+    return aborted.get();
+  }
+
   /**
    * Get the number of items that can be queued inside this pipe.
    *

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/ProcessingStage.java

@@ -35,6 +35,9 @@ public void run() {
     while (inputPipe.hasMore()) {
       processor.processNextInput(inputPipe, outputPipe);
     }
+    if (inputPipe.isAborted()) {
+      processor.abort();
+    }
     processor.finalize(outputPipe);
     outputPipe.close();
   }

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/Processor.java

@@ -16,4 +16,6 @@ interface Processor<I, O> {
   void processNextInput(final ReadPipe<I> inputPipe, final WritePipe<O> outputPipe);
 
   default void finalize(final WritePipe<O> outputPipe) {}
+
+  default void abort() {}
 }

services/pipeline/src/main/java/tech/pegasys/pantheon/services/pipeline/ReadPipe.java

@@ -30,6 +30,13 @@ public interface ReadPipe<T> {
    */
   boolean hasMore();
 
+  /**
+   * Determines if this pipeline this pipe is a part of has been aborted.
+   *
+   * @return true if the pipeline has been aborted, otherwise false.
+   */
+  boolean isAborted();
+
   /**
    * Get and remove the next item from this pipe. This method will block until the next item is
    * available but may still return <code>null</code> if the pipe is closed or the thread

services/pipeline/src/test/java/tech/pegasys/pantheon/services/pipeline/AsyncOperationProcessorTest.java

@@ -96,6 +96,21 @@ public void shouldOutputRemainingInProgressTasksWhenFinalizing() {
     verify(writePipe).put("b");
   }
 
+  @Test
+  public void shouldCancelInProgressTasksWhenAborted() {
+    final CompletableFuture<String> task1 = new CompletableFuture<>();
+    final CompletableFuture<String> task2 = new CompletableFuture<>();
+    when(readPipe.get()).thenReturn(task1).thenReturn(task2);
+
+    processor.processNextInput(readPipe, writePipe);
+    processor.processNextInput(readPipe, writePipe);
+
+    processor.abort();
+
+    assertThat(task1).isCancelled();
+    assertThat(task2).isCancelled();
+  }
+
   @Test
   public void shouldInterruptThreadWhenFutureCompletes() {
     // Ensures that if we're waiting for the next input we wake up and output completed tasks

services/pipeline/src/test/java/tech/pegasys/pantheon/services/pipeline/PipeTest.java

@@ -53,6 +53,7 @@ public void shouldNotHaveMoreWhenAbortedEvenIfNotEmpty() {
     pipe.abort();
 
     assertThat(pipe.hasMore()).isFalse();
+    assertThat(pipe.isAborted()).isTrue();
   }
 
   @Test
@@ -79,6 +80,7 @@ public void shouldNotWaitToReachMaximumSizeBeforeReturningBatch() {
   public void shouldNotBeOpenAfterAbort() {
     pipe.abort();
     assertThat(pipe.isOpen()).isFalse();
+    assertThat(pipe.isAborted()).isTrue();
   }
 
   @Test

services/pipeline/src/test/java/tech/pegasys/pantheon/services/pipeline/ProcessingStageTest.java

@@ -74,4 +74,14 @@ public void shouldFinalizeSingleStepStageAndCloseOutputPipeWhenInputCloses() {
     verifyNoMoreInteractions(singleStep);
     assertThat(outputPipe.isOpen()).isFalse();
   }
+
+  @Test
+  public void shouldAbortProcessorIfReadPipeIsAborted() {
+    inputPipe.abort();
+    stage.run();
+
+    verify(singleStep).abort();
+    verify(singleStep).finalize(outputPipe);
+    assertThat(outputPipe.isOpen()).isFalse();
+  }
 }