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

Example 51 with OperatorStateHandle

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

the class AbstractStreamOperatorTestHarness method repackageState.

/**
	 * Takes the different {@link OperatorStateHandles} created by calling {@link #snapshot(long, long)}
	 * on different instances of {@link AbstractStreamOperatorTestHarness} (each one representing one subtask)
	 * and repacks them into a single {@link OperatorStateHandles} so that the parallelism of the test
	 * can change arbitrarily (i.e. be able to scale both up and down).
	 *
	 * <p>
	 * After repacking the partial states, use {@link #initializeState(OperatorStateHandles)} to initialize
	 * a new instance with the resulting state. Bear in mind that for parallelism greater than one, you
	 * have to use the constructor {@link #AbstractStreamOperatorTestHarness(StreamOperator, int, int, int)}.
	 *
	 * <p>
	 * <b>NOTE: </b> each of the {@code handles} in the argument list is assumed to be from a single task of a single
	 * operator (i.e. chain length of one).
	 *
	 * <p>
	 * For an example of how to use it, have a look at
	 * {@link AbstractStreamOperatorTest#testStateAndTimerStateShufflingScalingDown()}.
	 *
	 * @param handles the different states to be merged.
	 * @return the resulting state, or {@code null} if no partial states are specified.
	 */
public static OperatorStateHandles repackageState(OperatorStateHandles... handles) throws Exception {
    if (handles.length < 1) {
        return null;
    } else if (handles.length == 1) {
        return handles[0];
    }
    List<OperatorStateHandle> mergedManagedOperatorState = new ArrayList<>(handles.length);
    List<OperatorStateHandle> mergedRawOperatorState = new ArrayList<>(handles.length);
    List<KeyGroupsStateHandle> mergedManagedKeyedState = new ArrayList<>(handles.length);
    List<KeyGroupsStateHandle> mergedRawKeyedState = new ArrayList<>(handles.length);
    for (OperatorStateHandles handle : handles) {
        Collection<OperatorStateHandle> managedOperatorState = handle.getManagedOperatorState();
        Collection<OperatorStateHandle> rawOperatorState = handle.getRawOperatorState();
        Collection<KeyGroupsStateHandle> managedKeyedState = handle.getManagedKeyedState();
        Collection<KeyGroupsStateHandle> rawKeyedState = handle.getRawKeyedState();
        if (managedOperatorState != null) {
            mergedManagedOperatorState.addAll(managedOperatorState);
        }
        if (rawOperatorState != null) {
            mergedRawOperatorState.addAll(rawOperatorState);
        }
        if (managedKeyedState != null) {
            mergedManagedKeyedState.addAll(managedKeyedState);
        }
        if (rawKeyedState != null) {
            mergedRawKeyedState.addAll(rawKeyedState);
        }
    }
    return new OperatorStateHandles(0, null, mergedManagedKeyedState, mergedRawKeyedState, mergedManagedOperatorState, mergedRawOperatorState);
}

Example 52 with OperatorStateHandle

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

the class InterruptSensitiveRestoreTest method createTask.

