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

Example 6 with StateInitializationContext

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