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

Example 11 with OperatorStateStore

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

the class FlinkKafkaConsumerBaseTest method checkRestoredNullCheckpointWhenFetcherNotReady.

/**
	 * Tests that no checkpoints happen when the fetcher is not running.
	 */
@Test
public void checkRestoredNullCheckpointWhenFetcherNotReady() throws Exception {
    FlinkKafkaConsumerBase<String> consumer = getConsumer(null, new LinkedMap(), true);
    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);
    when(initializationContext.isRestored()).thenReturn(false);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.snapshotState(new StateSnapshotContextSynchronousImpl(17, 17));
    assertFalse(listState.get().iterator().hasNext());
}

Example 12 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored.

/**
 * FLINK-8484: ensure that a state change in the StreamShardMetadata other than {@link
 * StreamShardMetadata#getShardId()} or {@link StreamShardMetadata#getStreamName()} does not
 * result in the shard not being able to be restored. This handles the corner case where the
 * stored shard metadata is open (no ending sequence number), but after the job restore, the
 * shard has been closed (ending number set) due to re-sharding, and we can no longer rely on
 * {@link StreamShardMetadata#equals(Object)} to find back the sequence number in the collection
 * of restored shard metadata.
 *
 * <p>Therefore, we will rely on synchronizing the snapshot's state with the Kinesis shard
 * before attempting to find back the sequence number to restore.
 */
@Test
public void testFindSequenceNumberToRestoreFromIfTheShardHasBeenClosedSinceTheStateWasStored() throws Exception {
    // ----------------------------------------------------------------------
    // setup initial state
    // ----------------------------------------------------------------------
    HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");
    // ----------------------------------------------------------------------
    // mock operator state backend and initial state for initializeState()
    // ----------------------------------------------------------------------
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
        listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
    }
    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 fetcher
    // ----------------------------------------------------------------------
    KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
    List<StreamShardHandle> shards = new ArrayList<>();
    // create a fake stream shard handle based on the first entry in the restored state
    final StreamShardHandle originalStreamShardHandle = fakeRestoredState.keySet().iterator().next();
    final StreamShardHandle closedStreamShardHandle = new StreamShardHandle(originalStreamShardHandle.getStreamName(), originalStreamShardHandle.getShard());
    // close the shard handle by setting an ending sequence number
    final SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
    sequenceNumberRange.setEndingSequenceNumber("1293844");
    closedStreamShardHandle.getShard().setSequenceNumberRange(sequenceNumberRange);
    shards.add(closedStreamShardHandle);
    when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    // ----------------------------------------------------------------------
    // start to test fetcher's initial state seeding
    // ----------------------------------------------------------------------
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(closedStreamShardHandle), closedStreamShardHandle, fakeRestoredState.get(closedStreamShardHandle)));
}

Example 13 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint.

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededIfRestoringFromCheckpoint() throws Exception {
    // ----------------------------------------------------------------------
    // setup initial state
    // ----------------------------------------------------------------------
    HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");
    // ----------------------------------------------------------------------
    // mock operator state backend and initial state for initializeState()
    // ----------------------------------------------------------------------
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
        listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
    }
    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 fetcher
    // ----------------------------------------------------------------------
    KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
    List<StreamShardHandle> shards = new ArrayList<>();
    shards.addAll(fakeRestoredState.keySet());
    when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    // ----------------------------------------------------------------------
    // start to test fetcher's initial state seeding
    // ----------------------------------------------------------------------
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
        Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()), restoredShard.getKey(), restoredShard.getValue()));
    }
}

Example 14 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testFetcherShouldBeCorrectlySeededOnlyItsOwnStates.

