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

Example 1 with StatePartitionStreamProvider

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

the class StateInitializationContextImplTest method getOperatorStateStore.

@Test
public void getOperatorStateStore() throws Exception {
    Set<Integer> readStatesCount = new HashSet<>();
    for (StatePartitionStreamProvider statePartitionStreamProvider : initializationContext.getRawOperatorStateInputs()) {
        Assert.assertNotNull(statePartitionStreamProvider);
        try (InputStream is = statePartitionStreamProvider.getStream()) {
            DataInputView div = new DataInputViewStreamWrapper(is);
            Assert.assertTrue(readStatesCount.add(div.readInt()));
        }
    }
    Assert.assertEquals(writtenOperatorStates, readStatesCount);
}

Example 2 with StatePartitionStreamProvider

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

the class StreamTaskStateInitializerImpl method rawOperatorStateInputs.

protected CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs(Iterator<StateObjectCollection<OperatorStateHandle>> restoreStateAlternatives) {
    if (restoreStateAlternatives.hasNext()) {
        Collection<OperatorStateHandle> rawOperatorState = restoreStateAlternatives.next();
        // TODO currently this does not support local state recovery, so we expect there is only
        // one handle.
        Preconditions.checkState(!restoreStateAlternatives.hasNext(), "Local recovery is currently not implemented for raw operator state, but found state alternative.");
        if (rawOperatorState != null) {
            return new CloseableIterable<StatePartitionStreamProvider>() {

                final CloseableRegistry closeableRegistry = new CloseableRegistry();

                @Override
                public void close() throws IOException {
                    closeableRegistry.close();
                }

                @Nonnull
                @Override
                public Iterator<StatePartitionStreamProvider> iterator() {
                    return new OperatorStateStreamIterator(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, rawOperatorState.iterator(), closeableRegistry);
                }
            };
        }
    }
    return CloseableIterable.empty();
}

Example 3 with StatePartitionStreamProvider

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

the class StreamTaskStateInitializerImpl method streamOperatorStateContext.

