Examples with OperatorStateStore org.apache.flink.api.common.state.OperatorStateStore used on opensource projects

Example 6 with OperatorStateStore

use of org.apache.flink.api.common.state.OperatorStateStore in project flink by apache.

the class SourceOperatorEventTimeTest method createTestOperator.

// ------------------------------------------------------------------------
// test setup helpers
// ------------------------------------------------------------------------
private static <T> SourceOperator<T, MockSourceSplit> createTestOperator(SourceReader<T, MockSourceSplit> reader, WatermarkStrategy<T> watermarkStrategy, ProcessingTimeService timeService, boolean emitProgressiveWatermarks) throws Exception {
    final OperatorStateStore operatorStateStore = new MemoryStateBackend().createOperatorStateBackend(new MockEnvironmentBuilder().build(), "test-operator", Collections.emptyList(), new CloseableRegistry());
    final StateInitializationContext stateContext = new StateInitializationContextImpl(null, operatorStateStore, null, null, null);
    final SourceOperator<T, MockSourceSplit> sourceOperator = new TestingSourceOperator<>(reader, watermarkStrategy, timeService, emitProgressiveWatermarks);
    sourceOperator.setup(new SourceOperatorStreamTask<Integer>(new StreamMockEnvironment(new Configuration(), new Configuration(), new ExecutionConfig(), 1L, new MockInputSplitProvider(), 1, new TestTaskStateManager())), new MockStreamConfig(new Configuration(), 1), new MockOutput<>(new ArrayList<>()));
    sourceOperator.initializeState(stateContext);
    sourceOperator.open();
    return sourceOperator;
}

Example 7 with OperatorStateStore

use of org.apache.flink.api.common.state.OperatorStateStore in project beam by apache.

the class ImpulseSourceFunctionTest method getMockOperatorState.

private static <T> OperatorStateStore getMockOperatorState(ListState<T> listState) throws Exception {
    OperatorStateStore mock = Mockito.mock(OperatorStateStore.class);
    when(mock.getListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);
    return mock;
}

Example 8 with OperatorStateStore

use of org.apache.flink.api.common.state.OperatorStateStore in project beam by apache.

the class UnboundedSourceWrapper method initializeState.

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    if (checkpointCoder == null) {
        // no checkpoint coder available in this source
        return;
    }
    OperatorStateStore stateStore = context.getOperatorStateStore();
    @SuppressWarnings("unchecked") CoderTypeInformation<KV<? extends UnboundedSource<OutputT, CheckpointMarkT>, CheckpointMarkT>> typeInformation = (CoderTypeInformation) new CoderTypeInformation<>(checkpointCoder, serializedOptions.get());
    stateForCheckpoint = stateStore.getListState(new ListStateDescriptor<>(DefaultOperatorStateBackend.DEFAULT_OPERATOR_STATE_NAME, typeInformation.createSerializer(new ExecutionConfig())));
    if (context.isRestored()) {
        isRestored = true;
        LOG.info("Restoring state in the UnboundedSourceWrapper.");
    } else {
        LOG.info("No restore state for UnboundedSourceWrapper.");
    }
}

Example 9 with OperatorStateStore

use of org.apache.flink.api.common.state.OperatorStateStore in project flink by apache.

the class RollingSink method initializeState.

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    Preconditions.checkArgument(this.restoredBucketStates == null, "The " + getClass().getSimpleName() + " has already been initialized.");
    try {
        initFileSystem();
    } catch (IOException e) {
        LOG.error("Error while creating FileSystem when initializing the state of the RollingSink.", e);
        throw new RuntimeException("Error while creating FileSystem when initializing the state of the RollingSink.", e);
    }
    if (this.refTruncate == null) {
        this.refTruncate = reflectTruncate(fs);
    }
    OperatorStateStore stateStore = context.getOperatorStateStore();
    restoredBucketStates = stateStore.getSerializableListState("rolling-states");
    int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
    if (context.isRestored()) {
        LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
        for (BucketState bucketState : restoredBucketStates.get()) {
            handleRestoredBucketState(bucketState);
        }
        if (LOG.isDebugEnabled()) {
            LOG.debug("{} (taskIdx= {}) restored {}", getClass().getSimpleName(), subtaskIndex, bucketState);
        }
    } else {
        LOG.info("No state to restore for the {} (taskIdx= {}).", getClass().getSimpleName(), subtaskIndex);
    }
}

Example 10 with OperatorStateStore

use of org.apache.flink.api.common.state.OperatorStateStore in project flink by apache.

the class BucketingSink method initializeState.

@Override
public void initializeState(FunctionInitializationContext context) throws Exception {
    Preconditions.checkArgument(this.restoredBucketStates == null, "The operator has already been initialized.");
    try {
        initFileSystem();
    } catch (IOException e) {
        LOG.error("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
        throw new RuntimeException("Error while creating FileSystem when initializing the state of the BucketingSink.", e);
    }
    if (this.refTruncate == null) {
        this.refTruncate = reflectTruncate(fs);
    }
    OperatorStateStore stateStore = context.getOperatorStateStore();
    restoredBucketStates = stateStore.getSerializableListState("bucket-states");
    int subtaskIndex = getRuntimeContext().getIndexOfThisSubtask();
    if (context.isRestored()) {
        LOG.info("Restoring state for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
        for (State<T> recoveredState : restoredBucketStates.get()) {
            handleRestoredBucketState(recoveredState);
            if (LOG.isDebugEnabled()) {
                LOG.debug("{} idx {} restored {}", getClass().getSimpleName(), subtaskIndex, recoveredState);
            }
        }
    } else {
        LOG.info("No state to restore for the {} (taskIdx={}).", getClass().getSimpleName(), subtaskIndex);
    }
}