Streaming state (#7348)

streaming/java/streaming-runtime/src/main/java/io/ray/streaming/runtime/core/processor/StreamProcessor.java

@@ -14,6 +14,7 @@ import org.slf4j.LoggerFactory;
  * @param <P> Type of the specific operator class.
  */
 public abstract class StreamProcessor<T, P extends Operator> implements Processor<T> {
+
   private static final Logger LOGGER = LoggerFactory.getLogger(StreamProcessor.class);
 
   protected List<Collector> collectors;

streaming/java/streaming-runtime/src/main/java/io/ray/streaming/runtime/worker/context/RayRuntimeContext.java

@@ -2,19 +2,39 @@ package io.ray.streaming.runtime.worker.context;
 
 import static io.ray.streaming.util.Config.STREAMING_BATCH_MAX_COUNT;
 
+import com.google.common.base.Preconditions;
 import io.ray.streaming.api.context.RuntimeContext;
 import io.ray.streaming.runtime.core.graph.ExecutionTask;
+import io.ray.streaming.state.backend.AbstractKeyStateBackend;
+import io.ray.streaming.state.backend.KeyStateBackend;
+import io.ray.streaming.state.backend.OperatorStateBackend;
+import io.ray.streaming.state.keystate.desc.AbstractStateDescriptor;
+import io.ray.streaming.state.keystate.desc.ListStateDescriptor;
+import io.ray.streaming.state.keystate.desc.MapStateDescriptor;
+import io.ray.streaming.state.keystate.desc.ValueStateDescriptor;
+import io.ray.streaming.state.keystate.state.ListState;
+import io.ray.streaming.state.keystate.state.MapState;
+import io.ray.streaming.state.keystate.state.ValueState;
 import java.util.Map;
 
 /**
  * Use Ray to implement RuntimeContext.
  */
 public class RayRuntimeContext implements RuntimeContext {
+
+  private final Long maxBatch;
+  /**
+   * Backend for keyed state. This might be empty if we're not on a keyed stream.
+   */
+  protected transient KeyStateBackend keyStateBackend;
+  /**
+   * Backend for operator state. This might be empty
+   */
+  protected transient OperatorStateBackend operatorStateBackend;
   private int taskId;
   private int taskIndex;
   private int parallelism;
-  private Long batchId;
-  private final Long maxBatch;
+  private Long checkpointId;
   private Map<String, String> config;
 
   public RayRuntimeContext(ExecutionTask executionTask, Map<String, String> config,
@@ -46,8 +66,19 @@ public class RayRuntimeContext implements RuntimeContext {
   }
 
   @Override
-  public Long getBatchId() {
-    return batchId;
+  public Long getCheckpointId() {
+    return checkpointId;
+  }
+
+  @Override
+  public void setCheckpointId(long checkpointId) {
+    if (this.keyStateBackend != null) {
+      this.keyStateBackend.setCheckpointId(checkpointId);
+    }
+    if (this.operatorStateBackend != null) {
+      this.operatorStateBackend.setCheckpointId(checkpointId);
+    }
+    this.checkpointId = checkpointId;
   }
 
   @Override
@@ -55,7 +86,48 @@ public class RayRuntimeContext implements RuntimeContext {
     return maxBatch;
   }
 
-  public void setBatchId(Long batchId) {
-    this.batchId = batchId;
+  @Override
+  public Map<String, String> getConfig() {
+    return config;
+  }
+
+
+  @Override
+  public void setCurrentKey(Object key) {
+    this.keyStateBackend.setCurrentKey(key);
+  }
+
+  @Override
+  public KeyStateBackend getKeyStateBackend() {
+    return keyStateBackend;
+  }
+
+  @Override
+  public void setKeyStateBackend(KeyStateBackend keyStateBackend) {
+    this.keyStateBackend = keyStateBackend;
+  }
+
+  @Override
+  public <T> ValueState<T> getValueState(ValueStateDescriptor<T> stateDescriptor) {
+    stateSanityCheck(stateDescriptor, this.keyStateBackend);
+    return this.keyStateBackend.getValueState(stateDescriptor);
+  }
+
+  @Override
+  public <T> ListState<T> getListState(ListStateDescriptor<T> stateDescriptor) {
+    stateSanityCheck(stateDescriptor, this.keyStateBackend);
+    return this.keyStateBackend.getListState(stateDescriptor);
+  }
+
+  @Override
+  public <S, T> MapState<S, T> getMapState(MapStateDescriptor<S, T> stateDescriptor) {
+    stateSanityCheck(stateDescriptor, this.keyStateBackend);
+    return this.keyStateBackend.getMapState(stateDescriptor);
+  }
+
+  protected void stateSanityCheck(AbstractStateDescriptor stateDescriptor,
+                                  AbstractKeyStateBackend backend) {
+    Preconditions.checkNotNull(stateDescriptor, "The state properties must not be null");
+    Preconditions.checkNotNull(backend, "backend must not be null");
   }
 }

streaming/java/streaming-runtime/src/main/java/io/ray/streaming/runtime/worker/tasks/StreamTask.java

@@ -26,6 +26,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public abstract class StreamTask implements Runnable {
+
   private static final Logger LOG = LoggerFactory.getLogger(StreamTask.class);
 
   protected int taskId;
@@ -41,8 +42,8 @@ public abstract class StreamTask implements Runnable {
     this.worker = worker;
     prepareTask();
 
-    this.thread = new Thread(Ray.wrapRunnable(this), this.getClass().getName()
-      + "-" + System.currentTimeMillis());
+    this.thread = new Thread(Ray.wrapRunnable(this),
+        this.getClass().getName() + "-" + System.currentTimeMillis());
     this.thread.setDaemon(true);
   }
 
@@ -105,8 +106,8 @@ public abstract class StreamTask implements Runnable {
       reader = new DataReader(channelIDs, inputActors, queueConf);
     }
 
-    RuntimeContext runtimeContext = new RayRuntimeContext(
-        worker.getExecutionTask(), worker.getConfig(), executionNode.getParallelism());
+    RuntimeContext runtimeContext = new RayRuntimeContext(worker.getExecutionTask(),
+        worker.getConfig(), executionNode.getParallelism());
 
     processor.open(collectors, runtimeContext);