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

private static HeapKeyedStateBackend<Long> createHeapKeyedStateBackend(File rootDir) throws IOException {
	File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
	KeyGroupRange keyGroupRange = new KeyGroupRange(0, 1);
	int numberOfKeyGroups = keyGroupRange.getNumberOfKeyGroups();
	ExecutionConfig executionConfig = new ExecutionConfig();
	HeapPriorityQueueSetFactory priorityQueueSetFactory =
		new HeapPriorityQueueSetFactory(keyGroupRange, numberOfKeyGroups, 128);
	HeapKeyedStateBackendBuilder<Long> backendBuilder = new HeapKeyedStateBackendBuilder<>(
		null,
		new LongSerializer(),
		Thread.currentThread().getContextClassLoader(),
		numberOfKeyGroups,
		keyGroupRange,
		executionConfig,
		TtlTimeProvider.DEFAULT,
		Collections.emptyList(),
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		new LocalRecoveryConfig(false, new LocalRecoveryDirectoryProviderImpl(recoveryBaseDir, new JobID(), new JobVertexID(), 0)),
		priorityQueueSetFactory,
		false,
		new CloseableRegistry()
	);
	return backendBuilder.build();
}
 

private static RocksDBKeyedStateBackend<Long> createRocksDBKeyedStateBackend(File rootDir) throws IOException {
	File recoveryBaseDir = prepareDirectory(recoveryDirName, rootDir);
	File dbPathFile = prepareDirectory(dbDirName, rootDir);
	ExecutionConfig executionConfig = new ExecutionConfig();
	RocksDBResourceContainer resourceContainer = new RocksDBResourceContainer();
	RocksDBKeyedStateBackendBuilder<Long> builder = new RocksDBKeyedStateBackendBuilder<>(
		"Test",
		Thread.currentThread().getContextClassLoader(),
		dbPathFile,
		resourceContainer,
		stateName -> resourceContainer.getColumnOptions(),
		null,
		LongSerializer.INSTANCE,
		2,
		new KeyGroupRange(0, 1),
		executionConfig,
		new LocalRecoveryConfig(false, new LocalRecoveryDirectoryProviderImpl(recoveryBaseDir, new JobID(), new JobVertexID(), 0)),
		RocksDBStateBackend.PriorityQueueStateType.ROCKSDB,
		TtlTimeProvider.DEFAULT,
		new UnregisteredMetricsGroup(),
		Collections.emptyList(),
		AbstractStateBackend.getCompressionDecorator(executionConfig),
		new CloseableRegistry());
	try {
		return builder.build();
	} catch (Exception e) {
		IOUtils.closeQuietly(resourceContainer);
		throw e;
	}
}
 

/**
 * This tests that state that was reported to the {@link org.apache.flink.runtime.state.TaskStateManager} is also
 * reported to {@link org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link TaskLocalStateStoreImpl}.
 */
