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

Example 1 with StateInitializationContextImpl

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

the class StateInitializationContextImplTest method setUp.

@Before
public void setUp() throws Exception {
    this.writtenKeyGroups = 0;
    this.writtenOperatorStates = new HashSet<>();
    this.closableRegistry = new CloseableRegistry();
    OperatorStateStore stateStore = mock(OperatorStateStore.class);
    ByteArrayOutputStreamWithPos out = new ByteArrayOutputStreamWithPos(64);
    List<KeyGroupsStateHandle> keyGroupsStateHandles = new ArrayList<>(NUM_HANDLES);
    int prev = 0;
    for (int i = 0; i < NUM_HANDLES; ++i) {
        out.reset();
        int size = i % 4;
        int end = prev + size;
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        KeyGroupRangeOffsets offsets = new KeyGroupRangeOffsets(i == 9 ? KeyGroupRange.EMPTY_KEY_GROUP_RANGE : new KeyGroupRange(prev, end));
        prev = end + 1;
        for (int kg : offsets.getKeyGroupRange()) {
            offsets.setKeyGroupOffset(kg, out.getPosition());
            dov.writeInt(kg);
            ++writtenKeyGroups;
        }
        KeyGroupsStateHandle handle = new KeyGroupsStateHandle(offsets, new ByteStateHandleCloseChecking("kg-" + i, out.toByteArray()));
        keyGroupsStateHandles.add(handle);
    }
    List<OperatorStateHandle> operatorStateHandles = new ArrayList<>(NUM_HANDLES);
    for (int i = 0; i < NUM_HANDLES; ++i) {
        int size = i % 4;
        out.reset();
        DataOutputView dov = new DataOutputViewStreamWrapper(out);
        LongArrayList offsets = new LongArrayList(size);
        for (int s = 0; s < size; ++s) {
            offsets.add(out.getPosition());
            int val = i * NUM_HANDLES + s;
            dov.writeInt(val);
            writtenOperatorStates.add(val);
        }
        Map<String, OperatorStateHandle.StateMetaInfo> offsetsMap = new HashMap<>();
        offsetsMap.put(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, new OperatorStateHandle.StateMetaInfo(offsets.toArray(), OperatorStateHandle.Mode.SPLIT_DISTRIBUTE));
        OperatorStateHandle operatorStateHandle = new OperatorStateHandle(offsetsMap, new ByteStateHandleCloseChecking("os-" + i, out.toByteArray()));
        operatorStateHandles.add(operatorStateHandle);
    }
    this.initializationContext = new StateInitializationContextImpl(true, stateStore, mock(KeyedStateStore.class), keyGroupsStateHandles, operatorStateHandles, closableRegistry);
}

Example 2 with StateInitializationContextImpl

use of org.apache.flink.runtime.state.StateInitializationContextImpl 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);
    }
}