// -----------------------------------------------------------------------------------------------------------------
@Override
public StreamOperatorStateContext streamOperatorStateContext(@Nonnull OperatorID operatorID, @Nonnull String operatorClassName, @Nonnull ProcessingTimeService processingTimeService, @Nonnull KeyContext keyContext, @Nullable TypeSerializer<?> keySerializer, @Nonnull CloseableRegistry streamTaskCloseableRegistry, @Nonnull MetricGroup metricGroup, double managedMemoryFraction, boolean isUsingCustomRawKeyedState) throws Exception {
    TaskInfo taskInfo = environment.getTaskInfo();
    OperatorSubtaskDescriptionText operatorSubtaskDescription = new OperatorSubtaskDescriptionText(operatorID, operatorClassName, taskInfo.getIndexOfThisSubtask(), taskInfo.getNumberOfParallelSubtasks());
    final String operatorIdentifierText = operatorSubtaskDescription.toString();
    final PrioritizedOperatorSubtaskState prioritizedOperatorSubtaskStates = taskStateManager.prioritizedOperatorState(operatorID);
    CheckpointableKeyedStateBackend<?> keyedStatedBackend = null;
    OperatorStateBackend operatorStateBackend = null;
    CloseableIterable<KeyGroupStatePartitionStreamProvider> rawKeyedStateInputs = null;
    CloseableIterable<StatePartitionStreamProvider> rawOperatorStateInputs = null;
    InternalTimeServiceManager<?> timeServiceManager;
    try {
        // -------------- Keyed State Backend --------------
        keyedStatedBackend = keyedStatedBackend(keySerializer, operatorIdentifierText, prioritizedOperatorSubtaskStates, streamTaskCloseableRegistry, metricGroup, managedMemoryFraction);
        // -------------- Operator State Backend --------------
        operatorStateBackend = operatorStateBackend(operatorIdentifierText, prioritizedOperatorSubtaskStates, streamTaskCloseableRegistry);
        // -------------- Raw State Streams --------------
        rawKeyedStateInputs = rawKeyedStateInputs(prioritizedOperatorSubtaskStates.getPrioritizedRawKeyedState().iterator());
        streamTaskCloseableRegistry.registerCloseable(rawKeyedStateInputs);
        rawOperatorStateInputs = rawOperatorStateInputs(prioritizedOperatorSubtaskStates.getPrioritizedRawOperatorState().iterator());
        streamTaskCloseableRegistry.registerCloseable(rawOperatorStateInputs);
        // -------------- Internal Timer Service Manager --------------
        if (keyedStatedBackend != null) {
            // if the operator indicates that it is using custom raw keyed state,
            // then whatever was written in the raw keyed state snapshot was NOT written
            // by the internal timer services (because there is only ever one user of raw keyed
            // state);
            // in this case, timers should not attempt to restore timers from the raw keyed
            // state.
            final Iterable<KeyGroupStatePartitionStreamProvider> restoredRawKeyedStateTimers = (prioritizedOperatorSubtaskStates.isRestored() && !isUsingCustomRawKeyedState) ? rawKeyedStateInputs : Collections.emptyList();
            timeServiceManager = timeServiceManagerProvider.create(keyedStatedBackend, environment.getUserCodeClassLoader().asClassLoader(), keyContext, processingTimeService, restoredRawKeyedStateTimers);
        } else {
            timeServiceManager = null;
        }
        return new StreamOperatorStateContextImpl(prioritizedOperatorSubtaskStates.getRestoredCheckpointId(), operatorStateBackend, keyedStatedBackend, timeServiceManager, rawOperatorStateInputs, rawKeyedStateInputs);
    } catch (Exception ex) {
        // cleanup if something went wrong before results got published.
        if (keyedStatedBackend != null) {
            if (streamTaskCloseableRegistry.unregisterCloseable(keyedStatedBackend)) {
                IOUtils.closeQuietly(keyedStatedBackend);
            }
            // release resource (e.g native resource)
            keyedStatedBackend.dispose();
        }
        if (operatorStateBackend != null) {
            if (streamTaskCloseableRegistry.unregisterCloseable(operatorStateBackend)) {
                IOUtils.closeQuietly(operatorStateBackend);
            }
            operatorStateBackend.dispose();
        }
        if (streamTaskCloseableRegistry.unregisterCloseable(rawKeyedStateInputs)) {
            IOUtils.closeQuietly(rawKeyedStateInputs);
        }
        if (streamTaskCloseableRegistry.unregisterCloseable(rawOperatorStateInputs)) {
            IOUtils.closeQuietly(rawOperatorStateInputs);
        }
        throw new Exception("Exception while creating StreamOperatorStateContext.", ex);
    }
}

Example 4 with StatePartitionStreamProvider

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

the class StreamOperatorStateHandler method initializeOperatorState.

public void initializeOperatorState(CheckpointedStreamOperator streamOperator) throws Exception {
    CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = context.rawKeyedStateInputs();
    CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = context.rawOperatorStateInputs();
    try {
        OptionalLong checkpointId = context.getRestoredCheckpointId();
        StateInitializationContext initializationContext = new StateInitializationContextImpl(checkpointId.isPresent() ? checkpointId.getAsLong() : null, // access to operator state backend
        operatorStateBackend, // access to keyed state backend
        keyedStateStore, // access to keyed state stream
        keyedStateInputs, // access to operator state stream
        operatorStateInputs);
        streamOperator.initializeState(initializationContext);
    } finally {
        closeFromRegistry(operatorStateInputs, closeableRegistry);
        closeFromRegistry(keyedStateInputs, closeableRegistry);
    }
}

Example 5 with StatePartitionStreamProvider

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

the class StreamTaskStateInitializerImplTest method testWithRestore.

