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

Example 1 with OperatorStateStore

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

the class FlinkKafkaConsumerBaseTest method checkUseFetcherWhenNoCheckpoint.

/**
	 * Tests that on snapshots, states and offsets to commit to Kafka are correct
	 */
@SuppressWarnings("unchecked")
@Test
public void checkUseFetcherWhenNoCheckpoint() throws Exception {
    FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
    List<KafkaTopicPartition> partitionList = new ArrayList<>(1);
    partitionList.add(new KafkaTopicPartition("test", 0));
    consumer.setSubscribedPartitions(partitionList);
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    TestingListState<Serializable> listState = new TestingListState<>();
    when(operatorStateStore.getSerializableListState(Matchers.any(String.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    // make the context signal that there is no restored state, then validate that
    when(initializationContext.isRestored()).thenReturn(false);
    consumer.initializeState(initializationContext);
    consumer.run(mock(SourceFunction.SourceContext.class));
}

Example 2 with OperatorStateStore

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

the class RMQSourceTest method getMockContext.

/**
 * Gets a mock context for initializing the source's state via {@link
 * org.apache.flink.streaming.api.checkpoint.CheckpointedFunction#initializeState}.
 *
 * @throws Exception
 */
FunctionInitializationContext getMockContext() throws Exception {
    OperatorStateStore mockStore = Mockito.mock(OperatorStateStore.class);
    FunctionInitializationContext mockContext = Mockito.mock(FunctionInitializationContext.class);
    Mockito.when(mockContext.getOperatorStateStore()).thenReturn(mockStore);
    Mockito.when(mockStore.getListState(any(ListStateDescriptor.class))).thenReturn(null);
    return mockContext;
}

Example 3 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testUseRestoredStateForSnapshotIfFetcherNotInitialized.

// ----------------------------------------------------------------------
// Tests related to state initialization
// ----------------------------------------------------------------------
@Test
public void testUseRestoredStateForSnapshotIfFetcherNotInitialized() throws Exception {
    Properties config = TestUtils.getStandardProperties();
    List<Tuple2<StreamShardMetadata, SequenceNumber>> globalUnionState = new ArrayList<>(4);
    globalUnionState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))), new SequenceNumber("1")));
    globalUnionState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1)))), new SequenceNumber("1")));
    globalUnionState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))), new SequenceNumber("1")));
    globalUnionState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(3)))), new SequenceNumber("1")));
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Tuple2<StreamShardMetadata, SequenceNumber> state : globalUnionState) {
        listState.add(state);
    }
    FlinkKinesisConsumer<String> consumer = new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
    RuntimeContext context = mock(RuntimeContext.class);
    when(context.getIndexOfThisSubtask()).thenReturn(0);
    when(context.getNumberOfParallelSubtasks()).thenReturn(2);
    consumer.setRuntimeContext(context);
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    when(initializationContext.isRestored()).thenReturn(true);
    consumer.initializeState(initializationContext);
    // only opened, not run
    consumer.open(new Configuration());
    // arbitrary checkpoint id and timestamp
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(123, 123));
    assertTrue(listState.isClearCalled());
    // the checkpointed list state should contain only the shards that it should subscribe to
    assertEquals(globalUnionState.size() / 2, listState.getList().size());
    assertTrue(listState.getList().contains(globalUnionState.get(0)));
    assertTrue(listState.getList().contains(globalUnionState.get(2)));
}

Example 4 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testListStateChangedAfterSnapshotState.

