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

Example 1 with KeyGroupsStateHandle

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

the class StreamTaskTest method testAsyncCheckpointingConcurrentCloseBeforeAcknowledge.

/**
	 * FLINK-5667
	 *
	 * Tests that a concurrent cancel operation discards the state handles of a not yet
	 * acknowledged checkpoint and prevents sending an acknowledge message to the
	 * CheckpointCoordinator. The situation can only happen if the cancel call is executed
	 * before Environment.acknowledgeCheckpoint().
	 */
@Test
public void testAsyncCheckpointingConcurrentCloseBeforeAcknowledge() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    final OneShotLatch createSubtask = new OneShotLatch();
    final OneShotLatch completeSubtask = 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);
    whenNew(SubtaskState.class).withAnyArguments().thenAnswer(new Answer<SubtaskState>() {

        @Override
        public SubtaskState answer(InvocationOnMock invocation) throws Throwable {
            createSubtask.trigger();
            completeSubtask.await();
            return new SubtaskState((ChainedStateHandle<StreamStateHandle>) invocation.getArguments()[0], (ChainedStateHandle<OperatorStateHandle>) invocation.getArguments()[1], (ChainedStateHandle<OperatorStateHandle>) invocation.getArguments()[2], (KeyGroupsStateHandle) invocation.getArguments()[3], (KeyGroupsStateHandle) invocation.getArguments()[4]);
        }
    });
    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);
    ExecutorService executor = Executors.newFixedThreadPool(1);
    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", executor);
    Whitebox.setInternalState(streamTask, "configuration", new StreamConfig(new Configuration()));
    Whitebox.setInternalState(streamTask, "stateBackend", mockStateBackend);
    streamTask.triggerCheckpoint(checkpointMetaData, CheckpointOptions.forFullCheckpoint());
    createSubtask.await();
    streamTask.cancel();
    completeSubtask.trigger();
    // wait for the completion of the async task
    executor.shutdown();
    if (!executor.awaitTermination(10000L, TimeUnit.MILLISECONDS)) {
        fail("Executor did not shut down within the given timeout. This indicates that the " + "checkpointing did not resume.");
    }
    // check that the checkpoint has not been acknowledged
    verify(mockEnvironment, never()).acknowledgeCheckpoint(eq(checkpointId), any(CheckpointMetrics.class), any(SubtaskState.class));
    // check that the state handles have been discarded
    verify(managedKeyedStateHandle).discardState();
    verify(rawKeyedStateHandle).discardState();
    verify(managedOperatorStateHandle).discardState();
    verify(rawOperatorStateHandle).discardState();
}

Example 2 with KeyGroupsStateHandle

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

the class AbstractStreamOperatorTestHarness method initializeStateFromLegacyCheckpoint.

public void initializeStateFromLegacyCheckpoint(String checkpointFilename) throws Exception {
    FileInputStream fin = new FileInputStream(checkpointFilename);
    StreamTaskState state = MigrationInstantiationUtil.deserializeObject(fin, ClassLoader.getSystemClassLoader());
    fin.close();
    if (!setupCalled) {
        setup();
    }
    StreamStateHandle stateHandle = SavepointV0Serializer.convertOperatorAndFunctionState(state);
    List<KeyGroupsStateHandle> keyGroupStatesList = new ArrayList<>();
    if (state.getKvStates() != null) {
        KeyGroupsStateHandle keyedStateHandle = SavepointV0Serializer.convertKeyedBackendState(state.getKvStates(), environment.getTaskInfo().getIndexOfThisSubtask(), 0);
        keyGroupStatesList.add(keyedStateHandle);
    }
    // finally calling the initializeState() with the legacy operatorStateHandles
    initializeState(new OperatorStateHandles(0, stateHandle, keyGroupStatesList, Collections.<KeyGroupsStateHandle>emptyList(), Collections.<OperatorStateHandle>emptyList(), Collections.<OperatorStateHandle>emptyList()));
}

Example 3 with KeyGroupsStateHandle

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

the class AbstractStreamOperatorTestHarness method initializeState.

/**
	 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#initializeState(OperatorStateHandles)}.
	 * Calls {@link org.apache.flink.streaming.api.operators.StreamOperator#setup(StreamTask, StreamConfig, Output)}
	 * if it was not called before.
	 *
	 * <p>This will reshape the state handles to include only those key-group states
	 * in the local key-group range and the operator states that would be assigned to the local
	 * subtask.
	 */
