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

Example 26 with OperatorStateHandle

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

the class AbstractStreamOperator method initializeState.

@Override
public final void initializeState(OperatorStateHandles stateHandles) throws Exception {
    Collection<KeyGroupsStateHandle> keyedStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesRaw = null;
    Collection<OperatorStateHandle> operatorStateHandlesBackend = null;
    boolean restoring = null != stateHandles;
    //TODO we should move the actual initialization of this from StreamTask to this class
    initKeyedState();
    if (getKeyedStateBackend() != null && timeServiceManager == null) {
        timeServiceManager = new InternalTimeServiceManager<>(getKeyedStateBackend().getNumberOfKeyGroups(), getKeyedStateBackend().getKeyGroupRange(), this, getRuntimeContext().getProcessingTimeService());
    }
    if (restoring) {
        //pass directly
        operatorStateHandlesBackend = stateHandles.getManagedOperatorState();
        operatorStateHandlesRaw = stateHandles.getRawOperatorState();
        if (null != getKeyedStateBackend()) {
            //only use the keyed state if it is meant for us (aka head operator)
            keyedStateHandlesRaw = stateHandles.getRawKeyedState();
        }
    }
    checkpointStreamFactory = container.createCheckpointStreamFactory(this);
    initOperatorState(operatorStateHandlesBackend);
    StateInitializationContext initializationContext = new StateInitializationContextImpl(// information whether we restore or start for the first time
    restoring, // access to operator state backend
    operatorStateBackend, // access to keyed state backend
    keyedStateStore, // access to keyed state stream
    keyedStateHandlesRaw, // access to operator state stream
    operatorStateHandlesRaw, // access to register streams for canceling
    getContainingTask().getCancelables());
    initializeState(initializationContext);
    if (restoring) {
        // finally restore the legacy state in case we are
        // migrating from a previous Flink version.
        restoreStreamCheckpointed(stateHandles);
    }
}

Example 27 with OperatorStateHandle

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

the class OperatorSnapshotResultTest method testCancelAndCleanup.

/**
	 * Tests that all runnable futures in an OperatorSnapshotResult are properly cancelled and if
	 * the StreamStateHandle result is retrievable that the state handle are discarded.
	 */
@Test
public void testCancelAndCleanup() throws Exception {
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult();
    operatorSnapshotResult.cancel();
    KeyGroupsStateHandle keyedManagedStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateManagedFuture = mock(RunnableFuture.class);
    when(keyedStateManagedFuture.get()).thenReturn(keyedManagedStateHandle);
    KeyGroupsStateHandle keyedRawStateHandle = mock(KeyGroupsStateHandle.class);
    RunnableFuture<KeyGroupsStateHandle> keyedStateRawFuture = mock(RunnableFuture.class);
    when(keyedStateRawFuture.get()).thenReturn(keyedRawStateHandle);
    OperatorStateHandle operatorManagedStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateManagedFuture = mock(RunnableFuture.class);
    when(operatorStateManagedFuture.get()).thenReturn(operatorManagedStateHandle);
    OperatorStateHandle operatorRawStateHandle = mock(OperatorStateHandle.class);
    RunnableFuture<OperatorStateHandle> operatorStateRawFuture = mock(RunnableFuture.class);
    when(operatorStateRawFuture.get()).thenReturn(operatorRawStateHandle);
    operatorSnapshotResult = new OperatorSnapshotResult(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture);
    operatorSnapshotResult.cancel();
    verify(keyedStateManagedFuture).cancel(true);
    verify(keyedStateRawFuture).cancel(true);
    verify(operatorStateManagedFuture).cancel(true);
    verify(operatorStateRawFuture).cancel(true);
    verify(keyedManagedStateHandle).discardState();
    verify(keyedRawStateHandle).discardState();
    verify(operatorManagedStateHandle).discardState();
    verify(operatorRawStateHandle).discardState();
}

Example 28 with OperatorStateHandle

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

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseAfterAcknowledge.