@Test
public void testListStateChangedAfterSnapshotState() throws Exception {
    // ----------------------------------------------------------------------
    // setup config, initial state and expected state snapshot
    // ----------------------------------------------------------------------
    Properties config = TestUtils.getStandardProperties();
    ArrayList<Tuple2<StreamShardMetadata, SequenceNumber>> initialState = new ArrayList<>(1);
    initialState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))), new SequenceNumber("1")));
    ArrayList<Tuple2<StreamShardMetadata, SequenceNumber>> expectedStateSnapshot = new ArrayList<>(3);
    expectedStateSnapshot.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(0)))), new SequenceNumber("12")));
    expectedStateSnapshot.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(1)))), new SequenceNumber("11")));
    expectedStateSnapshot.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(new StreamShardHandle("fakeStream1", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2)))), new SequenceNumber("31")));
    // ----------------------------------------------------------------------
    // mock operator state backend and initial state for initializeState()
    // ----------------------------------------------------------------------
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Tuple2<StreamShardMetadata, SequenceNumber> state : initialState) {
        listState.add(state);
    }
    OperatorStateStore operatorStateStore = mock(OperatorStateStore.class);
    when(operatorStateStore.getUnionListState(Matchers.any(ListStateDescriptor.class))).thenReturn(listState);
    StateInitializationContext initializationContext = mock(StateInitializationContext.class);
    when(initializationContext.getOperatorStateStore()).thenReturn(operatorStateStore);
    when(initializationContext.isRestored()).thenReturn(true);
    // ----------------------------------------------------------------------
    // mock a running fetcher and its state for snapshot
    // ----------------------------------------------------------------------
    HashMap<StreamShardMetadata, SequenceNumber> stateSnapshot = new HashMap<>();
    for (Tuple2<StreamShardMetadata, SequenceNumber> tuple : expectedStateSnapshot) {
        stateSnapshot.put(tuple.f0, tuple.f1);
    }
    KinesisDataFetcher mockedFetcher = mock(KinesisDataFetcher.class);
    when(mockedFetcher.snapshotState()).thenReturn(stateSnapshot);
    // ----------------------------------------------------------------------
    // create a consumer and test the snapshotState()
    // ----------------------------------------------------------------------
    FlinkKinesisConsumer<String> consumer = new FlinkKinesisConsumer<>("fakeStream", new SimpleStringSchema(), config);
    FlinkKinesisConsumer<?> mockedConsumer = spy(consumer);
    RuntimeContext context = mock(RuntimeContext.class);
    when(context.getIndexOfThisSubtask()).thenReturn(1);
    mockedConsumer.setRuntimeContext(context);
    mockedConsumer.initializeState(initializationContext);
    mockedConsumer.open(new Configuration());
    Whitebox.setInternalState(mockedConsumer, "fetcher", // mock consumer as running.
    mockedFetcher);
    mockedConsumer.snapshotState(mock(FunctionSnapshotContext.class));
    assertEquals(true, listState.clearCalled);
    assertEquals(3, listState.getList().size());
    for (Tuple2<StreamShardMetadata, SequenceNumber> state : initialState) {
        for (Tuple2<StreamShardMetadata, SequenceNumber> currentState : listState.getList()) {
            assertNotEquals(state, currentState);
        }
    }
    for (Tuple2<StreamShardMetadata, SequenceNumber> state : expectedStateSnapshot) {
        boolean hasOneIsSame = false;
        for (Tuple2<StreamShardMetadata, SequenceNumber> currentState : listState.getList()) {
            hasOneIsSame = hasOneIsSame || state.equals(currentState);
        }
        assertEquals(true, hasOneIsSame);
    }
}

Example 5 with OperatorStateStore

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

the class FlinkKafkaConsumerBase method initializeState.

// ------------------------------------------------------------------------
// Checkpoint and restore
// ------------------------------------------------------------------------
@Override
public final void initializeState(FunctionInitializationContext context) throws Exception {
    OperatorStateStore stateStore = context.getOperatorStateStore();
    this.unionOffsetStates = stateStore.getUnionListState(new ListStateDescriptor<>(OFFSETS_STATE_NAME, createStateSerializer(getRuntimeContext().getExecutionConfig())));
    if (context.isRestored()) {
        restoredState = new TreeMap<>(new KafkaTopicPartition.Comparator());
        // populate actual holder for restored state
        for (Tuple2<KafkaTopicPartition, Long> kafkaOffset : unionOffsetStates.get()) {
            restoredState.put(kafkaOffset.f0, kafkaOffset.f1);
        }
        LOG.info("Consumer subtask {} restored state: {}.", getRuntimeContext().getIndexOfThisSubtask(), restoredState);
    } else {
        LOG.info("Consumer subtask {} has no restore state.", getRuntimeContext().getIndexOfThisSubtask());
    }
}