Examples with LocalRecoveryDirectoryProviderImpl org.apache.flink.runtime.state.LocalRecoveryDirectoryProviderImpl used on opensource projects

Example 1 with LocalRecoveryDirectoryProviderImpl

use of org.apache.flink.runtime.state.LocalRecoveryDirectoryProviderImpl in project flink by apache.

the class LocalStateForwardingTest method testReportingFromTaskStateManagerToResponderAndTaskLocalStateStore.

/**
 * 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(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);
        }
    };
    StateChangelogStorage<?> stateChangelogStorage = new InMemoryStateChangelogStorage();
    TaskStateManagerImpl taskStateManager = new TaskStateManagerImpl(jobID, executionAttemptID, taskLocalStateStore, stateChangelogStorage, 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());
}

Example 2 with LocalRecoveryDirectoryProviderImpl

use of org.apache.flink.runtime.state.LocalRecoveryDirectoryProviderImpl in project flink by apache.

the class StreamOperatorSnapshotRestoreTest method testOperatorStatesSnapshotRestoreInternal.

private void testOperatorStatesSnapshotRestoreInternal(final int mode) throws Exception {
    // -------------------------------------------------------------------------- snapshot
    StateBackend stateBackend;
    FsStateBackend fsstateBackend = createStateBackendInternal();
    switch(stateBackendEnum) {
        case FILE:
            stateBackend = fsstateBackend;
            break;
        case ROCKSDB_FULLY_ASYNC:
            stateBackend = new RocksDBStateBackend(fsstateBackend, TernaryBoolean.FALSE);
            break;
        case ROCKSDB_INCREMENTAL:
            stateBackend = new RocksDBStateBackend(fsstateBackend, TernaryBoolean.TRUE);
            break;
        default:
            throw new IllegalStateException(String.format("Do not support statebackend type %s", stateBackendEnum));
    }
    TestOneInputStreamOperator op = new TestOneInputStreamOperator(false);
    JobID jobID = new JobID();
    JobVertexID jobVertexID = new JobVertexID();
    int subtaskIdx = 0;
    LocalRecoveryDirectoryProvider directoryProvider = mode == ONLY_JM_RECOVERY ? null : new LocalRecoveryDirectoryProviderImpl(temporaryFolder.newFolder(), jobID, jobVertexID, subtaskIdx);
    LocalRecoveryConfig localRecoveryConfig = new LocalRecoveryConfig(directoryProvider);
    MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setJobID(jobID).setJobVertexID(jobVertexID).setTaskName("test").setManagedMemorySize(1024L * 1024L).setInputSplitProvider(new MockInputSplitProvider()).setBufferSize(1024 * 1024).setTaskStateManager(new TestTaskStateManager(localRecoveryConfig)).setMaxParallelism(MAX_PARALLELISM).setSubtaskIndex(subtaskIdx).setUserCodeClassLoader(getClass().getClassLoader()).build();
    KeyedOneInputStreamOperatorTestHarness<Integer, Integer, Integer> testHarness = new KeyedOneInputStreamOperatorTestHarness<>(op, (KeySelector<Integer, Integer>) value -> value, TypeInformation.of(Integer.class), mockEnvironment);
    testHarness.setStateBackend(stateBackend);
    testHarness.open();
    for (int i = 0; i < 10; ++i) {
        testHarness.processElement(new StreamRecord<>(i));
    }
    OperatorSnapshotFinalizer snapshotWithLocalState = testHarness.snapshotWithLocalState(1L, 1L);
    testHarness.close();
    // -------------------------------------------------------------------------- restore
    op = new TestOneInputStreamOperator(true);
    testHarness = new KeyedOneInputStreamOperatorTestHarness<>(op, (KeySelector<Integer, Integer>) value -> value, TypeInformation.of(Integer.class), MAX_PARALLELISM, 1, /* num subtasks */
    0);
    testHarness.setTimeServiceManagerProvider(new InternalTimeServiceManager.Provider() {

        @Override
        public <K> InternalTimeServiceManager<K> create(CheckpointableKeyedStateBackend<K> keyedStatedBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws IOException {
            return null;
        }
    });
    testHarness.setStateBackend(stateBackend);
    OperatorSubtaskState jobManagerOwnedState = snapshotWithLocalState.getJobManagerOwnedState();
    OperatorSubtaskState taskLocalState = snapshotWithLocalState.getTaskLocalState();
    // We check if local state was created when we enabled local recovery
    Assert.assertTrue(mode > ONLY_JM_RECOVERY == (taskLocalState != null && taskLocalState.hasState()));
    if (mode == TM_REMOVE_JM_RECOVERY) {
        jobManagerOwnedState.getManagedKeyedState().discardState();
    } else if (mode == JM_REMOVE_TM_RECOVERY) {
        taskLocalState.getManagedKeyedState().discardState();
    }
    testHarness.initializeState(jobManagerOwnedState, taskLocalState);
    testHarness.open();
    for (int i = 0; i < 10; ++i) {
        testHarness.processElement(new StreamRecord<>(i));
    }
    testHarness.close();
}