/**
	 * FLINK-5667
	 *
	 * 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 long checkpointId = 42L;
    final long timestamp = 1L;
    final OneShotLatch acknowledgeCheckpointLatch = new OneShotLatch();
    final OneShotLatch completeAcknowledge = new OneShotLatch();
    TaskInfo mockTaskInfo = mock(TaskInfo.class);
    when(mockTaskInfo.getTaskNameWithSubtasks()).thenReturn("foobar");
    when(mockTaskInfo.getIndexOfThisSubtask()).thenReturn(0);
    Environment mockEnvironment = mock(Environment.class);
    when(mockEnvironment.getTaskInfo()).thenReturn(mockTaskInfo);
    doAnswer(new Answer() {

        @Override
        public Object answer(InvocationOnMock invocation) throws Throwable {
            acknowledgeCheckpointLatch.trigger();
            // block here so that we can issue the concurrent cancel call
            completeAcknowledge.await();
            return null;
        }
    }).when(mockEnvironment).acknowledgeCheckpoint(anyLong(), any(CheckpointMetrics.class), any(SubtaskState.class));
    StreamTask<?, AbstractStreamOperator<?>> streamTask = mock(StreamTask.class, Mockito.CALLS_REAL_METHODS);
    CheckpointMetaData checkpointMetaData = new CheckpointMetaData(checkpointId, timestamp);
    streamTask.setEnvironment(mockEnvironment);
    StreamOperator<?> streamOperator = mock(StreamOperator.class, withSettings().extraInterfaces(StreamCheckpointedOperator.class));
    KeyGroupsStateHandle managedKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    KeyGroupsStateHandle rawKeyedStateHandle = mock(KeyGroupsStateHandle.class);
    OperatorStateHandle managedOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorStateHandle rawOperatorStateHandle = mock(OperatorStateHandle.class);
    OperatorSnapshotResult operatorSnapshotResult = new OperatorSnapshotResult(new DoneFuture<>(managedKeyedStateHandle), new DoneFuture<>(rawKeyedStateHandle), new DoneFuture<>(managedOperatorStateHandle), new DoneFuture<>(rawOperatorStateHandle));
    when(streamOperator.snapshotState(anyLong(), anyLong(), any(CheckpointOptions.class))).thenReturn(operatorSnapshotResult);
    StreamOperator<?>[] streamOperators = { streamOperator };
    OperatorChain<Void, AbstractStreamOperator<Void>> operatorChain = mock(OperatorChain.class);
    when(operatorChain.getAllOperators()).thenReturn(streamOperators);
    StreamStateHandle streamStateHandle = mock(StreamStateHandle.class);
    CheckpointStreamFactory.CheckpointStateOutputStream outStream = mock(CheckpointStreamFactory.CheckpointStateOutputStream.class);
    when(outStream.closeAndGetHandle()).thenReturn(streamStateHandle);
    CheckpointStreamFactory mockStreamFactory = mock(CheckpointStreamFactory.class);
    when(mockStreamFactory.createCheckpointStateOutputStream(anyLong(), anyLong())).thenReturn(outStream);
    AbstractStateBackend mockStateBackend = mock(AbstractStateBackend.class);
    when(mockStateBackend.createStreamFactory(any(JobID.class), anyString())).thenReturn(mockStreamFactory);
    Whitebox.setInternalState(streamTask, "isRunning", true);
    Whitebox.setInternalState(streamTask, "lock", new Object());
    Whitebox.setInternalState(streamTask, "operatorChain", operatorChain);
    Whitebox.setInternalState(streamTask, "cancelables", new CloseableRegistry());
    Whitebox.setInternalState(streamTask, "asyncOperationsThreadPool", Executors.newFixedThreadPool(1));
    Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
    Whitebox.setInternalState(streamTask, "stateBackend", mockStateBackend);
    streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
    acknowledgeCheckpointLatch.await();
    ArgumentCaptor<SubtaskState> subtaskStateCaptor = ArgumentCaptor.forClass(SubtaskState.class);
    // check that the checkpoint has been completed
    verify(mockEnvironment).acknowledgeCheckpoint(eq(checkpointId), any(CheckpointMetrics.class), subtaskStateCaptor.capture());
    SubtaskState subtaskState = subtaskStateCaptor.getValue();
    // check that the subtask state contains the expected state handles
    assertEquals(managedKeyedStateHandle, subtaskState.getManagedKeyedState());
    assertEquals(rawKeyedStateHandle, subtaskState.getRawKeyedState());
    assertEquals(new ChainedStateHandle<>(Collections.singletonList(managedOperatorStateHandle)), subtaskState.getManagedOperatorState());
    assertEquals(new ChainedStateHandle<>(Collections.singletonList(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();
}

Example 29 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 30 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));
}