[FLINK-6640] Ensure registration of shared state happens before exter…

…nalizing a checkpoint

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

@@ -40,11 +40,11 @@
 import org.apache.flink.runtime.jobgraph.tasks.ExternalizedCheckpointSettings;
 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;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -173,6 +173,9 @@ public class CheckpointCoordinator {
 	@Nullable
 	private CheckpointStatsTracker statsTracker;
 
+	/** Registry that tracks state which is shared across (incremental) checkpoints */
+	private final SharedStateRegistry sharedStateRegistry;
+
 	// --------------------------------------------------------------------------------------------
 
 	public CheckpointCoordinator(
@@ -226,6 +229,7 @@ public CheckpointCoordinator(
 		this.completedCheckpointStore = checkNotNull(completedCheckpointStore);
 		this.checkpointDirectory = checkpointDirectory;
 		this.executor = checkNotNull(executor);
+		this.sharedStateRegistry = new SharedStateRegistry(executor);
 
 		this.recentPendingCheckpoints = new ArrayDeque<>(NUM_GHOST_CHECKPOINT_IDS);
 		this.masterHooks = new HashMap<>();
@@ -836,6 +840,10 @@ private void completePendingCheckpoint(PendingCheckpoint pendingCheckpoint) thro
 		final long checkpointId = pendingCheckpoint.getCheckpointId();
 		final CompletedCheckpoint completedCheckpoint;
 
+		// As a first step to complete the checkpoint, we register its state with the registry
+		Map<OperatorID, OperatorState> operatorStates = pendingCheckpoint.getOperatorStates();
+		sharedStateRegistry.registerAll(operatorStates.values());
+
 		try {
 			try {
 				// externalize the checkpoint if required
@@ -1003,7 +1011,7 @@ public boolean restoreLatestCheckpointedState(
 			}
 
 			// Recover the checkpoints
-			completedCheckpointStore.recover();
+			completedCheckpointStore.recover(sharedStateRegistry);
 
 			// restore from the latest checkpoint
 			CompletedCheckpoint latest = completedCheckpointStore.getLatestCheckpoint();
@@ -1120,11 +1128,6 @@ public CompletedCheckpointStore getCheckpointStore() {
 		return completedCheckpointStore;
 	}
 
-//	@VisibleForTesting
-//	SharedStateRegistry getSharedStateRegistry() {
-//		return sharedStateRegistry;
-//	}
-
 	public CheckpointIDCounter getCheckpointIdCounter() {
 		return checkpointIdCounter;
 	}

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

@@ -178,7 +178,7 @@ public void discardOnFailedStoring() throws Exception {
 		doDiscard();
 	}
 
-	public boolean discardOnSubsume(SharedStateRegistry sharedStateRegistry) throws Exception {
+	public boolean discardOnSubsume() throws Exception {
 
 		if (props.discardOnSubsumed()) {
 			doDiscard();
@@ -188,7 +188,7 @@ public boolean discardOnSubsume(SharedStateRegistry sharedStateRegistry) throws
 		return false;
 	}
 
-	public boolean discardOnShutdown(JobStatus jobStatus, SharedStateRegistry sharedStateRegistry) throws Exception {
+	public boolean discardOnShutdown(JobStatus jobStatus) throws Exception {
 
 		if (jobStatus == JobStatus.FINISHED && props.discardOnJobFinished() ||
 				jobStatus == JobStatus.CANCELED && props.discardOnJobCancelled() ||
@@ -290,7 +290,7 @@ void setDiscardCallback(@Nullable CompletedCheckpointStats.DiscardCallback disca
 	 *
 	 * @param sharedStateRegistry The registry where shared states are registered
 	 */
-	public void registerSharedStates(SharedStateRegistry sharedStateRegistry) {
+	public void registerSharedStatesAfterRestored(SharedStateRegistry sharedStateRegistry) {
 		sharedStateRegistry.registerAll(operatorStates.values());
 	}
 

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

@@ -19,6 +19,7 @@
 package org.apache.flink.runtime.checkpoint;
 
 import org.apache.flink.runtime.jobgraph.JobStatus;
+import org.apache.flink.runtime.state.SharedStateRegistry;
 
 import java.util.List;
 
@@ -32,8 +33,10 @@ public interface CompletedCheckpointStore {
 	 *
 	 * <p>After a call to this method, {@link #getLatestCheckpoint()} returns the latest
 	 * available checkpoint.
+	 *
+	 * @param sharedStateRegistry the shared state registry to register recovered states.
 	 */
-	void recover() throws Exception;
+	void recover(SharedStateRegistry sharedStateRegistry) throws Exception;
 
 	/**
 	 * Adds a {@link CompletedCheckpoint} instance to the list of completed checkpoints.

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

@@ -35,7 +35,6 @@
 import org.apache.flink.runtime.state.filesystem.FileStateHandle;
 import org.apache.flink.util.ExceptionUtils;
 import org.apache.flink.util.Preconditions;
-
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 

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

@@ -20,6 +20,7 @@
 
 import org.apache.flink.runtime.jobgraph.JobStatus;
 import org.apache.flink.runtime.jobmanager.HighAvailabilityMode;
+import org.apache.flink.runtime.state.SharedStateRegistry;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,7 +33,7 @@
 /**
  * {@link CompletedCheckpointStore} for JobManagers running in {@link HighAvailabilityMode#NONE}.
  */
-public class StandaloneCompletedCheckpointStore extends AbstractCompletedCheckpointStore {
+public class StandaloneCompletedCheckpointStore implements CompletedCheckpointStore {
 
 	private static final Logger LOG = LoggerFactory.getLogger(StandaloneCompletedCheckpointStore.class);
 
@@ -56,21 +57,19 @@ public StandaloneCompletedCheckpointStore(int maxNumberOfCheckpointsToRetain) {
 	}
 
 	@Override
-	public void recover() throws Exception {
+	public void recover(SharedStateRegistry sharedStateRegistry) throws Exception {
 		// Nothing to do
 	}
 
 	@Override
 	public void addCheckpoint(CompletedCheckpoint checkpoint) throws Exception {
-
-		checkpoint.registerSharedStates(sharedStateRegistry);
 
 		checkpoints.addLast(checkpoint);
 
 		if (checkpoints.size() > maxNumberOfCheckpointsToRetain) {
 			try {
 				CompletedCheckpoint checkpointToSubsume = checkpoints.removeFirst();
-				checkpointToSubsume.discardOnSubsume(sharedStateRegistry);
+				checkpointToSubsume.discardOnSubsume();
 			} catch (Exception e) {
 				LOG.warn("Fail to subsume the old checkpoint.", e);
 			}
@@ -103,7 +102,7 @@ public void shutdown(JobStatus jobStatus) throws Exception {
 			LOG.info("Shutting down");
 
 			for (CompletedCheckpoint checkpoint : checkpoints) {
-				checkpoint.discardOnShutdown(jobStatus, sharedStateRegistry);
+				checkpoint.discardOnShutdown(jobStatus);
 			}
 		} finally {
 			checkpoints.clear();

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

@@ -64,7 +64,7 @@
  * checkpoints is consistent. Currently, after recovery we start out with only a single
  * checkpoint to circumvent those situations.
  */
-public class ZooKeeperCompletedCheckpointStore extends AbstractCompletedCheckpointStore {
+public class ZooKeeperCompletedCheckpointStore implements CompletedCheckpointStore {
 
 	private static final Logger LOG = LoggerFactory.getLogger(ZooKeeperCompletedCheckpointStore.class);
 
@@ -102,8 +102,6 @@ public ZooKeeperCompletedCheckpointStore(
 			RetrievableStateStorageHelper<CompletedCheckpoint> stateStorage,
 			Executor executor) throws Exception {
 
-		super(executor);
-
 		checkArgument(maxNumberOfCheckpointsToRetain >= 1, "Must retain at least one checkpoint.");
 		checkNotNull(stateStorage, "State storage");
 
@@ -139,13 +137,14 @@ public boolean requiresExternalizedCheckpoints() {
 	 * that the history of checkpoints is consistent.
 	 */
 	@Override
-	public void recover() throws Exception {
+	public void recover(SharedStateRegistry sharedStateRegistry) throws Exception {
 		LOG.info("Recovering checkpoints from ZooKeeper.");
 
 		// Clear local handles in order to prevent duplicates on
 		// recovery. The local handles should reflect the state
 		// of ZooKeeper.
 		completedCheckpoints.clear();
+		sharedStateRegistry.clear();
 
 		// Get all there is first
 		List<Tuple2<RetrievableStateHandle<CompletedCheckpoint>, String>> initialCheckpoints;
@@ -171,7 +170,7 @@ public void recover() throws Exception {
 				completedCheckpoint = retrieveCompletedCheckpoint(checkpointStateHandle);
 				if (completedCheckpoint != null) {
 					// Re-register all shared states in the checkpoint.
-					completedCheckpoint.registerSharedStates(sharedStateRegistry);
+					completedCheckpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
 					completedCheckpoints.add(completedCheckpoint);
 				}
 			} catch (Exception e) {
@@ -195,9 +194,6 @@ public void addCheckpoint(final CompletedCheckpoint checkpoint) throws Exception
 
 		final String path = checkpointIdToPath(checkpoint.getCheckpointID());
 
-		// First, register all shared states in the checkpoint to consolidates placeholder.
-		checkpoint.registerSharedStates(sharedStateRegistry);
-
 		// Now add the new one. If it fails, we don't want to loose existing data.
 		checkpointsInZooKeeper.addAndLock(path, checkpoint);
 
@@ -206,7 +202,7 @@ public void addCheckpoint(final CompletedCheckpoint checkpoint) throws Exception
 		// Everything worked, let's remove a previous checkpoint if necessary.
 		while (completedCheckpoints.size() > maxNumberOfCheckpointsToRetain) {
 			try {
-				removeSubsumed(completedCheckpoints.removeFirst(), sharedStateRegistry);
+				removeSubsumed(completedCheckpoints.removeFirst());
 			} catch (Exception e) {
 				LOG.warn("Failed to subsume the old checkpoint", e);
 			}
@@ -248,7 +244,7 @@ public void shutdown(JobStatus jobStatus) throws Exception {
 
 			for (CompletedCheckpoint checkpoint : completedCheckpoints) {
 				try {
-					removeShutdown(checkpoint, jobStatus, sharedStateRegistry);
+					removeShutdown(checkpoint, jobStatus);
 				} catch (Exception e) {
 					LOG.error("Failed to discard checkpoint.", e);
 				}
@@ -274,8 +270,7 @@ public void shutdown(JobStatus jobStatus) throws Exception {
 	// ------------------------------------------------------------------------
 
 	private void removeSubsumed(
-		final CompletedCheckpoint completedCheckpoint,
-		final SharedStateRegistry sharedStateRegistry) throws Exception {
+		final CompletedCheckpoint completedCheckpoint) throws Exception {
 
 		if(completedCheckpoint == null) {
 			return;
@@ -287,7 +282,7 @@ private void removeSubsumed(
 				public void apply(@Nullable RetrievableStateHandle<CompletedCheckpoint> value) throws FlinkException {
 					if (value != null) {
 						try {
-							completedCheckpoint.discardOnSubsume(sharedStateRegistry);
+							completedCheckpoint.discardOnSubsume();
 						} catch (Exception e) {
 							throw new FlinkException("Could not discard the completed checkpoint on subsume.", e);
 						}
@@ -302,8 +297,7 @@ public void apply(@Nullable RetrievableStateHandle<CompletedCheckpoint> value) t
 
 	private void removeShutdown(
 			final CompletedCheckpoint completedCheckpoint,
-			final JobStatus jobStatus,
-			final SharedStateRegistry sharedStateRegistry) throws Exception {
+			final JobStatus jobStatus) throws Exception {
 
 		if(completedCheckpoint == null) {
 			return;
@@ -313,7 +307,7 @@ private void removeShutdown(
 			@Override
 			public void apply(@Nullable RetrievableStateHandle<CompletedCheckpoint> value) throws FlinkException {
 				try {
-					completedCheckpoint.discardOnShutdown(jobStatus, sharedStateRegistry);
+					completedCheckpoint.discardOnShutdown(jobStatus);
 				} catch (Exception e) {
 					throw new FlinkException("Could not discard the completed checkpoint on subsume.", e);
 				}

flink-runtime/src/main/java/org/apache/flink/runtime/state/SharedStateRegistry.java

@@ -29,10 +29,13 @@
 import java.util.concurrent.Executor;
 
 /**
+ * This registry manages state that is shared across (incremental) checkpoints, and is responsible
+ * for deleting shared state that is no longer used in any valid checkpoint.
+ *
  * A {@code SharedStateRegistry} will be deployed in the 
- * {@link org.apache.flink.runtime.checkpoint.CompletedCheckpointStore} to
- * maintain the reference count of {@link StreamStateHandle}s which are shared
- * among different incremental checkpoints.
+ * {@link org.apache.flink.runtime.checkpoint.CheckpointCoordinator} to
+ * maintain the reference count of {@link StreamStateHandle}s by a key that (logically) identifies
+ * them.
  */
 public class SharedStateRegistry {
 
@@ -247,4 +250,11 @@ public void run() {
 			}
 		}
 	}
+
+	/**
+	 * Clears the registry.
+	 */
+	public void clear() {
+		registeredStates.clear();
+	}
 }

flink-runtime/src/test/java/org/apache/flink/runtime/checkpoint/CheckpointCoordinatorFailureTest.java

@@ -28,6 +28,7 @@
 import org.apache.flink.runtime.state.ChainedStateHandle;
 import org.apache.flink.runtime.state.KeyedStateHandle;
 import org.apache.flink.runtime.state.OperatorStateHandle;
+import org.apache.flink.runtime.state.SharedStateRegistry;
 import org.apache.flink.runtime.state.StreamStateHandle;
 import org.apache.flink.util.TestLogger;
 import org.junit.Test;
@@ -135,7 +136,7 @@ public void testFailingCompletedCheckpointStoreAdd() throws Exception {
 	private static final class FailingCompletedCheckpointStore implements CompletedCheckpointStore {
 
 		@Override
-		public void recover() throws Exception {
+		public void recover(SharedStateRegistry sharedStateRegistry) throws Exception {
 			throw new UnsupportedOperationException("Not implemented.");
 		}