@SuppressWarnings("unchecked")
@Test
public void testWithRestore() throws Exception {
    StateBackend mockingBackend = spy(new StateBackend() {

        @Override
        public <K> AbstractKeyedStateBackend<K> createKeyedStateBackend(Environment env, JobID jobID, String operatorIdentifier, TypeSerializer<K> keySerializer, int numberOfKeyGroups, KeyGroupRange keyGroupRange, TaskKvStateRegistry kvStateRegistry, TtlTimeProvider ttlTimeProvider, MetricGroup metricGroup, @Nonnull Collection<KeyedStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(AbstractKeyedStateBackend.class);
        }

        @Override
        public OperatorStateBackend createOperatorStateBackend(Environment env, String operatorIdentifier, @Nonnull Collection<OperatorStateHandle> stateHandles, CloseableRegistry cancelStreamRegistry) throws Exception {
            return mock(OperatorStateBackend.class);
        }
    });
    OperatorID operatorID = new OperatorID(47L, 11L);
    TaskStateSnapshot taskStateSnapshot = new TaskStateSnapshot();
    Random random = new Random(0x42);
    OperatorSubtaskState operatorSubtaskState = OperatorSubtaskState.builder().setManagedOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("a", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 10 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setRawOperatorState(new OperatorStreamStateHandle(Collections.singletonMap("_default_", new OperatorStateHandle.StateMetaInfo(new long[] { 0, 20, 30 }, SPLIT_DISTRIBUTE)), CheckpointTestUtils.createDummyStreamStateHandle(random, null))).setManagedKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setRawKeyedState(CheckpointTestUtils.createDummyKeyGroupStateHandle(random, null)).setInputChannelState(singleton(createNewInputChannelStateHandle(10, random))).setResultSubpartitionState(singleton(createNewResultSubpartitionStateHandle(10, random))).build();
    taskStateSnapshot.putSubtaskStateByOperatorID(operatorID, operatorSubtaskState);
    JobManagerTaskRestore jobManagerTaskRestore = new JobManagerTaskRestore(42L, taskStateSnapshot);
    StreamTaskStateInitializer streamTaskStateManager = streamTaskStateManager(mockingBackend, jobManagerTaskRestore, false);
    AbstractStreamOperator<?> streamOperator = mock(AbstractStreamOperator.class);
    when(streamOperator.getOperatorID()).thenReturn(operatorID);
    TypeSerializer<?> typeSerializer = new IntSerializer();
    CloseableRegistry closeableRegistry = new CloseableRegistry();
    StreamOperatorStateContext stateContext = streamTaskStateManager.streamOperatorStateContext(streamOperator.getOperatorID(), streamOperator.getClass().getSimpleName(), new TestProcessingTimeService(), streamOperator, typeSerializer, closeableRegistry, new UnregisteredMetricsGroup(), 1.0, false);
    OperatorStateBackend operatorStateBackend = stateContext.operatorStateBackend();
    CheckpointableKeyedStateBackend<?> keyedStateBackend = stateContext.keyedStateBackend();
    InternalTimeServiceManager<?> timeServiceManager = stateContext.internalTimerServiceManager();
    CloseableIterable<KeyGroupStatePartitionStreamProvider> keyedStateInputs = stateContext.rawKeyedStateInputs();
    CloseableIterable<StatePartitionStreamProvider> operatorStateInputs = stateContext.rawOperatorStateInputs();
    Assert.assertTrue("Expected the context to be restored", stateContext.isRestored());
    Assert.assertEquals(OptionalLong.of(42L), stateContext.getRestoredCheckpointId());
    Assert.assertNotNull(operatorStateBackend);
    Assert.assertNotNull(keyedStateBackend);
    // this is deactivated on purpose so that it does not attempt to consume the raw keyed
    // state.
    Assert.assertNull(timeServiceManager);
    Assert.assertNotNull(keyedStateInputs);
    Assert.assertNotNull(operatorStateInputs);
    int count = 0;
    for (KeyGroupStatePartitionStreamProvider keyedStateInput : keyedStateInputs) {
        ++count;
    }
    Assert.assertEquals(1, count);
    count = 0;
    for (StatePartitionStreamProvider operatorStateInput : operatorStateInputs) {
        ++count;
    }
    Assert.assertEquals(3, count);
    checkCloseablesRegistered(closeableRegistry, operatorStateBackend, keyedStateBackend, keyedStateInputs, operatorStateInputs);
}