[FLINK-7213] Introduce state management by OperatorID in TaskManager

flink-contrib/flink-statebackend-rocksdb/src/test/java/org/apache/flink/contrib/streaming/state/RocksDBAsyncSnapshotTest.java

@@ -32,8 +32,10 @@
 import org.apache.flink.runtime.checkpoint.CheckpointMetaData;
 import org.apache.flink.runtime.checkpoint.CheckpointMetrics;
 import org.apache.flink.runtime.checkpoint.CheckpointOptions;
-import org.apache.flink.runtime.checkpoint.SubtaskState;
+import org.apache.flink.runtime.checkpoint.OperatorSubtaskState;
+import org.apache.flink.runtime.checkpoint.TaskStateSnapshot;
 import org.apache.flink.runtime.execution.Environment;
+import org.apache.flink.runtime.jobgraph.OperatorID;
 import org.apache.flink.runtime.operators.testutils.DummyEnvironment;
 import org.apache.flink.runtime.operators.testutils.MockInputSplitProvider;
 import org.apache.flink.runtime.state.AbstractKeyedStateBackend;
@@ -74,14 +76,15 @@
 import java.lang.reflect.Field;
 import java.net.URI;
 import java.util.Arrays;
+import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.RunnableFuture;
 import java.util.concurrent.TimeUnit;
 
-import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
@@ -137,6 +140,7 @@ public String getKey(String value) throws Exception {
 		streamConfig.setStateBackend(backend);
 
 		streamConfig.setStreamOperator(new AsyncCheckpointOperator());
+		streamConfig.setOperatorID(new OperatorID());
 
 		final OneShotLatch delayCheckpointLatch = new OneShotLatch();
 		final OneShotLatch ensureCheckpointLatch = new OneShotLatch();
@@ -152,7 +156,7 @@ public String getKey(String value) throws Exception {
 			public void acknowledgeCheckpoint(
 					long checkpointId,
 					CheckpointMetrics checkpointMetrics,
-					SubtaskState checkpointStateHandles) {
+					TaskStateSnapshot checkpointStateHandles) {
 
 				super.acknowledgeCheckpoint(checkpointId, checkpointMetrics);
 
@@ -164,8 +168,16 @@ public void acknowledgeCheckpoint(
 					throw new RuntimeException(e);
 				}
 
+				boolean hasManagedKeyedState = false;
+				for (Map.Entry<OperatorID, OperatorSubtaskState> entry : checkpointStateHandles.getSubtaskStateMappings()) {
+					OperatorSubtaskState state = entry.getValue();
+					if (state != null) {
+						hasManagedKeyedState |= state.getManagedKeyedState() != null;
+					}
+				}
+
 				// should be one k/v state
-				assertNotNull(checkpointStateHandles.getManagedKeyedState());
+				assertTrue(hasManagedKeyedState);
 
 				// we now know that the checkpoint went through
 				ensureCheckpointLatch.trigger();
@@ -241,6 +253,7 @@ public String getKey(String value) throws Exception {
 		streamConfig.setStateBackend(backend);
 
 		streamConfig.setStreamOperator(new AsyncCheckpointOperator());
+		streamConfig.setOperatorID(new OperatorID());
 
 		StreamMockEnvironment mockEnv = new StreamMockEnvironment(
 				testHarness.jobConfig,

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinator.java

@@ -40,7 +40,6 @@
 import org.apache.flink.runtime.messages.checkpoint.AcknowledgeCheckpoint;
 import org.apache.flink.runtime.messages.checkpoint.DeclineCheckpoint;
 import org.apache.flink.runtime.state.SharedStateRegistry;
-import org.apache.flink.runtime.state.TaskStateHandles;
 import org.apache.flink.runtime.taskmanager.DispatcherThreadFactory;
 import org.apache.flink.util.Preconditions;
 import org.apache.flink.util.StringUtils;
@@ -1016,7 +1015,7 @@ int getNumScheduledTasks() {
 	 * Restores the latest checkpointed state.
 	 *
 	 * @param tasks Map of job vertices to restore. State for these vertices is
-	 * restored via {@link Execution#setInitialState(TaskStateHandles)}.
+	 * restored via {@link Execution#setInitialState(TaskStateSnapshot)}.
 	 * @param errorIfNoCheckpoint Fail if no completed checkpoint is available to
 	 * restore from.
 	 * @param allowNonRestoredState Allow checkpoint state that cannot be mapped
@@ -1102,7 +1101,7 @@ public boolean restoreLatestCheckpointedState(
 	 *                         mapped to any job vertex in tasks.
 	 * @param tasks            Map of job vertices to restore. State for these 
 	 *                         vertices is restored via 
-	 *                         {@link Execution#setInitialState(TaskStateHandles)}.
+	 *                         {@link Execution#setInitialState(TaskStateSnapshot)}.
 	 * @param userClassLoader  The class loader to resolve serialized classes in 
 	 *                         legacy savepoint versions. 
 	 */
@@ -1256,7 +1255,7 @@ private void discardSubtaskState(
 			final JobID jobId,
 			final ExecutionAttemptID executionAttemptID,
 			final long checkpointId,
-			final SubtaskState subtaskState) {
+			final TaskStateSnapshot subtaskState) {
 
 		if (subtaskState != null) {
 			executor.execute(new Runnable() {

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorGateway.java

@@ -29,7 +29,7 @@ void acknowledgeCheckpoint(
 			final ExecutionAttemptID executionAttemptID,
 			final long checkpointId,
 			final CheckpointMetrics checkpointMetrics,
-			final SubtaskState subtaskState);
+			final TaskStateSnapshot subtaskState);
 
 	void declineCheckpoint(
 			JobID jobID,

flink-runtime/src/main/java/org/apache/flink/runtime/checkpoint/OperatorState.java

@@ -30,8 +30,8 @@
 import java.util.Objects;
 
 /**
- * Simple container class which contains the raw/managed/legacy operator state and key-group state handles for the sub
- * tasks of an operator.
+ * Simple container class which contains the raw/managed/legacy operator state and key-group state handles from all sub
+ * tasks of an operator and therefore represents the complete state of a logical operator.
  */
 public class OperatorState implements CompositeStateHandle {