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

Example 1 with ResultSubpartitionStateHandle

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

the class OperatorSnapshotFinalizerTest method testRunAndExtract.

/**
 * Test that the runnable futures are executed and the result is correctly extracted.
 */
@Test
public void testRunAndExtract() throws Exception {
    Random random = new Random(0x42);
    KeyedStateHandle keyedTemplate = StateHandleDummyUtil.createNewKeyedStateHandle(new KeyGroupRange(0, 0));
    OperatorStateHandle operatorTemplate = StateHandleDummyUtil.createNewOperatorStateHandle(2, random);
    InputChannelStateHandle inputChannelTemplate = StateHandleDummyUtil.createNewInputChannelStateHandle(2, random);
    ResultSubpartitionStateHandle resultSubpartitionTemplate = StateHandleDummyUtil.createNewResultSubpartitionStateHandle(2, random);
    SnapshotResult<KeyedStateHandle> manKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<KeyedStateHandle> rawKeyed = withLocalState(deepDummyCopy(keyedTemplate), deepDummyCopy(keyedTemplate));
    SnapshotResult<OperatorStateHandle> manOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<OperatorStateHandle> rawOper = withLocalState(deepDummyCopy(operatorTemplate), deepDummyCopy(operatorTemplate));
    SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannel = withLocalState(singleton(deepDummyCopy(inputChannelTemplate)), singleton(deepDummyCopy(inputChannelTemplate)));
    SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartition = withLocalState(singleton(deepDummyCopy(resultSubpartitionTemplate)), singleton(deepDummyCopy(resultSubpartitionTemplate)));
    OperatorSnapshotFutures snapshotFutures = new OperatorSnapshotFutures(new PseudoNotDoneFuture<>(manKeyed), new PseudoNotDoneFuture<>(rawKeyed), new PseudoNotDoneFuture<>(manOper), new PseudoNotDoneFuture<>(rawOper), new PseudoNotDoneFuture<>(inputChannel), new PseudoNotDoneFuture<>(resultSubpartition));
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertFalse(f.isDone());
    }
    OperatorSnapshotFinalizer finalizer = new OperatorSnapshotFinalizer(snapshotFutures);
    for (Future<?> f : snapshotFutures.getAllFutures()) {
        assertTrue(f.isDone());
    }
    Map<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> map = new HashMap<>();
    map.put(manKeyed, headExtractor(OperatorSubtaskState::getManagedKeyedState));
    map.put(rawKeyed, headExtractor(OperatorSubtaskState::getRawKeyedState));
    map.put(manOper, headExtractor(OperatorSubtaskState::getManagedOperatorState));
    map.put(rawOper, headExtractor(OperatorSubtaskState::getRawOperatorState));
    map.put(inputChannel, OperatorSubtaskState::getInputChannelState);
    map.put(resultSubpartition, OperatorSubtaskState::getResultSubpartitionState);
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getJobManagerOwnedSnapshot(), e.getValue().apply(finalizer.getJobManagerOwnedState()));
    }
    for (Map.Entry<SnapshotResult<?>, Function<OperatorSubtaskState, ? extends StateObject>> e : map.entrySet()) {
        assertEquals(e.getKey().getTaskLocalSnapshot(), e.getValue().apply(finalizer.getTaskLocalState()));
    }
}

Example 2 with ResultSubpartitionStateHandle

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

