org.apache.flink.runtime.state.SharedStateRegistry Java Examples

/**
 * Tests adding and getting a checkpoint.
 */
@Test
public void testAddAndGetLatestCheckpoint() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);
	
	// Empty state
	assertEquals(0, checkpoints.getNumberOfRetainedCheckpoints());
	assertEquals(0, checkpoints.getAllCheckpoints().size());

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry) };

	// Add and get latest
	checkpoints.addCheckpoint(expected[0]);
	assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[0], checkpoints.getLatestCheckpoint(false));

	checkpoints.addCheckpoint(expected[1]);
	assertEquals(2, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[1], checkpoints.getLatestCheckpoint(false));
}
 

/**
 * Tests that all added checkpoints are returned.
 */
@Test
public void testGetAllCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry),
			createCheckpoint(2, sharedStateRegistry), createCheckpoint(3, sharedStateRegistry)
	};

	for (TestCompletedCheckpoint checkpoint : expected) {
		checkpoints.addCheckpoint(checkpoint);
	}

	List<CompletedCheckpoint> actual = checkpoints.getAllCheckpoints();

	assertEquals(expected.length, actual.size());

	for (int i = 0; i < expected.length; i++) {
		assertEquals(expected[i], actual.get(i));
	}
}
 

/**
 * FLINK-6284
 *
 * Tests that the latest recovered checkpoint is the one with the highest checkpoint id
 */
@Test
public void testLatestCheckpointRecovery() throws Exception {
	final int numCheckpoints = 3;
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpointStore = createCompletedCheckpoints(numCheckpoints);
	List<CompletedCheckpoint> checkpoints = new ArrayList<>(numCheckpoints);

	checkpoints.add(createCheckpoint(9, sharedStateRegistry));
	checkpoints.add(createCheckpoint(10, sharedStateRegistry));
	checkpoints.add(createCheckpoint(11, sharedStateRegistry));

	for (CompletedCheckpoint checkpoint : checkpoints) {
		checkpointStore.addCheckpoint(checkpoint);
	}

	sharedStateRegistry.close();
	checkpointStore.recover();

	CompletedCheckpoint latestCheckpoint = checkpointStore.getLatestCheckpoint(false);

	assertEquals(checkpoints.get(checkpoints.size() -1), latestCheckpoint);
}
 

/**
 * FLINK-6284
 *
 * Tests that the latest recovered checkpoint is the one with the highest checkpoint id
 */
@Test
public void testLatestCheckpointRecovery() throws Exception {
	final int numCheckpoints = 3;
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpointStore = createCompletedCheckpoints(numCheckpoints);
	List<CompletedCheckpoint> checkpoints = new ArrayList<>(numCheckpoints);

	checkpoints.add(createCheckpoint(9, sharedStateRegistry));
	checkpoints.add(createCheckpoint(10, sharedStateRegistry));
	checkpoints.add(createCheckpoint(11, sharedStateRegistry));

	for (CompletedCheckpoint checkpoint : checkpoints) {
		checkpointStore.addCheckpoint(checkpoint);
	}

	sharedStateRegistry.close();
	checkpointStore.recover();

	CompletedCheckpoint latestCheckpoint = checkpointStore.getLatestCheckpoint(false);

	assertEquals(checkpoints.get(checkpoints.size() -1), latestCheckpoint);
}
 

/**
 * Tests that shutdown discards all checkpoints.
 */
@Test
public void testShutdownDiscardsCheckpoints() throws Exception {
	CuratorFramework client = ZOOKEEPER.getClient();

	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	assertNotNull(client.checkExists().forPath(CHECKPOINT_PATH + ZooKeeperCompletedCheckpointStore.checkpointIdToPath(checkpoint.getCheckpointID())));

	store.shutdown(JobStatus.FINISHED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertNull(client.checkExists().forPath(CHECKPOINT_PATH + ZooKeeperCompletedCheckpointStore.checkpointIdToPath(checkpoint.getCheckpointID())));

	sharedStateRegistry.close();
	store.recover();

	assertEquals(0, store.getNumberOfRetainedCheckpoints());
}
 

/**
 * Tests adding and getting a checkpoint.
 */
@Test
public void testAddAndGetLatestCheckpoint() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);
	
	// Empty state
	assertEquals(0, checkpoints.getNumberOfRetainedCheckpoints());
	assertEquals(0, checkpoints.getAllCheckpoints().size());

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry) };

	// Add and get latest
	checkpoints.addCheckpoint(expected[0]);
	assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[0], checkpoints.getLatestCheckpoint());

	checkpoints.addCheckpoint(expected[1]);
	assertEquals(2, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[1], checkpoints.getLatestCheckpoint());
}
 