@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededOnlyItsOwnStates() throws Exception {
    // ----------------------------------------------------------------------
    // setup initial state
    // ----------------------------------------------------------------------
    HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("fakeStream1");
    HashMap<StreamShardHandle, SequenceNumber> fakeRestoredStateForOthers = getFakeRestoredStore("fakeStream2");
    // ----------------------------------------------------------------------
    // mock operator state backend and initial state for initializeState()
    // ----------------------------------------------------------------------
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
        listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
    }
    for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredStateForOthers.entrySet()) {
        listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
    }
    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 fetcher
    // ----------------------------------------------------------------------
    KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
    List<StreamShardHandle> shards = new ArrayList<>();
    shards.addAll(fakeRestoredState.keySet());
    when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    // ----------------------------------------------------------------------
    // start to test fetcher's initial state seeding
    // ----------------------------------------------------------------------
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredStateForOthers.entrySet()) {
        // should never get restored state not belonging to itself
        Mockito.verify(mockedFetcher, never()).registerNewSubscribedShardState(new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()), restoredShard.getKey(), restoredShard.getValue()));
    }
    for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
        // should get restored state belonging to itself
        Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()), restoredShard.getKey(), restoredShard.getValue()));
    }
}

Example 15 with OperatorStateStore

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

the class FlinkKinesisConsumerTest method testFetcherShouldBeCorrectlySeededWithNewDiscoveredKinesisStreamShard.

/*
     * This tests that the consumer correctly picks up shards that were not discovered on the previous run.
     *
     * Case under test:
     *
     * If the original parallelism is 2 and states are:
     *   Consumer subtask 1:
     *     stream1, shard1, SequentialNumber(xxx)
     *   Consumer subtask 2:
     *     stream1, shard2, SequentialNumber(yyy)
     *
     * After discoverNewShardsToSubscribe() if there were two shards (shard3, shard4) created:
     *   Consumer subtask 1 (late for discoverNewShardsToSubscribe()):
     *     stream1, shard1, SequentialNumber(xxx)
     *   Consumer subtask 2:
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, shard4, SequentialNumber(zzz)
     *
     * If snapshotState() occurs and parallelism is changed to 1:
     *   Union state will be:
     *     stream1, shard1, SequentialNumber(xxx)
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, shard4, SequentialNumber(zzz)
     *   Fetcher should be seeded with:
     *     stream1, shard1, SequentialNumber(xxx)
     *     stream1, shard2, SequentialNumber(yyy)
     *     stream1, share3, SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM
     *     stream1, shard4, SequentialNumber(zzz)
     */
@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldBeCorrectlySeededWithNewDiscoveredKinesisStreamShard() throws Exception {
    // ----------------------------------------------------------------------
    // setup initial state
    // ----------------------------------------------------------------------
    HashMap<StreamShardHandle, SequenceNumber> fakeRestoredState = getFakeRestoredStore("all");
    // ----------------------------------------------------------------------
    // mock operator state backend and initial state for initializeState()
    // ----------------------------------------------------------------------
    TestingListState<Tuple2<StreamShardMetadata, SequenceNumber>> listState = new TestingListState<>();
    for (Map.Entry<StreamShardHandle, SequenceNumber> state : fakeRestoredState.entrySet()) {
        listState.add(Tuple2.of(KinesisDataFetcher.convertToStreamShardMetadata(state.getKey()), state.getValue()));
    }
    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 fetcher
    // ----------------------------------------------------------------------
    KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
    List<StreamShardHandle> shards = new ArrayList<>();
    shards.addAll(fakeRestoredState.keySet());
    shards.add(new StreamShardHandle("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))));
    when(mockedFetcher.discoverNewShardsToSubscribe()).thenReturn(shards);
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    // ----------------------------------------------------------------------
    // start to test fetcher's initial state seeding
    // ----------------------------------------------------------------------
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.initializeState(initializationContext);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
    fakeRestoredState.put(new StreamShardHandle("fakeStream2", new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))), SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get());
    for (Map.Entry<StreamShardHandle, SequenceNumber> restoredShard : fakeRestoredState.entrySet()) {
        Mockito.verify(mockedFetcher).registerNewSubscribedShardState(new KinesisStreamShardState(KinesisDataFetcher.convertToStreamShardMetadata(restoredShard.getKey()), restoredShard.getKey(), restoredShard.getValue()));
    }
}