// ------------------------------------------------------------------------
//  Utilities
// ------------------------------------------------------------------------
private static Task createTask(Configuration taskConfig, StreamStateHandle state, int mode) throws IOException {
    NetworkEnvironment networkEnvironment = mock(NetworkEnvironment.class);
    when(networkEnvironment.createKvStateTaskRegistry(any(JobID.class), any(JobVertexID.class))).thenReturn(mock(TaskKvStateRegistry.class));
    ChainedStateHandle<StreamStateHandle> operatorState = null;
    List<KeyGroupsStateHandle> keyGroupStateFromBackend = Collections.emptyList();
    List<KeyGroupsStateHandle> keyGroupStateFromStream = Collections.emptyList();
    List<Collection<OperatorStateHandle>> operatorStateBackend = Collections.emptyList();
    List<Collection<OperatorStateHandle>> operatorStateStream = Collections.emptyList();
    Map<String, OperatorStateHandle.StateMetaInfo> operatorStateMetadata = new HashMap<>(1);
    OperatorStateHandle.StateMetaInfo metaInfo = new OperatorStateHandle.StateMetaInfo(new long[] { 0 }, OperatorStateHandle.Mode.SPLIT_DISTRIBUTE);
    operatorStateMetadata.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, metaInfo);
    KeyGroupRangeOffsets keyGroupRangeOffsets = new KeyGroupRangeOffsets(new KeyGroupRange(0, 0));
    Collection<OperatorStateHandle> operatorStateHandles = Collections.singletonList(new OperatorStateHandle(operatorStateMetadata, state));
    List<KeyGroupsStateHandle> keyGroupsStateHandles = Collections.singletonList(new KeyGroupsStateHandle(keyGroupRangeOffsets, state));
    switch(mode) {
        case OPERATOR_MANAGED:
            operatorStateBackend = Collections.singletonList(operatorStateHandles);
            break;
        case OPERATOR_RAW:
            operatorStateStream = Collections.singletonList(operatorStateHandles);
            break;
        case KEYED_MANAGED:
            keyGroupStateFromBackend = keyGroupsStateHandles;
            break;
        case KEYED_RAW:
            keyGroupStateFromStream = keyGroupsStateHandles;
            break;
        case LEGACY:
            operatorState = new ChainedStateHandle<>(Collections.singletonList(state));
            break;
        default:
            throw new IllegalArgumentException();
    }
    TaskStateHandles taskStateHandles = new TaskStateHandles(operatorState, operatorStateBackend, operatorStateStream, keyGroupStateFromBackend, keyGroupStateFromStream);
    JobInformation jobInformation = new JobInformation(new JobID(), "test job name", new SerializedValue<>(new ExecutionConfig()), new Configuration(), Collections.<BlobKey>emptyList(), Collections.<URL>emptyList());
    TaskInformation taskInformation = new TaskInformation(new JobVertexID(), "test task name", 1, 1, SourceStreamTask.class.getName(), taskConfig);
    return new Task(jobInformation, taskInformation, new ExecutionAttemptID(), new AllocationID(), 0, 0, Collections.<ResultPartitionDeploymentDescriptor>emptyList(), Collections.<InputGateDeploymentDescriptor>emptyList(), 0, taskStateHandles, mock(MemoryManager.class), mock(IOManager.class), networkEnvironment, mock(BroadcastVariableManager.class), mock(TaskManagerActions.class), mock(InputSplitProvider.class), mock(CheckpointResponder.class), new FallbackLibraryCacheManager(), new FileCache(new String[] { EnvironmentInformation.getTemporaryFileDirectory() }), new TestingTaskManagerRuntimeInfo(), new UnregisteredTaskMetricsGroup(), mock(ResultPartitionConsumableNotifier.class), mock(PartitionProducerStateChecker.class), mock(Executor.class));
}

Example 53 with OperatorStateHandle

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

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseAfterAcknowledge.

/**
 * FLINK-5667
 *
 * <p>Tests that a concurrent cancel operation does not discard the state handles of an
 * acknowledged checkpoint. The situation can only happen if the cancel call is executed after
 * Environment.acknowledgeCheckpoint() and before the CloseableRegistry.unregisterClosable()
 * call.
 */
@Test
public void testAsyncCheckpointingConcurrentCloseAfterAcknowledge() throws Exception {
    final OneShotLatch acknowledgeCheckpointLatch = new OneShotLatch();
    final OneShotLatch completeAcknowledge = new OneShotLatch();
    CheckpointResponder checkpointResponder = mock(CheckpointResponder.class);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) {
            acknowledgeCheckpointLatch.trigger();
            // block here so that we can issue the concurrent cancel call
            while (true) {
                try {
                    // wait until we successfully await (no pun intended)
                    completeAcknowledge.await();
                    // when await() returns normally, we break out of the loop
                    break;
                } catch (InterruptedException e) {
                // survive interruptions that arise from thread pool
                // shutdown
                // production code cannot actually throw
                // InterruptedException from
                // checkpoint acknowledgement
                }
            }
            return null;
        }
    }).when(checkpointResponder).acknowledgeCheckpoint(any(JobID.class), any(ExecutionAttemptID.class), anyLong(), any(CheckpointMetrics.class), any(TaskStateSnapshot.class));
    TaskStateManager taskStateManager = new TaskStateManagerImpl(new JobID(1L, 2L), new ExecutionAttemptID(), mock(TaskLocalStateStoreImpl.class), new InMemoryStateChangelogStorage(), null, checkpointResponder);
    KeyedStateHandle managedKeyedStateHandle = mock(KeyedStateHandle.class);
    KeyedStateHandle rawKeyedStateHandle = mock(KeyedStateHandle.class);
    OperatorStateHandle managedOperatorStateHandle = mock(OperatorStreamStateHandle.class);
    OperatorStateHandle rawOperatorStateHandle = mock(OperatorStreamStateHandle.class);
    OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(DoneFuture.of(SnapshotResult.of(managedKeyedStateHandle)), DoneFuture.of(SnapshotResult.of(rawKeyedStateHandle)), DoneFuture.of(SnapshotResult.of(managedOperatorStateHandle)), DoneFuture.of(SnapshotResult.of(rawOperatorStateHandle)), DoneFuture.of(SnapshotResult.empty()), DoneFuture.of(SnapshotResult.empty()));
    try (MockEnvironment mockEnvironment = new MockEnvironmentBuilder().setTaskName("mock-task").setTaskStateManager(taskStateManager).build()) {
        RunningTask<MockStreamTask> task = runTask(() -> createMockStreamTask(mockEnvironment, operatorChain(streamOperatorWithSnapshot(operatorSnapshotResult))));
        MockStreamTask streamTask = task.streamTask;
        waitTaskIsRunning(streamTask, task.invocationFuture);
        final long checkpointId = 42L;
        streamTask.triggerCheckpointAsync(new CheckpointMetaData(checkpointId, 1L), CheckpointOptions.forCheckpointWithDefaultLocation());
        acknowledgeCheckpointLatch.await();
        ArgumentCaptor<TaskStateSnapshot> subtaskStateCaptor = ArgumentCaptor.forClass(TaskStateSnapshot.class);
        // check that the checkpoint has been completed
        verify(checkpointResponder).acknowledgeCheckpoint(any(JobID.class), any(ExecutionAttemptID.class), eq(checkpointId), any(CheckpointMetrics.class), subtaskStateCaptor.capture());
        TaskStateSnapshot subtaskStates = subtaskStateCaptor.getValue();
        OperatorSubtaskState subtaskState = subtaskStates.getSubtaskStateMappings().iterator().next().getValue();
        // check that the subtask state contains the expected state handles
        assertEquals(singleton(managedKeyedStateHandle), subtaskState.getManagedKeyedState());
        assertEquals(singleton(rawKeyedStateHandle), subtaskState.getRawKeyedState());
        assertEquals(singleton(managedOperatorStateHandle), subtaskState.getManagedOperatorState());
        assertEquals(singleton(rawOperatorStateHandle), subtaskState.getRawOperatorState());
        // check that the state handles have not been discarded
        verify(managedKeyedStateHandle, never()).discardState();
        verify(rawKeyedStateHandle, never()).discardState();
        verify(managedOperatorStateHandle, never()).discardState();
        verify(rawOperatorStateHandle, never()).discardState();
        streamTask.cancel();
        completeAcknowledge.trigger();
        // canceling the stream task after it has acknowledged the checkpoint should not discard
        // the state handles
        verify(managedKeyedStateHandle, never()).discardState();
        verify(rawKeyedStateHandle, never()).discardState();
        verify(managedOperatorStateHandle, never()).discardState();
        verify(rawOperatorStateHandle, never()).discardState();
        task.waitForTaskCompletion(true);
    }
}

