org.apache.flink.runtime.state.testutils.TestCompletedCheckpointStorageLocation Java Examples

@Test
public void testRegisterStatesAtRegistry() {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
			new TestCompletedCheckpointStorageLocation());

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);
}
 

/**
 * Tests that the stats callbacks happen if the callback is registered.
 */
@Test
public void testCompletedCheckpointStatsCallbacks() throws Exception {
	CompletedCheckpoint completed = new CompletedCheckpoint(
		new JobID(),
		0,
		0,
		1,
		Collections.emptyMap(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpointStats.DiscardCallback callback = mock(CompletedCheckpointStats.DiscardCallback.class);
	completed.setDiscardCallback(callback);

	completed.discardOnShutdown(JobStatus.FINISHED);
	verify(callback, times(1)).notifyDiscardedCheckpoint();
}
 

/**
 * Tests that the stats callbacks happen if the callback is registered.
 */
@Test
public void testCompletedCheckpointStatsCallbacks() throws Exception {
	CompletedCheckpoint completed = new CompletedCheckpoint(
		new JobID(),
		0,
		0,
		1,
		Collections.emptyMap(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpointStats.DiscardCallback callback = mock(CompletedCheckpointStats.DiscardCallback.class);
	completed.setDiscardCallback(callback);

	completed.discardOnShutdown(JobStatus.FINISHED);
	verify(callback, times(1)).notifyDiscardedCheckpoint();
}
 

/**
 * Tests that the stats callbacks happen if the callback is registered.
 */
@Test
public void testCompletedCheckpointStatsCallbacks() throws Exception {
	CompletedCheckpoint completed = new CompletedCheckpoint(
		new JobID(),
		0,
		0,
		1,
		Collections.emptyMap(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpointStats.DiscardCallback callback = mock(CompletedCheckpointStats.DiscardCallback.class);
	completed.setDiscardCallback(callback);

	completed.discardOnShutdown(JobStatus.FINISHED);
	verify(callback, times(1)).notifyDiscardedCheckpoint();
}
 

@Test
public void testRegisterStatesAtRegistry() {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
			new TestCompletedCheckpointStorageLocation());

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);
}
 

@Test
public void testRegisterStatesAtRegistry() {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
			new TestCompletedCheckpointStorageLocation());

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);
}
 

@Test
public void testCompareCheckpointsWithDifferentOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

@Test
public void testCompareCheckpointsWithSameOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);

	assertTrue(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

@Test
public void testCompareCheckpointsWithDifferentOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Tests that the restore callbacks are called if registered.
 */
@Test
public void testCheckpointStatsTrackerRestoreCallback() throws Exception {
	StandaloneCompletedCheckpointStore store = new StandaloneCompletedCheckpointStore(1);

	// set up the coordinator and validate the initial state
	CheckpointCoordinator coord =
		new CheckpointCoordinatorBuilder()
			.setCompletedCheckpointStore(store)
			.setTimer(manuallyTriggeredScheduledExecutor)
			.build();

	store.addCheckpoint(new CompletedCheckpoint(
		new JobID(),
		0,
		0,
		0,
		Collections.<OperatorID, OperatorState>emptyMap(),
		Collections.<MasterState>emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
		new TestCompletedCheckpointStorageLocation()));

	CheckpointStatsTracker tracker = mock(CheckpointStatsTracker.class);
	coord.setCheckpointStatsTracker(tracker);

	assertTrue(coord.restoreLatestCheckpointedStateToAll(Collections.emptySet(), true));

	verify(tracker, times(1))
		.reportRestoredCheckpoint(any(RestoredCheckpointStats.class));
}
 

public TestCompletedCheckpoint(
		JobID jobId,
		long checkpointId,
		long timestamp,
		Map<OperatorID, OperatorState> operatorGroupState,
		CheckpointProperties props) {

	super(jobId, checkpointId, timestamp, Long.MAX_VALUE, operatorGroupState, null, props,
			new TestCompletedCheckpointStorageLocation());
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameJobID() {
	JobID jobID = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID, 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameCheckpointId() {
	JobID jobID1 = new JobID();
	JobID jobID2 = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID1, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID2, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Tests that the garbage collection properties are respected when subsuming checkpoints.
 */
@Test
public void testCleanUpOnSubsume() throws Exception {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	EmptyStreamStateHandle metadata = new EmptyStreamStateHandle();
	TestCompletedCheckpointStorageLocation location =
			new TestCompletedCheckpointStorageLocation(metadata, "ptr");

	CheckpointProperties props = new CheckpointProperties(false, CheckpointType.CHECKPOINT, true, false, false, false, false);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			props,
			location);

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);

	// Subsume
	checkpoint.discardOnSubsume();

	verify(state, times(1)).discardState();
	assertTrue(location.isDisposed());
	assertTrue(metadata.isDisposed());
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameCheckpointId() {
	JobID jobID1 = new JobID();
	JobID jobID2 = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID1, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID2, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameJobID() {
	JobID jobID = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID, 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

@Test
public void testCompareCheckpointsWithSameOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);

	assertTrue(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Tests that the garbage collection properties are respected when subsuming checkpoints.
 */
@Test
public void testCleanUpOnSubsume() throws Exception {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	EmptyStreamStateHandle metadata = new EmptyStreamStateHandle();
	TestCompletedCheckpointStorageLocation location =
			new TestCompletedCheckpointStorageLocation(metadata, "ptr");

	CheckpointProperties props = new CheckpointProperties(false, CheckpointType.CHECKPOINT, true, false, false, false, false);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			props,
			location);

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);

	// Subsume
	checkpoint.discardOnSubsume();

	verify(state, times(1)).discardState();
	assertTrue(location.isDisposed());
	assertTrue(metadata.isDisposed());
}
 

public TestCompletedCheckpoint(
		JobID jobId,
		long checkpointId,
		long timestamp,
		Map<OperatorID, OperatorState> operatorGroupState,
		CheckpointProperties props) {

	super(jobId, checkpointId, timestamp, Long.MAX_VALUE, operatorGroupState, null, props,
			new TestCompletedCheckpointStorageLocation());
}
 

/**
 * Tests that the garbage collection properties are respected when subsuming checkpoints.
 */
@Test
public void testCleanUpOnSubsume() throws Exception {
	OperatorState state = mock(OperatorState.class);
	Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
	operatorStates.put(new OperatorID(), state);

	EmptyStreamStateHandle metadata = new EmptyStreamStateHandle();
	TestCompletedCheckpointStorageLocation location =
			new TestCompletedCheckpointStorageLocation(metadata, "ptr");

	CheckpointProperties props = new CheckpointProperties(false, CheckpointType.CHECKPOINT, true, false, false, false, false);

	CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			new JobID(), 0, 0, 1,
			operatorStates,
			Collections.emptyList(),
			props,
			location);

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	checkpoint.registerSharedStatesAfterRestored(sharedStateRegistry);
	verify(state, times(1)).registerSharedStates(sharedStateRegistry);

	// Subsume
	checkpoint.discardOnSubsume();

	verify(state, times(1)).discardState();
	assertTrue(location.isDisposed());
	assertTrue(metadata.isDisposed());
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameCheckpointId() {
	JobID jobID1 = new JobID();
	JobID jobID2 = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID1, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID2, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

/**
 * Verify that both JobID and checkpoint id are taken into account when comparing.
 */
@Test
public void testCompareCheckpointsWithSameJobID() {
	JobID jobID = new JobID();

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		jobID, 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		jobID, 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

@Test
public void testCompareCheckpointsWithSameOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);

	assertTrue(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

@Test
public void testCompareCheckpointsWithDifferentOrder() {

	CompletedCheckpoint checkpoint1 = new CompletedCheckpoint(
		new JobID(), 0, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	CompletedCheckpoint checkpoint2 = new CompletedCheckpoint(
		new JobID(), 1, 0, 1,
		new HashMap<>(),
		Collections.emptyList(),
		CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.RETAIN_ON_FAILURE),
		new TestCompletedCheckpointStorageLocation());

	List<CompletedCheckpoint> checkpoints1= new ArrayList<>();
	checkpoints1.add(checkpoint1);
	checkpoints1.add(checkpoint2);
	checkpoints1.add(checkpoint1);

	List<CompletedCheckpoint> checkpoints2 = new ArrayList<>();
	checkpoints2.add(checkpoint2);
	checkpoints2.add(checkpoint1);
	checkpoints2.add(checkpoint2);

	assertFalse(CompletedCheckpoint.checkpointsMatch(checkpoints1, checkpoints2));
}
 

public TestCompletedCheckpoint(
		JobID jobId,
		long checkpointId,
		long timestamp,
		Map<OperatorID, OperatorState> operatorGroupState,
		CheckpointProperties props) {

	super(jobId, checkpointId, timestamp, Long.MAX_VALUE, operatorGroupState, null, props,
			new TestCompletedCheckpointStorageLocation());
}
 

@Test
public void checkUnMatchedStateOnRestore() throws Exception {
	final String id1 = "id1";
	final String id2 = "id2";

	final String state1 = "the-test-string-state";
	final byte[] state1serialized = new StringSerializer().serialize(state1);

	final long state2 = 987654321L;
	final byte[] state2serialized = new LongSerializer().serialize(state2);

	final List<MasterState> masterHookStates = Arrays.asList(
			new MasterState(id1, state1serialized, StringSerializer.VERSION),
			new MasterState(id2, state2serialized, LongSerializer.VERSION));

	final MasterTriggerRestoreHook<String> statefulHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook.getIdentifier()).thenReturn(id1);
	when(statefulHook.createCheckpointDataSerializer()).thenReturn(new StringSerializer());
	when(statefulHook.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Void> statelessHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statelessHook.getIdentifier()).thenReturn("some-id");

	final JobID jid = new JobID();
	final long checkpointId = 44L;

	final CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			jid, checkpointId, 123L, 125L,
			Collections.<OperatorID, OperatorState>emptyMap(),
			masterHookStates,
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
			new TestCompletedCheckpointStorageLocation());

	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(jid, ackVertex);

	cc.addMasterHook(statefulHook);
	cc.addMasterHook(statelessHook);

	cc.getCheckpointStore().addCheckpoint(checkpoint);

	// since we have unmatched state, this should fail
	try {
		cc.restoreLatestCheckpointedStateToAll(
				Collections.emptySet(),
				false);
		fail("exception expected");
	}
	catch (IllegalStateException ignored) {}

	// permitting unmatched state should succeed
	cc.restoreLatestCheckpointedStateToAll(
			Collections.emptySet(),
			true);

	verify(statefulHook, times(1)).restoreCheckpoint(eq(checkpointId), eq(state1));
	verify(statelessHook, times(1)).restoreCheckpoint(eq(checkpointId), isNull(Void.class));
}
 

/**
 * Tests that the garbage collection properties are respected when shutting down.
 */
@Test
public void testCleanUpOnShutdown() throws Exception {
	JobStatus[] terminalStates = new JobStatus[] {
			JobStatus.FINISHED, JobStatus.CANCELED, JobStatus.FAILED, JobStatus.SUSPENDED
	};

	for (JobStatus status : terminalStates) {

		OperatorState state = mock(OperatorState.class);
		Map<OperatorID, OperatorState> operatorStates = new HashMap<>();
		operatorStates.put(new OperatorID(), state);

		EmptyStreamStateHandle retainedHandle = new EmptyStreamStateHandle();
		TestCompletedCheckpointStorageLocation retainedLocation =
				new TestCompletedCheckpointStorageLocation(retainedHandle, "ptr");

		// Keep
		CheckpointProperties retainProps = new CheckpointProperties(false, CheckpointType.CHECKPOINT, false, false, false, false, false);
		CompletedCheckpoint checkpoint = new CompletedCheckpoint(
				new JobID(), 0, 0, 1,
				new HashMap<>(operatorStates),
				Collections.emptyList(),
				retainProps,
				retainedLocation);

		checkpoint.discardOnShutdown(status);

		verify(state, times(0)).discardState();
		assertFalse(retainedLocation.isDisposed());
		assertFalse(retainedHandle.isDisposed());

		// Discard
		EmptyStreamStateHandle discardHandle = new EmptyStreamStateHandle();
		TestCompletedCheckpointStorageLocation discardLocation =
				new TestCompletedCheckpointStorageLocation(discardHandle, "ptr");

		// Keep
		CheckpointProperties discardProps = new CheckpointProperties(false, CheckpointType.CHECKPOINT, true, true, true, true, true);

		checkpoint = new CompletedCheckpoint(
				new JobID(), 0, 0, 1,
				new HashMap<>(operatorStates),
				Collections.emptyList(),
				discardProps,
				discardLocation);

		checkpoint.discardOnShutdown(status);

		verify(state, times(1)).discardState();
		assertTrue(discardLocation.isDisposed());
		assertTrue(discardHandle.isDisposed());
	}
}
 

@Test
public void testHooksAreCalledOnRestore() throws Exception {
	final String id1 = "id1";
	final String id2 = "id2";

	final String state1 = "the-test-string-state";
	final byte[] state1serialized = new StringSerializer().serialize(state1);

	final long state2 = 987654321L;
	final byte[] state2serialized = new LongSerializer().serialize(state2);

	final List<MasterState> masterHookStates = Arrays.asList(
			new MasterState(id1, state1serialized, StringSerializer.VERSION),
			new MasterState(id2, state2serialized, LongSerializer.VERSION));

	final MasterTriggerRestoreHook<String> statefulHook1 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook1.getIdentifier()).thenReturn(id1);
	when(statefulHook1.createCheckpointDataSerializer()).thenReturn(new StringSerializer());
	when(statefulHook1.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Long> statefulHook2 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook2.getIdentifier()).thenReturn(id2);
	when(statefulHook2.createCheckpointDataSerializer()).thenReturn(new LongSerializer());
	when(statefulHook2.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Void> statelessHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statelessHook.getIdentifier()).thenReturn("some-id");

	final JobID jid = new JobID();
	final long checkpointId = 13L;

	final CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			jid, checkpointId, 123L, 125L,
			Collections.<OperatorID, OperatorState>emptyMap(),
			masterHookStates,
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
			new TestCompletedCheckpointStorageLocation());
	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(jid, ackVertex);

	cc.addMasterHook(statefulHook1);
	cc.addMasterHook(statelessHook);
	cc.addMasterHook(statefulHook2);

	cc.getCheckpointStore().addCheckpoint(checkpoint);
	cc.restoreLatestCheckpointedStateToAll(
			Collections.emptySet(),
			false);

	verify(statefulHook1, times(1)).restoreCheckpoint(eq(checkpointId), eq(state1));
	verify(statefulHook2, times(1)).restoreCheckpoint(eq(checkpointId), eq(state2));
	verify(statelessHook, times(1)).restoreCheckpoint(eq(checkpointId), isNull(Void.class));
}
 

@Test
public void checkUnMatchedStateOnRestore() throws Exception {
	final String id1 = "id1";
	final String id2 = "id2";

	final String state1 = "the-test-string-state";
	final byte[] state1serialized = new StringSerializer().serialize(state1);

	final long state2 = 987654321L;
	final byte[] state2serialized = new LongSerializer().serialize(state2);

	final List<MasterState> masterHookStates = Arrays.asList(
			new MasterState(id1, state1serialized, StringSerializer.VERSION),
			new MasterState(id2, state2serialized, LongSerializer.VERSION));

	final MasterTriggerRestoreHook<String> statefulHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook.getIdentifier()).thenReturn(id1);
	when(statefulHook.createCheckpointDataSerializer()).thenReturn(new StringSerializer());
	when(statefulHook.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Void> statelessHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statelessHook.getIdentifier()).thenReturn("some-id");

	final JobID jid = new JobID();
	final long checkpointId = 44L;

	final CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			jid, checkpointId, 123L, 125L,
			Collections.<OperatorID, OperatorState>emptyMap(),
			masterHookStates,
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
			new TestCompletedCheckpointStorageLocation());

	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(jid, ackVertex);

	cc.addMasterHook(statefulHook);
	cc.addMasterHook(statelessHook);

	cc.getCheckpointStore().addCheckpoint(checkpoint);

	// since we have unmatched state, this should fail
	try {
		cc.restoreLatestCheckpointedState(
				Collections.<JobVertexID, ExecutionJobVertex>emptyMap(),
				true,
				false);
		fail("exception expected");
	}
	catch (IllegalStateException ignored) {}

	// permitting unmatched state should succeed
	cc.restoreLatestCheckpointedState(
			Collections.<JobVertexID, ExecutionJobVertex>emptyMap(),
			true,
			true);

	verify(statefulHook, times(1)).restoreCheckpoint(eq(checkpointId), eq(state1));
	verify(statelessHook, times(1)).restoreCheckpoint(eq(checkpointId), isNull(Void.class));
}
 