the class OperatorSnapshotFuturesTest 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 {
    OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures();
    operatorSnapshotResult.cancel();
    KeyedStateHandle keyedManagedStateHandle = mock(KeyedStateHandle.class);
    SnapshotResult<KeyedStateHandle> keyedStateManagedResult = SnapshotResult.of(keyedManagedStateHandle);
    RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture = spy(DoneFuture.of(keyedStateManagedResult));
    KeyedStateHandle keyedRawStateHandle = mock(KeyedStateHandle.class);
    SnapshotResult<KeyedStateHandle> keyedStateRawResult = SnapshotResult.of(keyedRawStateHandle);
    RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture = spy(DoneFuture.of(keyedStateRawResult));
    OperatorStateHandle operatorManagedStateHandle = mock(OperatorStreamStateHandle.class);
    SnapshotResult<OperatorStateHandle> operatorStateManagedResult = SnapshotResult.of(operatorManagedStateHandle);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture = spy(DoneFuture.of(operatorStateManagedResult));
    OperatorStateHandle operatorRawStateHandle = mock(OperatorStreamStateHandle.class);
    SnapshotResult<OperatorStateHandle> operatorStateRawResult = SnapshotResult.of(operatorRawStateHandle);
    RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture = spy(DoneFuture.of(operatorStateRawResult));
    InputChannelStateHandle inputChannelRawStateHandle = mock(InputChannelStateHandle.class);
    SnapshotResult<StateObjectCollection<InputChannelStateHandle>> inputChannelStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(inputChannelRawStateHandle));
    Future<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>> inputChannelStateRawFuture = spy(DoneFuture.of(inputChannelStateRawResult));
    ResultSubpartitionStateHandle resultSubpartitionRawStateHandle = mock(ResultSubpartitionStateHandle.class);
    SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>> resultSubpartitionStateRawResult = SnapshotResult.of(StateObjectCollection.singleton(resultSubpartitionRawStateHandle));
    Future<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>> resultSubpartitionStateRawFuture = spy(DoneFuture.of(resultSubpartitionStateRawResult));
    operatorSnapshotResult = new OperatorSnapshotFutures(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture, inputChannelStateRawFuture, resultSubpartitionStateRawFuture);
    operatorSnapshotResult.cancel();
    verify(keyedStateManagedFuture).cancel(true);
    verify(keyedStateRawFuture).cancel(true);
    verify(operatorStateManagedFuture).cancel(true);
    verify(operatorStateRawFuture).cancel(true);
    verify(inputChannelStateRawFuture).cancel(true);
    verify(resultSubpartitionStateRawFuture).cancel(true);
    verify(keyedManagedStateHandle).discardState();
    verify(keyedRawStateHandle).discardState();
    verify(operatorManagedStateHandle).discardState();
    verify(operatorRawStateHandle).discardState();
    verify(inputChannelRawStateHandle).discardState();
    verify(resultSubpartitionRawStateHandle).discardState();
}

Example 3 with ResultSubpartitionStateHandle

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

the class StreamOperatorStateHandlerTest method testFailingBackendSnapshotMethod.

/**
 * Tests that a failing snapshot method call to the keyed state backend will trigger the closing
 * of the StateSnapshotContextSynchronousImpl and the cancellation of the
 * OperatorSnapshotResult. The latter is supposed to also cancel all assigned futures.
 */