/**
 * Tests that suspends discards all checkpoints (as they cannot be
 * recovered later in standalone recovery mode).
 */
@Test
public void testSuspendDiscardsCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);
	Collection<OperatorState> taskStates = checkpoint.getOperatorStates().values();

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	verifyCheckpointRegistered(taskStates, sharedStateRegistry);

	store.shutdown(JobStatus.SUSPENDED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertTrue(checkpoint.isDiscarded());
	verifyCheckpointDiscarded(taskStates);
}
 

/**
 * Tests that all added checkpoints are returned.
 */
@Test
public void testGetAllCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry),
			createCheckpoint(2, sharedStateRegistry), createCheckpoint(3, sharedStateRegistry)
	};

	for (TestCompletedCheckpoint checkpoint : expected) {
		checkpoints.addCheckpoint(checkpoint);
	}

	List<CompletedCheckpoint> actual = checkpoints.getAllCheckpoints();

	assertEquals(expected.length, actual.size());

	for (int i = 0; i < expected.length; i++) {
		assertEquals(expected[i], actual.get(i));
	}
}
 

public static TestCompletedCheckpoint createCheckpoint(
	int id,
	SharedStateRegistry sharedStateRegistry) throws IOException {

	int numberOfStates = 4;
	CheckpointProperties props = CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION);

	OperatorID operatorID = new OperatorID();

	Map<OperatorID, OperatorState> operatorGroupState = new HashMap<>();
	OperatorState operatorState = new OperatorState(operatorID, numberOfStates, numberOfStates);
	operatorGroupState.put(operatorID, operatorState);

	for (int i = 0; i < numberOfStates; i++) {
		OperatorSubtaskState subtaskState =
			new TestOperatorSubtaskState();

		operatorState.putState(i, subtaskState);
	}

	operatorState.registerSharedStates(sharedStateRegistry);

	return new TestCompletedCheckpoint(new JobID(), id, 0, operatorGroupState, props);
}
 

/**
 * Tests that shutdown discards all checkpoints.
 */
@Test
public void testShutdownDiscardsCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);
	Collection<OperatorState> operatorStates = checkpoint.getOperatorStates().values();

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	verifyCheckpointRegistered(operatorStates, sharedStateRegistry);

	store.shutdown(JobStatus.FINISHED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertTrue(checkpoint.isDiscarded());
	verifyCheckpointDiscarded(operatorStates);
}
 

/**
 * Tests that checkpoints are discarded when the completed checkpoint store is shut
 * down with a globally terminal state.
 */
@Test
public void testDiscardingCheckpointsAtShutDown() throws Exception {
	final SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final Configuration configuration = new Configuration();
	configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, zooKeeperResource.getConnectString());

	final CuratorFramework client = ZooKeeperUtils.startCuratorFramework(configuration);
	final ZooKeeperCompletedCheckpointStore checkpointStore = createZooKeeperCheckpointStore(client);

	try {
		final CompletedCheckpointStoreTest.TestCompletedCheckpoint checkpoint1 = CompletedCheckpointStoreTest.createCheckpoint(0, sharedStateRegistry);

		checkpointStore.addCheckpoint(checkpoint1);
		assertThat(checkpointStore.getAllCheckpoints(), Matchers.contains(checkpoint1));

		checkpointStore.shutdown(JobStatus.FINISHED);

		// verify that the checkpoint is discarded
		CompletedCheckpointStoreTest.verifyCheckpointDiscarded(checkpoint1);
	} finally {
		client.close();
	}
}
 

/**
 * Tests that checkpoints are discarded when the completed checkpoint store is shut
 * down with a globally terminal state.
 */
@Test
public void testDiscardingCheckpointsAtShutDown() throws Exception {
	final SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final Configuration configuration = new Configuration();
	configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, zooKeeperResource.getConnectString());

	final CuratorFramework client = ZooKeeperUtils.startCuratorFramework(configuration);
	final ZooKeeperCompletedCheckpointStore checkpointStore = createZooKeeperCheckpointStore(client);

	try {
		final CompletedCheckpointStoreTest.TestCompletedCheckpoint checkpoint1 = CompletedCheckpointStoreTest.createCheckpoint(0, sharedStateRegistry);

		checkpointStore.addCheckpoint(checkpoint1);
		assertThat(checkpointStore.getAllCheckpoints(), Matchers.contains(checkpoint1));

		checkpointStore.shutdown(JobStatus.FINISHED);

		// verify that the checkpoint is discarded
		CompletedCheckpointStoreTest.verifyCheckpointDiscarded(checkpoint1);
	} finally {
		client.close();
	}
}
 

