Examples with KinesisDataFetcher org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher used on opensource projects

Example 1 with KinesisDataFetcher

use of org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher 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 2 with KinesisDataFetcher

use of org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher in project flink by apache.

the class FlinkKinesisConsumerTest method testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint.

// ----------------------------------------------------------------------
// Tests related to fetcher initialization
// ----------------------------------------------------------------------
@Test
@SuppressWarnings("unchecked")
public void testFetcherShouldNotBeRestoringFromFailureIfNotRestoringFromCheckpoint() throws Exception {
    KinesisDataFetcher mockedFetcher = mockKinesisDataFetcher();
    // assume the given config is correct
    PowerMockito.mockStatic(KinesisConfigUtil.class);
    PowerMockito.doNothing().when(KinesisConfigUtil.class);
    TestableFlinkKinesisConsumer consumer = new TestableFlinkKinesisConsumer("fakeStream", new Properties(), 10, 2);
    consumer.open(new Configuration());
    consumer.run(Mockito.mock(SourceFunction.SourceContext.class));
}

Example 3 with KinesisDataFetcher

use of org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher in project flink by apache.

the class FlinkKinesisConsumerTest method mockKinesisDataFetcher.

private static KinesisDataFetcher mockKinesisDataFetcher() throws Exception {
    KinesisDataFetcher mockedFetcher = Mockito.mock(KinesisDataFetcher.class);
    java.lang.reflect.Constructor<KinesisDataFetcher> ctor = (java.lang.reflect.Constructor<KinesisDataFetcher>) KinesisDataFetcher.class.getConstructors()[0];
    Class<?>[] otherParamTypes = new Class<?>[ctor.getParameterTypes().length - 1];
    System.arraycopy(ctor.getParameterTypes(), 1, otherParamTypes, 0, ctor.getParameterTypes().length - 1);
    Supplier<Object[]> argumentSupplier = () -> {
        Object[] otherParamArgs = new Object[otherParamTypes.length];
        for (int i = 0; i < otherParamTypes.length; i++) {
            otherParamArgs[i] = Mockito.nullable(otherParamTypes[i]);
        }
        return otherParamArgs;
    };
    PowerMockito.whenNew(ctor).withArguments(Mockito.any(ctor.getParameterTypes()[0]), argumentSupplier.get()).thenReturn(mockedFetcher);
    return mockedFetcher;
}

Example 4 with KinesisDataFetcher

use of org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher in project flink by apache.

the class FlinkKinesisConsumer method cancel.

@Override
public void cancel() {
    running = false;
    KinesisDataFetcher fetcher = this.fetcher;
    // so we must check if the fetcher is actually created
    if (fetcher != null) {
        try {
            // interrupt the fetcher of any work
            fetcher.shutdownFetcher();
        } catch (Exception e) {
            LOG.warn("Error while closing Kinesis data fetcher", e);
        }
    }
}

Example 5 with KinesisDataFetcher

use of org.apache.flink.streaming.connectors.kinesis.internals.KinesisDataFetcher 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)));
}