@Test
public void testFailingBackendSnapshotMethod() throws Exception {
    final long checkpointId = 42L;
    final long timestamp = 1L;
    try (CloseableRegistry closeableRegistry = new CloseableRegistry()) {
        RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateManagedFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<KeyedStateHandle>> keyedStateRawFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateManagedFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<OperatorStateHandle>> operatorStateRawFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<StateObjectCollection<InputChannelStateHandle>>> inputChannelStateFuture = new CancelableFuture<>();
        RunnableFuture<SnapshotResult<StateObjectCollection<ResultSubpartitionStateHandle>>> resultSubpartitionStateFuture = new CancelableFuture<>();
        OperatorSnapshotFutures operatorSnapshotResult = new OperatorSnapshotFutures(keyedStateManagedFuture, keyedStateRawFuture, operatorStateManagedFuture, operatorStateRawFuture, inputChannelStateFuture, resultSubpartitionStateFuture);
        StateSnapshotContextSynchronousImpl context = new TestStateSnapshotContextSynchronousImpl(checkpointId, timestamp, closeableRegistry);
        context.getRawKeyedOperatorStateOutput();
        context.getRawOperatorStateOutput();
        StreamTaskStateInitializerImpl stateInitializer = new StreamTaskStateInitializerImpl(new MockEnvironmentBuilder().build(), new MemoryStateBackend());
        StreamOperatorStateContext stateContext = stateInitializer.streamOperatorStateContext(new OperatorID(), "whatever", new TestProcessingTimeService(), new UnUsedKeyContext(), IntSerializer.INSTANCE, closeableRegistry, new InterceptingOperatorMetricGroup(), 1.0, false);
        StreamOperatorStateHandler stateHandler = new StreamOperatorStateHandler(stateContext, new ExecutionConfig(), closeableRegistry);
        final String keyedStateField = "keyedStateField";
        final String operatorStateField = "operatorStateField";
        CheckpointedStreamOperator checkpointedStreamOperator = new CheckpointedStreamOperator() {

            @Override
            public void initializeState(StateInitializationContext context) throws Exception {
                context.getKeyedStateStore().getState(new ValueStateDescriptor<>(keyedStateField, LongSerializer.INSTANCE)).update(42L);
                context.getOperatorStateStore().getListState(new ListStateDescriptor<>(operatorStateField, LongSerializer.INSTANCE)).add(42L);
            }

            @Override
            public void snapshotState(StateSnapshotContext context) throws Exception {
                throw new ExpectedTestException();
            }
        };
        stateHandler.setCurrentKey("44");
        stateHandler.initializeOperatorState(checkpointedStreamOperator);
        assertThat(stateContext.operatorStateBackend().getRegisteredStateNames(), is(not(empty())));
        assertThat(((AbstractKeyedStateBackend<?>) stateContext.keyedStateBackend()).numKeyValueStatesByName(), equalTo(1));
        try {
            stateHandler.snapshotState(checkpointedStreamOperator, Optional.of(stateContext.internalTimerServiceManager()), "42", 42, 42, CheckpointOptions.forCheckpointWithDefaultLocation(), new MemCheckpointStreamFactory(1024), operatorSnapshotResult, context, false);
            fail("Exception expected.");
        } catch (CheckpointException e) {
            // as CheckpointException is wrapping the cause with SerializedThrowable
            if (!ExceptionUtils.findThrowableWithMessage(e, ExpectedTestException.MESSAGE).isPresent()) {
                throw e;
            }
        }
        assertTrue(keyedStateManagedFuture.isCancelled());
        assertTrue(keyedStateRawFuture.isCancelled());
        assertTrue(context.getKeyedStateStreamFuture().isCancelled());
        assertTrue(operatorStateManagedFuture.isCancelled());
        assertTrue(operatorStateRawFuture.isCancelled());
        assertTrue(context.getOperatorStateStreamFuture().isCancelled());
        assertTrue(inputChannelStateFuture.isCancelled());
        assertTrue(resultSubpartitionStateFuture.isCancelled());
        stateHandler.dispose();
        assertThat(stateContext.operatorStateBackend().getRegisteredBroadcastStateNames(), is(empty()));
        assertThat(stateContext.operatorStateBackend().getRegisteredStateNames(), is(empty()));
        assertThat(((AbstractKeyedStateBackend<?>) stateContext.keyedStateBackend()).numKeyValueStatesByName(), equalTo(0));
    }
}

Example 4 with ResultSubpartitionStateHandle

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

the class SequentialChannelStateReaderImplTest method writePermuted.

private Tuple2<List<InputChannelStateHandle>, List<ResultSubpartitionStateHandle>> writePermuted(Map<InputChannelInfo, List<byte[]>> inputChannels, Map<ResultSubpartitionInfo, List<byte[]>> resultSubpartitions) throws IOException {
    try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
        DataOutputStream dataStream = new DataOutputStream(out);
        serializer.writeHeader(dataStream);
        Map<InputChannelInfo, List<Long>> icOffsets = write(dataStream, permute(inputChannels));
        Map<ResultSubpartitionInfo, List<Long>> rsOffsets = write(dataStream, permute(resultSubpartitions));
        ByteStreamStateHandle streamStateHandle = new ByteStreamStateHandle("", out.toByteArray());
        return Tuple2.of(icOffsets.entrySet().stream().map(e -> new InputChannelStateHandle(e.getKey(), streamStateHandle, e.getValue())).collect(toList()), rsOffsets.entrySet().stream().map(e -> new ResultSubpartitionStateHandle(e.getKey(), streamStateHandle, e.getValue())).collect(toList()));
    }
}

Example 5 with ResultSubpartitionStateHandle

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

the class TaskStateSnapshotTest method testSizeIncludesChannelState.

@Test
public void testSizeIncludesChannelState() {
    final Random random = new Random();
    InputChannelStateHandle inputChannelStateHandle = createNewInputChannelStateHandle(10, random);
    ResultSubpartitionStateHandle resultSubpartitionStateHandle = createNewResultSubpartitionStateHandle(10, random);
    final TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot(Collections.singletonMap(new OperatorID(), OperatorSubtaskState.builder().setInputChannelState(singleton(inputChannelStateHandle)).setResultSubpartitionState(singleton(resultSubpartitionStateHandle)).build()));
    Assert.assertEquals(inputChannelStateHandle.getStateSize() + resultSubpartitionStateHandle.getStateSize(), taskStateSnapshot.getStateSize());
    Assert.assertTrue(taskStateSnapshot.hasState());
}