@Test
public void testHooksAreCalledOnRestore() throws Exception {
	final String id1 = "id1";
	final String id2 = "id2";

	final String state1 = "the-test-string-state";
	final byte[] state1serialized = new StringSerializer().serialize(state1);

	final long state2 = 987654321L;
	final byte[] state2serialized = new LongSerializer().serialize(state2);

	final List<MasterState> masterHookStates = Arrays.asList(
			new MasterState(id1, state1serialized, StringSerializer.VERSION),
			new MasterState(id2, state2serialized, LongSerializer.VERSION));

	final MasterTriggerRestoreHook<String> statefulHook1 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook1.getIdentifier()).thenReturn(id1);
	when(statefulHook1.createCheckpointDataSerializer()).thenReturn(new StringSerializer());
	when(statefulHook1.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Long> statefulHook2 = mockGeneric(MasterTriggerRestoreHook.class);
	when(statefulHook2.getIdentifier()).thenReturn(id2);
	when(statefulHook2.createCheckpointDataSerializer()).thenReturn(new LongSerializer());
	when(statefulHook2.triggerCheckpoint(anyLong(), anyLong(), any(Executor.class)))
			.thenThrow(new Exception("not expected"));

	final MasterTriggerRestoreHook<Void> statelessHook = mockGeneric(MasterTriggerRestoreHook.class);
	when(statelessHook.getIdentifier()).thenReturn("some-id");

	final JobID jid = new JobID();
	final long checkpointId = 13L;

	final CompletedCheckpoint checkpoint = new CompletedCheckpoint(
			jid, checkpointId, 123L, 125L,
			Collections.<OperatorID, OperatorState>emptyMap(),
			masterHookStates,
			CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION),
			new TestCompletedCheckpointStorageLocation());
	final ExecutionAttemptID execId = new ExecutionAttemptID();
	final ExecutionVertex ackVertex = mockExecutionVertex(execId);
	final CheckpointCoordinator cc = instantiateCheckpointCoordinator(jid, ackVertex);

	cc.addMasterHook(statefulHook1);
	cc.addMasterHook(statelessHook);
	cc.addMasterHook(statefulHook2);

	cc.getCheckpointStore().addCheckpoint(checkpoint);
	cc.restoreLatestCheckpointedState(
			Collections.<JobVertexID, ExecutionJobVertex>emptyMap(),
			true,
			false);

	verify(statefulHook1, times(1)).restoreCheckpoint(eq(checkpointId), eq(state1));
	verify(statefulHook2, times(1)).restoreCheckpoint(eq(checkpointId), eq(state2));
	verify(statelessHook, times(1)).restoreCheckpoint(eq(checkpointId), isNull(Void.class));
}