public void initializeState(OperatorStateHandles operatorStateHandles) throws Exception {
    if (!setupCalled) {
        setup();
    }
    if (operatorStateHandles != null) {
        int numKeyGroups = getEnvironment().getTaskInfo().getMaxNumberOfParallelSubtasks();
        int numSubtasks = getEnvironment().getTaskInfo().getNumberOfParallelSubtasks();
        int subtaskIndex = getEnvironment().getTaskInfo().getIndexOfThisSubtask();
        // create a new OperatorStateHandles that only contains the state for our key-groups
        List<KeyGroupRange> keyGroupPartitions = StateAssignmentOperation.createKeyGroupPartitions(numKeyGroups, numSubtasks);
        KeyGroupRange localKeyGroupRange = keyGroupPartitions.get(subtaskIndex);
        List<KeyGroupsStateHandle> localManagedKeyGroupState = null;
        if (operatorStateHandles.getManagedKeyedState() != null) {
            localManagedKeyGroupState = StateAssignmentOperation.getKeyGroupsStateHandles(operatorStateHandles.getManagedKeyedState(), localKeyGroupRange);
        }
        List<KeyGroupsStateHandle> localRawKeyGroupState = null;
        if (operatorStateHandles.getRawKeyedState() != null) {
            localRawKeyGroupState = StateAssignmentOperation.getKeyGroupsStateHandles(operatorStateHandles.getRawKeyedState(), localKeyGroupRange);
        }
        List<OperatorStateHandle> managedOperatorState = new ArrayList<>();
        if (operatorStateHandles.getManagedOperatorState() != null) {
            managedOperatorState.addAll(operatorStateHandles.getManagedOperatorState());
        }
        Collection<OperatorStateHandle> localManagedOperatorState = operatorStateRepartitioner.repartitionState(managedOperatorState, numSubtasks).get(subtaskIndex);
        List<OperatorStateHandle> rawOperatorState = new ArrayList<>();
        if (operatorStateHandles.getRawOperatorState() != null) {
            rawOperatorState.addAll(operatorStateHandles.getRawOperatorState());
        }
        Collection<OperatorStateHandle> localRawOperatorState = operatorStateRepartitioner.repartitionState(rawOperatorState, numSubtasks).get(subtaskIndex);
        OperatorStateHandles massagedOperatorStateHandles = new OperatorStateHandles(0, operatorStateHandles.getLegacyOperatorState(), localManagedKeyGroupState, localRawKeyGroupState, localManagedOperatorState, localRawOperatorState);
        operator.initializeState(massagedOperatorStateHandles);
    } else {
        operator.initializeState(null);
    }
    initializeCalled = true;
}

Example 4 with KeyGroupsStateHandle

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

the class AbstractStreamOperatorTestHarness method snapshot.

/**
	 * Calls {@link StreamOperator#snapshotState(long, long, CheckpointOptions)}.
	 */
public OperatorStateHandles snapshot(long checkpointId, long timestamp) throws Exception {
    CheckpointStreamFactory streamFactory = stateBackend.createStreamFactory(new JobID(), "test_op");
    OperatorSnapshotResult operatorStateResult = operator.snapshotState(checkpointId, timestamp, CheckpointOptions.forFullCheckpoint());
    KeyGroupsStateHandle keyedManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateManagedFuture());
    KeyGroupsStateHandle keyedRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getKeyedStateRawFuture());
    OperatorStateHandle opManaged = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateManagedFuture());
    OperatorStateHandle opRaw = FutureUtil.runIfNotDoneAndGet(operatorStateResult.getOperatorStateRawFuture());
    return new OperatorStateHandles(0, null, keyedManaged != null ? Collections.singletonList(keyedManaged) : null, keyedRaw != null ? Collections.singletonList(keyedRaw) : null, opManaged != null ? Collections.singletonList(opManaged) : null, opRaw != null ? Collections.singletonList(opRaw) : null);
}

Example 5 with KeyGroupsStateHandle

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

the class SavepointV1Serializer method deserializeSubtaskState.

private static SubtaskState deserializeSubtaskState(DataInputStream dis) throws IOException {
    // Duration field has been removed from SubtaskState
    long ignoredDuration = dis.readLong();
    int len = dis.readInt();
    List<StreamStateHandle> nonPartitionableState = new ArrayList<>(len);
    for (int i = 0; i < len; ++i) {
        StreamStateHandle streamStateHandle = deserializeStreamStateHandle(dis);
        nonPartitionableState.add(streamStateHandle);
    }
    len = dis.readInt();
    List<OperatorStateHandle> operatorStateBackend = new ArrayList<>(len);
    for (int i = 0; i < len; ++i) {
        OperatorStateHandle streamStateHandle = deserializeOperatorStateHandle(dis);
        operatorStateBackend.add(streamStateHandle);
    }
    len = dis.readInt();
    List<OperatorStateHandle> operatorStateStream = new ArrayList<>(len);
    for (int i = 0; i < len; ++i) {
        OperatorStateHandle streamStateHandle = deserializeOperatorStateHandle(dis);
        operatorStateStream.add(streamStateHandle);
    }
    KeyGroupsStateHandle keyedStateBackend = deserializeKeyGroupStateHandle(dis);
    KeyGroupsStateHandle keyedStateStream = deserializeKeyGroupStateHandle(dis);
    ChainedStateHandle<StreamStateHandle> nonPartitionableStateChain = new ChainedStateHandle<>(nonPartitionableState);
    ChainedStateHandle<OperatorStateHandle> operatorStateBackendChain = new ChainedStateHandle<>(operatorStateBackend);
    ChainedStateHandle<OperatorStateHandle> operatorStateStreamChain = new ChainedStateHandle<>(operatorStateStream);
    return new SubtaskState(nonPartitionableStateChain, operatorStateBackendChain, operatorStateStreamChain, keyedStateBackend, keyedStateStream);
}