/**
 * Tests that shutdown discards all checkpoints.
 */
@Test
public void testShutdownDiscardsCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);
	Collection<OperatorState> operatorStates = checkpoint.getOperatorStates().values();

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	verifyCheckpointRegistered(operatorStates, sharedStateRegistry);

	store.shutdown(JobStatus.FINISHED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertTrue(checkpoint.isDiscarded());
	verifyCheckpointDiscarded(operatorStates);
}
 

/**
 * Tests that suspends discards all checkpoints (as they cannot be
 * recovered later in standalone recovery mode).
 */
@Test
public void testSuspendDiscardsCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);
	Collection<OperatorState> taskStates = checkpoint.getOperatorStates().values();

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	verifyCheckpointRegistered(taskStates, sharedStateRegistry);

	store.shutdown(JobStatus.SUSPENDED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertTrue(checkpoint.isDiscarded());
	verifyCheckpointDiscarded(taskStates);
}
 

private static CheckpointCoordinator instantiateCheckpointCoordinator(JobID jid, ExecutionVertex... ackVertices) {
	CheckpointCoordinatorConfiguration chkConfig = new CheckpointCoordinatorConfiguration(
		10000000L,
		600000L,
		0L,
		1,
		CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
		true,
		false,
		0);
	return new CheckpointCoordinator(
			jid,
			chkConfig,
			new ExecutionVertex[0],
			ackVertices,
			new ExecutionVertex[0],
			new StandaloneCheckpointIDCounter(),
			new StandaloneCompletedCheckpointStore(10),
			new MemoryStateBackend(),
			Executors.directExecutor(),
			SharedStateRegistry.DEFAULT_FACTORY,
			new CheckpointFailureManager(
				0,
				NoOpFailJobCall.INSTANCE));
}
 

/**
 * Tests that all added checkpoints are returned.
 */
@Test
public void testGetAllCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry),
			createCheckpoint(2, sharedStateRegistry), createCheckpoint(3, sharedStateRegistry)
	};

	for (TestCompletedCheckpoint checkpoint : expected) {
		checkpoints.addCheckpoint(checkpoint);
	}

	List<CompletedCheckpoint> actual = checkpoints.getAllCheckpoints();

	assertEquals(expected.length, actual.size());

	for (int i = 0; i < expected.length; i++) {
		assertEquals(expected[i], actual.get(i));
	}
}
 

/**
 * Tests that checkpoints are discarded when the completed checkpoint store is shut
 * down with a globally terminal state.
 */
@Test
public void testDiscardingCheckpointsAtShutDown() throws Exception {
	final SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final Configuration configuration = new Configuration();
	configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, zooKeeperResource.getConnectString());

	final CuratorFramework client = ZooKeeperUtils.startCuratorFramework(configuration);
	final ZooKeeperCompletedCheckpointStore checkpointStore = createZooKeeperCheckpointStore(client);

	try {
		final CompletedCheckpointStoreTest.TestCompletedCheckpoint checkpoint1 = CompletedCheckpointStoreTest.createCheckpoint(0, sharedStateRegistry);

		checkpointStore.addCheckpoint(checkpoint1);
		assertThat(checkpointStore.getAllCheckpoints(), Matchers.contains(checkpoint1));

		checkpointStore.shutdown(JobStatus.FINISHED);

		// verify that the checkpoint is discarded
		CompletedCheckpointStoreTest.verifyCheckpointDiscarded(checkpoint1);
	} finally {
		client.close();
	}
}
 

/**
 * Tests that adding more checkpoints than retained discards the correct checkpoints (using
 * the correct class loader).
 */
@Test
public void testAddCheckpointMoreThanMaxRetained() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(1);

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry),
			createCheckpoint(2, sharedStateRegistry), createCheckpoint(3, sharedStateRegistry)
	};

	// Add checkpoints
	checkpoints.addCheckpoint(expected[0]);
	assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());

	for (int i = 1; i < expected.length; i++) {
		Collection<OperatorState> taskStates = expected[i - 1].getOperatorStates().values();

		checkpoints.addCheckpoint(expected[i]);

		// The ZooKeeper implementation discards asynchronously
		expected[i - 1].awaitDiscard();
		assertTrue(expected[i - 1].isDiscarded());
		assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());
	}
}
 