Example 54 with OperatorStateHandle

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

the class RoundRobinOperatorStateRepartitioner method repartitionUnionState.

/**
 * Repartition UNION state.
 */
private void repartitionUnionState(Map<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> unionState, List<Map<StreamStateHandle, OperatorStateHandle>> mergeMapList) {
    for (Map<StreamStateHandle, OperatorStateHandle> mergeMap : mergeMapList) {
        for (Map.Entry<String, List<Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo>>> e : unionState.entrySet()) {
            for (Tuple2<StreamStateHandle, OperatorStateHandle.StateMetaInfo> handleWithMetaInfo : e.getValue()) {
                OperatorStateHandle operatorStateHandle = mergeMap.get(handleWithMetaInfo.f0);
                if (operatorStateHandle == null) {
                    operatorStateHandle = new OperatorStreamStateHandle(new HashMap<>(unionState.size()), handleWithMetaInfo.f0);
                    mergeMap.put(handleWithMetaInfo.f0, operatorStateHandle);
                }
                operatorStateHandle.getStateNameToPartitionOffsets().put(e.getKey(), handleWithMetaInfo.f1);
            }
        }
    }
}

Example 55 with OperatorStateHandle

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

the class StateAssignmentOperationTest method testRepartitionUnionState.

@Test
public void testRepartitionUnionState() {
    OperatorID operatorID = new OperatorID();
    OperatorState operatorState = new OperatorState(operatorID, 2, 4);
    Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap1 = new HashMap<>(2);
    metaInfoMap1.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[] { 0 }, OperatorStateHandle.Mode.UNION));
    metaInfoMap1.put("t-4", new OperatorStateHandle.StateMetaInfo(new long[] { 22, 44 }, OperatorStateHandle.Mode.UNION));
    OperatorStateHandle osh1 = new OperatorStreamStateHandle(metaInfoMap1, new ByteStreamStateHandle("test1", new byte[50]));
    operatorState.putState(0, OperatorSubtaskState.builder().setManagedOperatorState(osh1).build());
    Map<String, OperatorStateHandle.StateMetaInfo> metaInfoMap2 = new HashMap<>(1);
    metaInfoMap2.put("t-3", new OperatorStateHandle.StateMetaInfo(new long[] { 0 }, OperatorStateHandle.Mode.UNION));
    OperatorStateHandle osh2 = new OperatorStreamStateHandle(metaInfoMap2, new ByteStreamStateHandle("test2", new byte[20]));
    operatorState.putState(1, OperatorSubtaskState.builder().setManagedOperatorState(osh2).build());
    verifyOneKindPartitionableStateRescale(operatorState, operatorID);
}