SavepointCoordinator extends CheckpointCoordinator { public void restoreSavepoint( Map<JobVertexID, ExecutionJobVertex> tasks, String savepointPath) throws Exception { checkNotNull(savepointPath, "Savepoint path"); synchronized (lock) { if (isShutdown()) { throw new IllegalStateException("CheckpointCoordinator is shut down"); } long recoveryTimestamp = System.currentTimeMillis(); LOG.info("Rolling back to savepoint '{}'.", savepointPath); CompletedCheckpoint checkpoint = savepointStore.getState(savepointPath); LOG.info("Savepoint: {}@{}", checkpoint.getCheckpointID(), checkpoint.getTimestamp()); LOG.debug("Rolling back individual operators."); for (StateForTask state : checkpoint.getStates()) { LOG.debug("Rolling back subtask {} of operator {}.", state.getSubtask(), state.getOperatorId()); ExecutionJobVertex vertex = tasks.get(state.getOperatorId()); if (vertex == null) { String msg = String.format("Failed to rollback to savepoint %s. " + "Cannot map old state for task %s to the new program. " + "This indicates that the program has been changed in a " + "non-compatible way after the savepoint.", checkpoint, state.getOperatorId()); throw new IllegalStateException(msg); } if (state.getSubtask() >= vertex.getParallelism()) { String msg = String.format("Failed to rollback to savepoint %s. " + "Parallelism mismatch between savepoint state and new program. " + "Cannot map subtask %d of operator %s to new program with " + "parallelism %d. This indicates that the program has been changed " + "in a non-compatible way after the savepoint.", checkpoint, state.getSubtask(), state.getOperatorId(), vertex.getParallelism()); throw new IllegalStateException(msg); } Execution exec = vertex.getTaskVertices()[state.getSubtask()] .getCurrentExecutionAttempt(); exec.setInitialState(state.getState(), recoveryTimestamp); } long nextCheckpointId = checkpoint.getCheckpointID(); checkpointIdCounter.start(); checkpointIdCounter.setCount(nextCheckpointId + 1); LOG.info("Reset the checkpoint ID to {}", nextCheckpointId); if (savepointRestorePath == null) { savepointRestorePath = savepointPath; } } } SavepointCoordinator(
			JobID jobId,
			long baseInterval,
			long checkpointTimeout,
			ExecutionVertex[] tasksToTrigger,
			ExecutionVertex[] tasksToWaitFor,
			ExecutionVertex[] tasksToCommitTo,
			ClassLoader userClassLoader,
			CheckpointIDCounter checkpointIDCounter,
			StateStore<CompletedCheckpoint> savepointStore,
			CheckpointStatsTracker statsTracker); String getSavepointRestorePath(); Future<String> triggerSavepoint(long timestamp); void restoreSavepoint(
			Map<JobVertexID, ExecutionJobVertex> tasks,
			String savepointPath); @Override ActorGateway createActivatorDeactivator(
			ActorSystem actorSystem,
			UUID leaderSessionID);  }
@Test public void testRollbackSetsCheckpointID() throws Exception { CompletedCheckpoint savepoint = mock(CompletedCheckpoint.class); when(savepoint.getStates()).thenReturn(Collections.<StateForTask>emptyList()); when(savepoint.getCheckpointID()).thenReturn(12312312L); CheckpointIDCounter checkpointIdCounter = mock(CheckpointIDCounter.class); StateStore<CompletedCheckpoint> savepointStore = mock(StateStore.class); when(savepointStore.getState(anyString())).thenReturn(savepoint); SavepointCoordinator coordinator = createSavepointCoordinator( new JobID(), 60 * 1000, new ExecutionVertex[] {}, new ExecutionVertex[] {}, new ExecutionVertex[] {}, checkpointIdCounter, savepointStore); coordinator.restoreSavepoint(createExecutionJobVertexMap(), "any"); verify(checkpointIdCounter).setCount(eq(12312312L + 1)); coordinator.shutdown(); }