/**
 * Tests adding and getting a checkpoint.
 */
@Test
public void testAddAndGetLatestCheckpoint() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore checkpoints = createCompletedCheckpoints(4);
	
	// Empty state
	assertEquals(0, checkpoints.getNumberOfRetainedCheckpoints());
	assertEquals(0, checkpoints.getAllCheckpoints().size());

	TestCompletedCheckpoint[] expected = new TestCompletedCheckpoint[] {
			createCheckpoint(0, sharedStateRegistry), createCheckpoint(1, sharedStateRegistry) };

	// Add and get latest
	checkpoints.addCheckpoint(expected[0]);
	assertEquals(1, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[0], checkpoints.getLatestCheckpoint(false));

	checkpoints.addCheckpoint(expected[1]);
	assertEquals(2, checkpoints.getNumberOfRetainedCheckpoints());
	verifyCheckpoint(expected[1], checkpoints.getLatestCheckpoint(false));
}
 

@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);
}
 

public static TestCompletedCheckpoint createCheckpoint(
	int id,
	SharedStateRegistry sharedStateRegistry) throws IOException {

	int numberOfStates = 4;
	CheckpointProperties props = CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION);

	OperatorID operatorID = new OperatorID();

	Map<OperatorID, OperatorState> operatorGroupState = new HashMap<>();
	OperatorState operatorState = new OperatorState(operatorID, numberOfStates, numberOfStates);
	operatorGroupState.put(operatorID, operatorState);

	for (int i = 0; i < numberOfStates; i++) {
		OperatorSubtaskState subtaskState =
			new TestOperatorSubtaskState();

		operatorState.putState(i, subtaskState);
	}

	operatorState.registerSharedStates(sharedStateRegistry);

	return new TestCompletedCheckpoint(new JobID(), id, 0, operatorGroupState, props);
}
 

/**
 * Tests that shutdown discards all checkpoints.
 */
@Test
public void testShutdownDiscardsCheckpoints() throws Exception {
	SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	CompletedCheckpointStore store = createCompletedCheckpoints(1);
	TestCompletedCheckpoint checkpoint = createCheckpoint(0, sharedStateRegistry);
	Collection<OperatorState> operatorStates = checkpoint.getOperatorStates().values();

	store.addCheckpoint(checkpoint);
	assertEquals(1, store.getNumberOfRetainedCheckpoints());
	verifyCheckpointRegistered(operatorStates, sharedStateRegistry);

	store.shutdown(JobStatus.FINISHED);
	assertEquals(0, store.getNumberOfRetainedCheckpoints());
	assertTrue(checkpoint.isDiscarded());
	verifyCheckpointDiscarded(operatorStates);
}
 

private static CheckpointCoordinator instantiateCheckpointCoordinator(JobID jid, ExecutionVertex... ackVertices) {
	return new CheckpointCoordinator(
			jid,
			10000000L,
			600000L,
			0L,
			1,
			CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
			new ExecutionVertex[0],
			ackVertices,
			new ExecutionVertex[0],
			new StandaloneCheckpointIDCounter(),
			new StandaloneCompletedCheckpointStore(10),
			new MemoryStateBackend(),
			Executors.directExecutor(),
			SharedStateRegistry.DEFAULT_FACTORY);
}
 

/**
 * Tests that subsumed checkpoints are discarded.
 */
@Test
public void testDiscardingSubsumedCheckpoints() throws Exception {
	final SharedStateRegistry sharedStateRegistry = new SharedStateRegistry();
	final Configuration configuration = new Configuration();
	configuration.setString(HighAvailabilityOptions.HA_ZOOKEEPER_QUORUM, zooKeeperResource.getConnectString());

	final CuratorFramework client = ZooKeeperUtils.startCuratorFramework(configuration);
	final ZooKeeperCompletedCheckpointStore checkpointStore = createZooKeeperCheckpointStore(client);

	try {
		final CompletedCheckpointStoreTest.TestCompletedCheckpoint checkpoint1 = CompletedCheckpointStoreTest.createCheckpoint(0, sharedStateRegistry);

		checkpointStore.addCheckpoint(checkpoint1);
		assertThat(checkpointStore.getAllCheckpoints(), Matchers.contains(checkpoint1));

		final CompletedCheckpointStoreTest.TestCompletedCheckpoint checkpoint2 = CompletedCheckpointStoreTest.createCheckpoint(1, sharedStateRegistry);
		checkpointStore.addCheckpoint(checkpoint2);
		final List<CompletedCheckpoint> allCheckpoints = checkpointStore.getAllCheckpoints();
		assertThat(allCheckpoints, Matchers.contains(checkpoint2));
		assertThat(allCheckpoints, Matchers.not(Matchers.contains(checkpoint1)));

		// verify that the subsumed checkpoint is discarded
		CompletedCheckpointStoreTest.verifyCheckpointDiscarded(checkpoint1);
	} finally {
		client.close();
	}
}
 

