[Streaming] Fault Tolerance Implementation (#10595)

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/function/Function.java

@@ -7,4 +7,26 @@ import java.io.Serializable;
  */
 public interface Function extends Serializable {
 
+  /**
+   * This method will be called periodically by framework, you should return a a serializable
+   * object which represents function state, framework will help you to serialize this object, save
+   * it to storage, and load it back when in fail-over through.
+   * {@link Function#loadCheckpoint(Serializable)}.
+   *
+   * @return A serializable object which represents function state.
+   */
+  default Serializable saveCheckpoint() {
+    return null;
+  }
+
+  /**
+   * This method will be called by framework when a worker died and been restarted.
+   * We will pass the last object you returned in {@link Function#saveCheckpoint()} when
+   * doing checkpoint, you are responsible to load this object back to you function.
+   *
+   * @param checkpointObject the last object you returned in {@link Function#saveCheckpoint()}
+   */
+  default void loadCheckpoint(Serializable checkpointObject) {
+  }
+
 }

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/function/impl/SourceFunction.java

@@ -9,9 +9,9 @@ import io.ray.streaming.api.function.Function;
  */
 public interface SourceFunction<T> extends Function {
 
-  void init(int parallel, int index);
+  void init(int parallelism, int index);
 
-  void run(SourceContext<T> ctx) throws Exception;
+  void fetch(SourceContext<T> ctx) throws Exception;
 
   void close();
 

streaming/java/streaming-api/src/main/java/io/ray/streaming/api/function/internal/CollectionSourceFunction.java

@@ -1,7 +1,6 @@
 package io.ray.streaming.api.function.internal;
 
 import io.ray.streaming.api.function.impl.SourceFunction;
-import java.util.ArrayList;
 import java.util.Collection;
 
 /**
@@ -12,22 +11,25 @@ import java.util.Collection;
 public class CollectionSourceFunction<T> implements SourceFunction<T> {
 
   private Collection<T> values;
+  private boolean finished = false;
 
   public CollectionSourceFunction(Collection<T> values) {
     this.values = values;
   }
 
   @Override
-  public void init(int parallel, int index) {
+  public void init(int totalParallel, int currentIndex) {
   }
 
   @Override
-  public void run(SourceContext<T> ctx) throws Exception {
+  public void fetch(SourceContext<T> ctx) throws Exception {
+    if (finished) {
+      return;
+    }
     for (T value : values) {
       ctx.collect(value);
     }
-    // empty collection
-    values = new ArrayList<>();
+    finished = true;
   }
 
   @Override

streaming/java/streaming-api/src/main/java/io/ray/streaming/message/KeyRecord.java

@@ -1,6 +1,5 @@
 package io.ray.streaming.message;
 
-
 import java.util.Objects;
 
 public class KeyRecord<K, T> extends Record<T> {

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/Operator.java

@@ -25,4 +25,13 @@ public interface Operator extends Serializable {
 
   ChainStrategy getChainStrategy();
 
+  /**
+   * See {@link Function#saveCheckpoint()}.
+   */
+  Serializable saveCheckpoint();
+
+  /**
+   * See {@link Function#loadCheckpoint(Serializable)}.
+   */
+  void loadCheckpoint(Serializable checkpointObject);
 }

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/SourceOperator.java

@@ -4,7 +4,7 @@ import io.ray.streaming.api.function.impl.SourceFunction.SourceContext;
 
 public interface SourceOperator<T> extends Operator {
 
-  void run();
+  void fetch();
 
   SourceContext<T> getSourceContext();
 

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/StreamOperator.java

@@ -8,6 +8,7 @@ import io.ray.streaming.api.function.RichFunction;
 import io.ray.streaming.api.function.internal.Functions;
 import io.ray.streaming.message.KeyRecord;
 import io.ray.streaming.message.Record;
+import java.io.Serializable;
 import java.util.List;
 
 public abstract class StreamOperator<F extends Function> implements Operator {
@@ -72,6 +73,16 @@ public abstract class StreamOperator<F extends Function> implements Operator {
     }
   }
 
+  @Override
+  public Serializable saveCheckpoint() {
+    return function.saveCheckpoint();
+  }
+
+  @Override
+  public void loadCheckpoint(Serializable checkpointObject) {
+    function.loadCheckpoint(checkpointObject);
+  }
+
   @Override
   public String getName() {
     return name;

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/chain/ChainedOperator.java

@@ -13,6 +13,7 @@ import io.ray.streaming.operator.OperatorType;
 import io.ray.streaming.operator.SourceOperator;
 import io.ray.streaming.operator.StreamOperator;
 import io.ray.streaming.operator.TwoInputOperator;
+import java.io.Serializable;
 import java.lang.reflect.Proxy;
 import java.util.Collections;
 import java.util.List;
@@ -85,6 +86,23 @@ public abstract class ChainedOperator extends StreamOperator<Function> {
     return tailOperator;
   }
 
+  @Override
+  public Serializable saveCheckpoint() {
+    Object[] checkpoints = new Object[operators.size()];
+    for (int i = 0; i < operators.size(); ++i) {
+      checkpoints[i] = operators.get(i).saveCheckpoint();
+    }
+    return checkpoints;
+  }
+
+  @Override
+  public void loadCheckpoint(Serializable checkpointObject) {
+    Serializable[] checkpoints = (Serializable[]) checkpointObject;
+    for (int i = 0; i < operators.size(); ++i) {
+      operators.get(i).loadCheckpoint(checkpoints[i]);
+    }
+  }
+
   private RuntimeContext createRuntimeContext(RuntimeContext runtimeContext, int index) {
     return (RuntimeContext) Proxy.newProxyInstance(runtimeContext.getClass().getClassLoader(),
         new Class[] {RuntimeContext.class},
@@ -125,8 +143,8 @@ public abstract class ChainedOperator extends StreamOperator<Function> {
     }
 
     @Override
-    public void run() {
-      sourceOperator.run();
+    public void fetch() {
+      sourceOperator.fetch();
     }
 
     @Override

streaming/java/streaming-api/src/main/java/io/ray/streaming/operator/impl/SourceOperatorImpl.java

@@ -29,9 +29,9 @@ public class SourceOperatorImpl<T> extends StreamOperator<SourceFunction<T>>
   }
 
   @Override
-  public void run() {
+  public void fetch() {
     try {
-      this.function.run(this.sourceContext);
+      this.function.fetch(this.sourceContext);
     } catch (Exception e) {
       throw new RuntimeException(e);
     }