Search in sources :

Example 1 with TestTaskLocalStateStore

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

the class StateInitializationContextImplTest method setUp.

@Before
public void setUp() throws Exception {
    this.writtenKeyGroups = 0;
    this.writtenOperatorStates = new HashSet<>();
    this.closableRegistry = new CloseableRegistry();
    ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos(64);
    List<KeyedStateHandle> keyedStateHandles = new ArrayList<>(NUM_HANDLES);
    int prev = 0;
    for (int i = 0; i < NUM_HANDLES; ++i) {
        out.reset();
        int size = i % 4;
        int end = prev + size;
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(i == 9 ? KeyGroupRange.EMPTY_KEY_GROUP_RANGE : new KeyGroupRange(prev, end));
        prev = end + 1;
        for (int kg : offsets.getKeyGroupRange()) {
            offsets.setKeyGroupOffset(kg, out.getPosition());
            dov.writeInt(kg);
            ++writtenKeyGroups;
        }
        KeyedStateHandle handle = new KeyGroupsStateHandle(offsets, new ByteStateHandleCloseChecking("kg-" + i, out.toByteArray()));
        keyedStateHandles.add(handle);
    }
    List<OperatorStateHandle> operatorStateHandles = new ArrayList<>(NUM_HANDLES);
    for (int i = 0; i < NUM_HANDLES; ++i) {
        int size = i % 4;
        out.reset();
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        LongArrayList offsets = new LongArrayList(size);
        for (int s = 0; s < size; ++s) {
            offsets.add(out.getPosition());
            int val = i * NUM_HANDLES + s;
            dov.writeInt(val);
            writtenOperatorStates.add(val);
        }
        Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
        offsetsMap.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new OperatorStateHandle.StateMetaInfo(offsets.toArray(), OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
        OperatorStateHandle operatorStateHandle = new OperatorStreamStateHandle(offsetsMap, new ByteStateHandleCloseChecking("os-" + i, out.toByteArray()));
        operatorStateHandles.add(operatorStateHandle);
    }
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setRawOperatorState(new StateObjectCollection<>(operatorStateHandles)).setRawKeyedState(new StateObjectCollection<>(keyedStateHandles)).build();
    OperatorID operatorID = new OperatorID();
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(0L, taskStateSnapshot);
    TaskStateManager manager = new TaskStateManagerImpl(new JobID(), new ExecutionAttemptID(), new TestTaskLocalStateStore(), new InMemoryStateChangelogStorage(), jobManagerTaskRestore, mock(CheckpointResponder.class));
    DummyEnvironment environment = new DummyEnvironment("test", 1, 0, prev);
    environment.setTaskStateManager(manager);
    StateBackend stateBackend = new MemoryStateBackend(1024);
    StreamTaskStateInitializer streamTaskStateManager = new StreamTaskStateInitializerImpl(environment, stateBackend, TtlTimeProvider.DEFAULT, new InternalTimeServiceManager.Provider() {

        @Override
        public <K> InternalTimeServiceManager<K> create(CheckpointableKeyedStateBackend<K> keyedStatedBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
            // stream.
            return null;
        }
    });
    AbstractStreamOperator<?> mockOperator = mock(AbstractStreamOperator.class);
    when(mockOperator.getOperatorID()).thenReturn(operatorID);
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(operatorID, "TestOperatorClass", mock(ProcessingTimeService.class), mockOperator, // consumed by the timer service.
    IntSerializer.INSTANCE, closableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OptionalLong restoredCheckpointId = stateContext.getRestoredCheckpointId();
    this.initializationContext = new StateInitializationContextImpl(restoredCheckpointId.isPresent() ? restoredCheckpointId.getAsLong() : null, stateContext.operatorStateBackend(), mock(KeyedStateStore.class), stateContext.rawKeyedStateInputs(), stateContext.rawOperatorStateInputs());
}
Also used : HashMap(java.util.HashMap) KeyGroupRangeOffsets(org.apache.flink.runtime.state.KeyGroupRangeOffsets) LongArrayList(org.apache.flink.runtime.util.LongArrayList) ArrayList(java.util.ArrayList) KeyGroupRange(org.apache.flink.runtime.state.KeyGroupRange) CloseableRegistry(org.apache.flink.core.fs.CloseableRegistry) OperatorSubtaskState(org.apache.flink.runtime.checkpoint.OperatorSubtaskState) StateBackend(org.apache.flink.runtime.state.StateBackend) DefaultOperatorStateBackend(org.apache.flink.runtime.state.DefaultOperatorStateBackend) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) CheckpointableKeyedStateBackend(org.apache.flink.runtime.state.CheckpointableKeyedStateBackend) KeyGroupStatePartitionStreamProvider(org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider) InMemoryStateChangelogStorage(org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage) ByteArrayOutputStreamWithPos(org.apache.flink.core.memory.ByteArrayOutputStreamWithPos) TestTaskLocalStateStore(org.apache.flink.runtime.state.TestTaskLocalStateStore) DataOutputView(org.apache.flink.core.memory.DataOutputView) StateObjectCollection(org.apache.flink.runtime.checkpoint.StateObjectCollection) OperatorStreamStateHandle(org.apache.flink.runtime.state.OperatorStreamStateHandle) StateInitializationContextImpl(org.apache.flink.runtime.state.StateInitializationContextImpl) OperatorStateHandle(org.apache.flink.runtime.state.OperatorStateHandle) JobID(org.apache.flink.api.common.JobID) TaskStateManagerImpl(org.apache.flink.runtime.state.TaskStateManagerImpl) UnregisteredMetricsGroup(org.apache.flink.metrics.groups.UnregisteredMetricsGroup) MemoryStateBackend(org.apache.flink.runtime.state.memory.MemoryStateBackend) JobManagerTaskRestore(org.apache.flink.runtime.checkpoint.JobManagerTaskRestore) OperatorID(org.apache.flink.runtime.jobgraph.OperatorID) KeyedStateHandle(org.apache.flink.runtime.state.KeyedStateHandle) KeyGroupsStateHandle(org.apache.flink.runtime.state.KeyGroupsStateHandle) TaskStateSnapshot(org.apache.flink.runtime.checkpoint.TaskStateSnapshot) ProcessingTimeService(org.apache.flink.streaming.runtime.tasks.ProcessingTimeService) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) LongArrayList(org.apache.flink.runtime.util.LongArrayList) CheckpointResponder(org.apache.flink.runtime.taskmanager.CheckpointResponder) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) TaskStateManager(org.apache.flink.runtime.state.TaskStateManager) IOException(java.io.IOException) DataOutputViewStreamWrapper(org.apache.flink.core.memory.DataOutputViewStreamWrapper) OptionalLong(java.util.OptionalLong) Before(org.junit.Before)