@Test
public void testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore() throws Exception {

	final JobID jobID = new JobID();
	final AllocationID allocationID = new AllocationID();
	final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
	final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
	final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
	final int subtaskIdx = 42;
	JobVertexID jobVertexID = new JobVertexID();

	TaskStateSnapshot jmSnapshot = new TaskStateSnapshot();
	TaskStateSnapshot tmSnapshot = new TaskStateSnapshot();

	final AtomicBoolean jmReported = new AtomicBoolean(false);
	final AtomicBoolean tmReported = new AtomicBoolean(false);

	TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {

		@Override
		public void acknowledgeCheckpoint(
			JobID lJobID,
			ExecutionAttemptID lExecutionAttemptID,
			long lCheckpointId,
			CheckpointMetrics lCheckpointMetrics,
			TaskStateSnapshot lSubtaskState) {

			Assert.assertEquals(jobID, lJobID);
			Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
			Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
			Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
			jmReported.set(true);
		}
	};

	Executor executor = Executors.directExecutor();

	LocalRecoveryDirectoryProviderImpl directoryProvider = new LocalRecoveryDirectoryProviderImpl(
		temporaryFolder.newFolder(),
		jobID,
		jobVertexID,
		subtaskIdx);

	LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(true, directoryProvider);

	TaskLocalStateStore taskLocalStateStore =
		new TaskLocalStateStoreImpl(jobID, allocationID, jobVertexID, subtaskIdx, localRecoveryConfig, executor) {
			@Override
			public void storeLocalState(
				@Nonnegative long checkpointId,
				@Nullable TaskStateSnapshot localState) {

				Assert.assertEquals(tmSnapshot, localState);
				tmReported.set(true);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			null,
			checkpointResponder);

	taskStateManager.reportTaskStateSnapshots(
		checkpointMetaData,
		checkpointMetrics,
		jmSnapshot,
		tmSnapshot);

	Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
	Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
}
 

public StreamTaskTestHarness(
	Function<Environment, ? extends StreamTask<OUT, ?>> taskFactory,
	TypeInformation<OUT> outputType,
	File localRootDir) {
	this(taskFactory, outputType, new LocalRecoveryConfig(true, new LocalRecoveryDirectoryProviderImpl(localRootDir, new JobID(), new JobVertexID(), 0)));
}
 

/**
 * This tests that state that was reported to the {@link org.apache.flink.runtime.state.TaskStateManager} is also
 * reported to {@link org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link TaskLocalStateStoreImpl}.
 */
@Test
public void testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore() throws Exception {

	final JobID jobID = new JobID();
	final AllocationID allocationID = new AllocationID();
	final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
	final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
	final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
	final int subtaskIdx = 42;
	JobVertexID jobVertexID = new JobVertexID();

	TaskStateSnapshot jmSnapshot = new TaskStateSnapshot();
	TaskStateSnapshot tmSnapshot = new TaskStateSnapshot();

	final AtomicBoolean jmReported = new AtomicBoolean(false);
	final AtomicBoolean tmReported = new AtomicBoolean(false);

	TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {

		@Override
		public void acknowledgeCheckpoint(
			JobID lJobID,
			ExecutionAttemptID lExecutionAttemptID,
			long lCheckpointId,
			CheckpointMetrics lCheckpointMetrics,
			TaskStateSnapshot lSubtaskState) {

			Assert.assertEquals(jobID, lJobID);
			Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
			Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
			Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
			jmReported.set(true);
		}
	};

	Executor executor = Executors.directExecutor();

	LocalRecoveryDirectoryProviderImpl directoryProvider = new LocalRecoveryDirectoryProviderImpl(
		temporaryFolder.newFolder(),
		jobID,
		jobVertexID,
		subtaskIdx);

	LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(true, directoryProvider);

	TaskLocalStateStore taskLocalStateStore =
		new TaskLocalStateStoreImpl(jobID, allocationID, jobVertexID, subtaskIdx, localRecoveryConfig, executor) {
			@Override
			public void storeLocalState(
				@Nonnegative long checkpointId,
				@Nullable TaskStateSnapshot localState) {

				Assert.assertEquals(tmSnapshot, localState);
				tmReported.set(true);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			null,
			checkpointResponder);

	taskStateManager.reportTaskStateSnapshots(
		checkpointMetaData,
		checkpointMetrics,
		jmSnapshot,
		tmSnapshot);

	Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
	Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
}
 

public StreamTaskTestHarness(
	Function<Environment, ? extends StreamTask<OUT, ?>> taskFactory,
	TypeInformation<OUT> outputType,
	File localRootDir) {
	this(taskFactory, outputType, new LocalRecoveryConfig(true, new LocalRecoveryDirectoryProviderImpl(localRootDir, new JobID(), new JobVertexID(), 0)));
}
 

/**
 * This tests that state that was reported to the {@link org.apache.flink.runtime.state.TaskStateManager} is also
 * reported to {@link org.apache.flink.runtime.taskmanager.CheckpointResponder} and {@link TaskLocalStateStoreImpl}.
 */
@Test
public void testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore() throws Exception {

	final JobID jobID = new JobID();
	final AllocationID allocationID = new AllocationID();
	final ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
	final CheckpointMetaData checkpointMetaData = new CheckpointMetaData(42L, 4711L);
	final CheckpointMetrics checkpointMetrics = new CheckpointMetrics();
	final int subtaskIdx = 42;
	JobVertexID jobVertexID = new JobVertexID();

	TaskStateSnapshot jmSnapshot = new TaskStateSnapshot();
	TaskStateSnapshot tmSnapshot = new TaskStateSnapshot();

	final AtomicBoolean jmReported = new AtomicBoolean(false);
	final AtomicBoolean tmReported = new AtomicBoolean(false);

	TestCheckpointResponder checkpointResponder = new TestCheckpointResponder() {

		@Override
		public void acknowledgeCheckpoint(
			JobID lJobID,
			ExecutionAttemptID lExecutionAttemptID,
			long lCheckpointId,
			CheckpointMetrics lCheckpointMetrics,
			TaskStateSnapshot lSubtaskState) {

			Assert.assertEquals(jobID, lJobID);
			Assert.assertEquals(executionAttemptID, lExecutionAttemptID);
			Assert.assertEquals(checkpointMetaData.getCheckpointId(), lCheckpointId);
			Assert.assertEquals(checkpointMetrics, lCheckpointMetrics);
			jmReported.set(true);
		}
	};

	Executor executor = Executors.directExecutor();

	LocalRecoveryDirectoryProviderImpl directoryProvider = new LocalRecoveryDirectoryProviderImpl(
		temporaryFolder.newFolder(),
		jobID,
		jobVertexID,
		subtaskIdx);

	LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(true, directoryProvider);

	TaskLocalStateStore taskLocalStateStore =
		new TaskLocalStateStoreImpl(jobID, allocationID, jobVertexID, subtaskIdx, localRecoveryConfig, executor) {
			@Override
			public void storeLocalState(
				@Nonnegative long checkpointId,
				@Nullable TaskStateSnapshot localState) {

				Assert.assertEquals(tmSnapshot, localState);
				tmReported.set(true);
			}
		};

	TaskStateManagerImpl taskStateManager =
		new TaskStateManagerImpl(
			jobID,
			executionAttemptID,
			taskLocalStateStore,
			null,
			checkpointResponder);

	taskStateManager.reportTaskStateSnapshots(
		checkpointMetaData,
		checkpointMetrics,
		jmSnapshot,
		tmSnapshot);

	Assert.assertTrue("Reporting for JM state was not called.", jmReported.get());
	Assert.assertTrue("Reporting for TM state was not called.", tmReported.get());
}
 

public StreamTaskTestHarness(
	FunctionWithException<Environment, ? extends StreamTask<OUT, ?>, Exception> taskFactory,
	TypeInformation<OUT> outputType,
	File localRootDir) {
	this(taskFactory, outputType, new LocalRecoveryConfig(true, new LocalRecoveryDirectoryProviderImpl(localRootDir, new JobID(), new JobVertexID(), 0)));
}