private CheckpointCoordinator instantiateCheckpointCoordinator(
	JobID jid,
	ScheduledExecutor testingScheduledExecutor,
	ExecutionVertex... ackVertices) {

	CheckpointCoordinatorConfiguration chkConfig = new CheckpointCoordinatorConfiguration(
		10000000L,
		600000L,
		0L,
		1,
		CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
		true,
		false,
		false,
		0);
	return new CheckpointCoordinator(
			jid,
			chkConfig,
			new ExecutionVertex[0],
			ackVertices,
			new ExecutionVertex[0],
			Collections.emptyList(),
			new StandaloneCheckpointIDCounter(),
			new StandaloneCompletedCheckpointStore(10),
			new MemoryStateBackend(),
			Executors.directExecutor(),
			testingScheduledExecutor,
			SharedStateRegistry.DEFAULT_FACTORY,
			new CheckpointFailureManager(
				0,
				NoOpFailJobCall.INSTANCE));
}
 

private CheckpointCoordinator getCheckpointCoordinator(
		final JobID jobId,
		final ExecutionVertex vertex1,
		final ExecutionVertex vertex2,
		final CheckpointFailureManager failureManager) {

	final CheckpointCoordinatorConfiguration chkConfig = new CheckpointCoordinatorConfiguration(
			600000,
			600000,
			0,
			Integer.MAX_VALUE,
			CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
			true,
			false,
			0);

	return new CheckpointCoordinator(
			jobId,
			chkConfig,
			new ExecutionVertex[]{vertex1, vertex2},
			new ExecutionVertex[]{vertex1, vertex2},
			new ExecutionVertex[]{vertex1, vertex2},
			new StandaloneCheckpointIDCounter(),
			new StandaloneCompletedCheckpointStore(1),
			new MemoryStateBackend(),
			Executors.directExecutor(),
			SharedStateRegistry.DEFAULT_FACTORY,
			failureManager);
}
 

/**
 * 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());
}
 

/**
 * 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());
}
 

@Override
public void registerSharedStates(SharedStateRegistry sharedStateRegistry) {
	if (managedKeyedState != null) {
		managedKeyedState.registerSharedStates(sharedStateRegistry);
	}

	if (rawKeyedState != null) {
		rawKeyedState.registerSharedStates(sharedStateRegistry);
	}
}
 

/**
 * Tests that the pending checkpoint stats callbacks are created.
 */
@Test
public void testCheckpointStatsTrackerPendingCheckpointCallback() {
	final long timestamp = System.currentTimeMillis();
	ExecutionVertex vertex1 = mockExecutionVertex(new ExecutionAttemptID());

	// set up the coordinator and validate the initial state
	CheckpointCoordinator coord = new CheckpointCoordinator(
		new JobID(),
		600000,
		600000,
		0,
		Integer.MAX_VALUE,
		CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION,
		new ExecutionVertex[]{vertex1},
		new ExecutionVertex[]{vertex1},
		new ExecutionVertex[]{vertex1},
		new StandaloneCheckpointIDCounter(),
		new StandaloneCompletedCheckpointStore(1),
		new MemoryStateBackend(),
		Executors.directExecutor(),
		SharedStateRegistry.DEFAULT_FACTORY);

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

	when(tracker.reportPendingCheckpoint(anyLong(), anyLong(), any(CheckpointProperties.class)))
		.thenReturn(mock(PendingCheckpointStats.class));

	// Trigger a checkpoint and verify callback
	assertTrue(coord.triggerCheckpoint(timestamp, false));

	verify(tracker, times(1))
		.reportPendingCheckpoint(eq(1L), eq(timestamp), eq(CheckpointProperties.forCheckpoint(CheckpointRetentionPolicy.NEVER_RETAIN_AFTER_TERMINATION)));
}