Example 2 with TestTaskLocalStateStore

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

the class StreamTaskStateInitializerImplTest method streamTaskStateManager.

private StreamTaskStateInitializer streamTaskStateManager(StateBackend stateBackend, JobManagerTaskRestore jobManagerTaskRestore, boolean createTimerServiceManager) {
    JobID jobID = new JobID(42L, 43L);
    ExecutionAttemptID executionAttemptID = new ExecutionAttemptID();
    TestCheckpointResponder checkpointResponderMock = new TestCheckpointResponder();
    TaskLocalStateStore taskLocalStateStore = new TestTaskLocalStateStore();
    InMemoryStateChangelogStorage changelogStorage = new InMemoryStateChangelogStorage();
    TaskStateManager taskStateManager = TaskStateManagerImplTest.taskStateManager(jobID, executionAttemptID, checkpointResponderMock, jobManagerTaskRestore, taskLocalStateStore, changelogStorage);
    DummyEnvironment dummyEnvironment = new DummyEnvironment("test-task", 1, 0);
    dummyEnvironment.setTaskStateManager(taskStateManager);
    if (createTimerServiceManager) {
        return new StreamTaskStateInitializerImpl(dummyEnvironment, stateBackend);
    } else {
        return new StreamTaskStateInitializerImpl(dummyEnvironment, stateBackend, TtlTimeProvider.DEFAULT, new InternalTimeServiceManager.Provider() {

            @Override
            public <K> InternalTimeServiceManager<K> create(CheckpointableKeyedStateBackend<K> keyedStatedBackend, ClassLoader userClassloader, KeyContext keyContext, ProcessingTimeService processingTimeService, Iterable<KeyGroupStatePartitionStreamProvider> rawKeyedStates) throws Exception {
                return null;
            }
        });
    }
}
Also used : TestTaskLocalStateStore(org.apache.flink.runtime.state.TestTaskLocalStateStore) ExecutionAttemptID(org.apache.flink.runtime.executiongraph.ExecutionAttemptID) TaskLocalStateStore(org.apache.flink.runtime.state.TaskLocalStateStore) TestTaskLocalStateStore(org.apache.flink.runtime.state.TestTaskLocalStateStore) DummyEnvironment(org.apache.flink.runtime.operators.testutils.DummyEnvironment) TaskStateManager(org.apache.flink.runtime.state.TaskStateManager) KeyGroupStatePartitionStreamProvider(org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider) InMemoryStateChangelogStorage(org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage) TestProcessingTimeService(org.apache.flink.streaming.runtime.tasks.TestProcessingTimeService) ProcessingTimeService(org.apache.flink.streaming.runtime.tasks.ProcessingTimeService) TestCheckpointResponder(org.apache.flink.runtime.taskmanager.TestCheckpointResponder) JobID(org.apache.flink.api.common.JobID)

Aggregations

JobID (org.apache.flink.api.common.JobID)2 ExecutionAttemptID (org.apache.flink.runtime.executiongraph.ExecutionAttemptID)2 DummyEnvironment (org.apache.flink.runtime.operators.testutils.DummyEnvironment)2 KeyGroupStatePartitionStreamProvider (org.apache.flink.runtime.state.KeyGroupStatePartitionStreamProvider)2 TaskStateManager (org.apache.flink.runtime.state.TaskStateManager)2 TestTaskLocalStateStore (org.apache.flink.runtime.state.TestTaskLocalStateStore)2 InMemoryStateChangelogStorage (org.apache.flink.runtime.state.changelog.inmemory.InMemoryStateChangelogStorage)2 ProcessingTimeService (org.apache.flink.streaming.runtime.tasks.ProcessingTimeService)2 IOException (java.io.IOException)1 ArrayList (java.util.ArrayList)1 HashMap (java.util.HashMap)1 OptionalLong (java.util.OptionalLong)1 CloseableRegistry (org.apache.flink.core.fs.CloseableRegistry)1 ByteArrayOutputStreamWithPos (org.apache.flink.core.memory.ByteArrayOutputStreamWithPos)1 DataOutputView (org.apache.flink.core.memory.DataOutputView)1 DataOutputViewStreamWrapper (org.apache.flink.core.memory.DataOutputViewStreamWrapper)1 UnregisteredMetricsGroup (org.apache.flink.metrics.groups.UnregisteredMetricsGroup)1 JobManagerTaskRestore (org.apache.flink.runtime.checkpoint.JobManagerTaskRestore)1 OperatorSubtaskState (org.apache.flink.runtime.checkpoint.OperatorSubtaskState)1 StateObjectCollection (org.apache.flink.runtime.checkpoint.StateObjectCollection)1