Cancellable Executions (#1920)

* add failing tests

* add another failing test to test writes

* wip try to thread CancellationHandler through

* finish threading cancellation handler through

* try to fix test

* wip cfuture/cpromise

* test passes, everything compiles

* reenable tests

* remove sys error

* fix indenting

* don't recover on flow stop

* move failFastSequence to CFuture; add CFuture helper methods

* uncancellable

* move CFuture/CPromise

* CFuture.fromFuture

* add typeclass

* remove cec

* execution changes

* future cache backwards compatible

* AsyncFlowDef fixes

* change jdk

* add failing test

* delete failing test

* add flatmap test

* add flow step listener

* tryFailure

* handle null throwable

* respond to review

* stop flow in promise

Author: stephbian

.travis.yml

@@ -1,5 +1,5 @@
 language: scala
-jdk: oraclejdk8
+jdk: openjdk8
 sudo: false
 
 before_install:

scalding-core/src/main/java/com/twitter/scalding/cascading_interop/FlowListenerPromise.java

@@ -16,6 +16,9 @@
 package com.twitter.scalding.cascading_interop;
 
 import cascading.flow.FlowListener;
+import cascading.flow.FlowException;
+import cascading.flow.FlowStepListener;
+import cascading.flow.FlowStep;
 import cascading.flow.Flow;
 import cascading.stats.CascadingStats;
 
@@ -28,6 +31,12 @@
  * deal with in scala
  */
 public class FlowListenerPromise {
+  public static class FlowStopException extends Exception {
+    public FlowStopException(String message) {
+      super(message);
+    }
+  }
+
   /*
    * This starts the flow and applies a mapping function fn in
    * the same thread that completion happens
@@ -37,7 +46,7 @@ public static <Config, T> Future<T> start(Flow<Config> flow, final scala.Functio
     flow.addListener(new FlowListener() {
       public void onStarting(Flow f) { } // ignore
       public void onStopping(Flow f) { // in case of runtime exception cascading call onStopping
-        result.tryFailure(new Exception("Flow was stopped"));
+        result.tryFailure(new FlowStopException("Flow was stopped"));
       }
       public void onCompleted(Flow f) {
         // This is always called, but onThrowable is called first
@@ -48,15 +57,30 @@ public void onCompleted(Flow f) {
               T toPut = (T) fn.apply(f);
               result.success(toPut);
             } catch (Throwable t) {
-              result.failure(t);
+              result.tryFailure(t);
             }
           } else {
-            result.failure(new Exception("Flow was not successfully finished"));
+            result.tryFailure(new Exception("Flow was not successfully finished"));
           }
         }
       }
       public boolean onThrowable(Flow f, Throwable t) {
-        result.failure(t);
+        result.tryFailure(t);
+        // The exception is handled by the owner of the promise and should not be rethrown
+        return true;
+      }
+    });
+    flow.addStepListener(new FlowStepListener() {
+      public void onStepStarting(FlowStep flowStep) { } // ignore
+      public void onStepRunning(FlowStep flowStep) { } // ignore
+      public void onStepCompleted(FlowStep flowStep) { } // ignore
+      public void onStepStopping(FlowStep f) { result.tryFailure(new FlowStopException("Flow step was stopped")); }
+      public boolean onStepThrowable(FlowStep f, Throwable t) {
+        if (t != null) {
+          result.tryFailure(t);
+        } else {
+          result.tryFailure(new FlowException("Flow step failed: " + f.getName()));
+        }
         // The exception is handled by the owner of the promise and should not be rethrown
         return true;
       }

scalding-core/src/main/scala/com/twitter/scalding/CFuture.scala

@@ -0,0 +1,53 @@
+package com.twitter.scalding
+
+import scala.concurrent.{ Future, ExecutionContext => ConcurrentExecutionContext }
+
+/**
+ * Represents a cancellable future.
+ */
+case class CFuture[+T](future: Future[T], cancellationHandler: CancellationHandler) {
+  def map[S](fn: T => S)(implicit cec: ConcurrentExecutionContext): CFuture[S] = {
+    val mapped = future.map(fn)
+    CFuture(mapped, cancellationHandler)
+  }
+
+  def mapFuture[S](fn: Future[T] => Future[S]): CFuture[S] = {
+    val transformed = fn(future)
+    CFuture(transformed, cancellationHandler)
+  }
+
+  def zip[U](other: CFuture[U])(implicit cec: ConcurrentExecutionContext): CFuture[(T, U)] = {
+    val zippedFut: Future[(T, U)] = Execution.failFastZip(future, other.future)
+    val cancelHandler = cancellationHandler.compose(other.cancellationHandler)
+
+    CFuture(zippedFut, cancelHandler)
+  }
+}
+
+object CFuture {
+  def successful[T](result: T): CFuture[T] = {
+    CFuture(Future.successful(result), CancellationHandler.empty)
+  }
+
+  def failed(t: Throwable): CFuture[Nothing] = {
+    val f = Future.failed(t)
+    CFuture(f, CancellationHandler.empty)
+  }
+
+  def uncancellable[T](fut: Future[T]): CFuture[T] = {
+    CFuture(fut, CancellationHandler.empty)
+  }
+
+  def fromFuture[T](fut: Future[CFuture[T]])(implicit cec: ConcurrentExecutionContext): CFuture[T] = {
+    CFuture(fut.flatMap(_.future), CancellationHandler.fromFuture(fut.map(_.cancellationHandler)))
+  }
+
+  /**
+   * Use our internal faster failing zip function rather than the standard one due to waiting
+   */
+  def failFastSequence[T](t: Iterable[CFuture[T]])(implicit cec: ConcurrentExecutionContext): CFuture[List[T]] = {
+    t.foldLeft(CFuture.successful(Nil: List[T])) { (f, i) =>
+      f.zip(i).map { case (tail, h) => h :: tail }
+    }.map(_.reverse)
+  }
+}

scalding-core/src/main/scala/com/twitter/scalding/CPromise.scala

@@ -0,0 +1,25 @@
+package com.twitter.scalding
+
+import scala.concurrent.{ Future, Promise, ExecutionContext => ConcurrentExecutionContext }
+
+/**
+ * Represents a cancellable promise.
+ */
+case class CPromise[T](promise: Promise[T], cancellationHandler: Promise[CancellationHandler]) {
+  /**
+   * Creates a CFuture using the given promises.
+   */
+  def cfuture: CFuture[T] = {
+    CFuture(promise.future, CancellationHandler.fromFuture(cancellationHandler.future))
+  }
+
+  def completeWith(other: CFuture[T]): this.type = {
+    // fullfill the main and cancellation handler promises
+    promise.completeWith(other.future)
+    cancellationHandler.completeWith(Future.successful(other.cancellationHandler))
+    this
+  }
+}
+object CPromise {
+  def apply[T](): CPromise[T] = CPromise(Promise[T](), Promise[CancellationHandler]())
+}

scalding-core/src/main/scala/com/twitter/scalding/CancellationHandler.scala

@@ -0,0 +1,28 @@
+package com.twitter.scalding
+
+import scala.concurrent.{ Future, ExecutionContext => ConcurrentExecutionContext }
+
+sealed trait CancellationHandler { outer =>
+  def stop()(implicit ec: ConcurrentExecutionContext): Future[Unit]
+  def compose(other: CancellationHandler): CancellationHandler = new CancellationHandler {
+    override def stop()(implicit ec: ConcurrentExecutionContext): Future[Unit] = {
+      other.stop().zip(outer.stop()).map(_ => ())
+    }
+  }
+}
+
+object CancellationHandler {
+  val empty: CancellationHandler = new CancellationHandler {
+    def stop()(implicit ec: ConcurrentExecutionContext): Future[Unit] = Future.successful(())
+  }
+
+  def fromFn(fn: ConcurrentExecutionContext => Future[Unit]): CancellationHandler = new CancellationHandler {
+    override def stop()(implicit ec: ConcurrentExecutionContext): Future[Unit] = fn(ec)
+  }
+
+  def fromFuture(f: Future[CancellationHandler]): CancellationHandler = new CancellationHandler {
+    override def stop()(implicit ec: ConcurrentExecutionContext): Future[Unit] = {
+      f.flatMap(_.stop())
+    }
+  }
+}