Examples with StreamShardHandle org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle used on opensource projects

Example 16 with StreamShardHandle

use of org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle in project flink by apache.

the class KinesisProxyTest method testGetShardList.

@Test
public void testGetShardList() throws Exception {
    List<String> shardIds = Arrays.asList("shardId-000000000000", "shardId-000000000001", "shardId-000000000002", "shardId-000000000003");
    String nextToken = "NextToken";
    String fakeStreamName = "fake-stream";
    List<Shard> shards = shardIds.stream().map(shardId -> new Shard().withShardId(shardId)).collect(Collectors.toList());
    AmazonKinesis mockClient = mock(AmazonKinesis.class);
    KinesisProxy kinesisProxy = getProxy(mockClient);
    ListShardsResult responseWithMoreData = new ListShardsResult().withShards(shards.subList(0, 2)).withNextToken(nextToken);
    ListShardsResult responseFinal = new ListShardsResult().withShards(shards.subList(2, shards.size())).withNextToken(null);
    doReturn(responseWithMoreData).when(mockClient).listShards(argThat(initialListShardsRequestMatcher()));
    doReturn(responseFinal).when(mockClient).listShards(argThat(listShardsNextToken(nextToken)));
    HashMap<String, String> streamHashMap = createInitialSubscribedStreamsToLastDiscoveredShardsState(Arrays.asList(fakeStreamName));
    GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);
    Assert.assertEquals(shardListResult.hasRetrievedShards(), true);
    Set<String> expectedStreams = new HashSet<>();
    expectedStreams.add(fakeStreamName);
    Assert.assertEquals(shardListResult.getStreamsWithRetrievedShards(), expectedStreams);
    List<StreamShardHandle> actualShardList = shardListResult.getRetrievedShardListOfStream(fakeStreamName);
    List<StreamShardHandle> expectedStreamShard = new ArrayList<>();
    assertThat(actualShardList, hasSize(4));
    for (int i = 0; i < 4; i++) {
        StreamShardHandle shardHandle = new StreamShardHandle(fakeStreamName, new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i)));
        expectedStreamShard.add(shardHandle);
    }
    Assert.assertThat(actualShardList, containsInAnyOrder(expectedStreamShard.toArray(new StreamShardHandle[actualShardList.size()])));
}

Example 17 with StreamShardHandle

use of org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle in project flink by apache.

the class KinesisProxyTest method testGetShardListWithNewShardsOnSecondRun.

@Test
public void testGetShardListWithNewShardsOnSecondRun() throws Exception {
    // given
    List<String> shardIds = Arrays.asList(KinesisShardIdGenerator.generateFromShardOrder(0), KinesisShardIdGenerator.generateFromShardOrder(1));
    String fakeStreamName = "fake-stream";
    List<Shard> shards = shardIds.stream().map(shardId -> new Shard().withShardId(shardId)).collect(Collectors.toList());
    AmazonKinesis mockClient = mock(AmazonKinesis.class);
    KinesisProxy kinesisProxy = getProxy(mockClient);
    ListShardsResult responseFirst = new ListShardsResult().withShards(shards).withNextToken(null);
    doReturn(responseFirst).when(mockClient).listShards(argThat(initialListShardsRequestMatcher()));
    HashMap<String, String> streamHashMap = createInitialSubscribedStreamsToLastDiscoveredShardsState(Collections.singletonList(fakeStreamName));
    // when
    GetShardListResult shardListResult = kinesisProxy.getShardList(streamHashMap);
    // then
    Assert.assertTrue(shardListResult.hasRetrievedShards());
    Set<String> expectedStreams = new HashSet<>();
    expectedStreams.add(fakeStreamName);
    Assert.assertEquals(shardListResult.getStreamsWithRetrievedShards(), expectedStreams);
    List<StreamShardHandle> actualShardList = shardListResult.getRetrievedShardListOfStream(fakeStreamName);
    Assert.assertThat(actualShardList, hasSize(2));
    List<StreamShardHandle> expectedStreamShard = IntStream.range(0, actualShardList.size()).mapToObj(i -> new StreamShardHandle(fakeStreamName, new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(i)))).collect(Collectors.toList());
    Assert.assertThat(actualShardList, containsInAnyOrder(expectedStreamShard.toArray(new StreamShardHandle[actualShardList.size()])));
    // given new shards
    ListShardsResult responseSecond = new ListShardsResult().withShards(new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))).withNextToken(null);
    doReturn(responseSecond).when(mockClient).listShards(argThat(initialListShardsRequestMatcher()));
    // when new shards
    GetShardListResult newShardListResult = kinesisProxy.getShardList(streamHashMap);
    // then new shards
    Assert.assertTrue(newShardListResult.hasRetrievedShards());
    Assert.assertEquals(newShardListResult.getStreamsWithRetrievedShards(), expectedStreams);
    List<StreamShardHandle> newActualShardList = newShardListResult.getRetrievedShardListOfStream(fakeStreamName);
    Assert.assertThat(newActualShardList, hasSize(1));
    List<StreamShardHandle> newExpectedStreamShard = Collections.singletonList(new StreamShardHandle(fakeStreamName, new Shard().withShardId(KinesisShardIdGenerator.generateFromShardOrder(2))));
    Assert.assertThat(newActualShardList, containsInAnyOrder(newExpectedStreamShard.toArray(new StreamShardHandle[newActualShardList.size()])));
}

Example 18 with StreamShardHandle

use of org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle in project flink by apache.

the class FlinkKinesisConsumerMigrationTest method writeSnapshot.

// ------------------------------------------------------------------------
@SuppressWarnings("unchecked")
private void writeSnapshot(String path, HashMap<StreamShardMetadata, SequenceNumber> state) throws Exception {
    final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(state.size());
    for (StreamShardMetadata shardMetadata : state.keySet()) {
        Shard shard = new Shard();
        shard.setShardId(shardMetadata.getShardId());
        SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
        sequenceNumberRange.withStartingSequenceNumber("1");
        shard.setSequenceNumberRange(sequenceNumberRange);
        initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
    }
    final TestFetcher<String> fetcher = new TestFetcher<>(Collections.singletonList(TEST_STREAM_NAME), new TestSourceContext<>(), new TestRuntimeContext(true, 1, 0), TestUtils.getStandardProperties(), new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()), state, initialDiscoveryShards);
    final DummyFlinkKinesisConsumer<String> consumer = new DummyFlinkKinesisConsumer<>(fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));
    StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumer);
    final AbstractStreamOperatorTestHarness<String> testHarness = new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);
    testHarness.setTimeCharacteristic(TimeCharacteristic.ProcessingTime);
    testHarness.setup();
    testHarness.open();
    final AtomicReference<Throwable> error = new AtomicReference<>();
    // run the source asynchronously
    Thread runner = new Thread() {

        @Override
        public void run() {
            try {
                consumer.run(new TestSourceContext<>());
            } catch (Throwable t) {
                t.printStackTrace();
                error.set(t);
            }
        }
    };
    runner.start();
    fetcher.waitUntilRun();
    final OperatorSubtaskState snapshot;
    synchronized (testHarness.getCheckpointLock()) {
        snapshot = testHarness.snapshot(0L, 0L);
    }
    OperatorSnapshotUtil.writeStateHandle(snapshot, path);
    consumerOperator.close();
    runner.join();
}

Example 19 with StreamShardHandle

use of org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle in project flink by apache.

the class FlinkKinesisConsumerMigrationTest method testRestoreWithEmptyState.

@Test
public void testRestoreWithEmptyState() throws Exception {
    final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(TEST_STATE.size());
    for (StreamShardMetadata shardMetadata : TEST_STATE.keySet()) {
        Shard shard = new Shard();
        shard.setShardId(shardMetadata.getShardId());
        SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
        sequenceNumberRange.withStartingSequenceNumber("1");
        shard.setSequenceNumberRange(sequenceNumberRange);
        initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
    }
    final TestFetcher<String> fetcher = new TestFetcher<>(Collections.singletonList(TEST_STREAM_NAME), new TestSourceContext<>(), new TestRuntimeContext(true, 1, 0), TestUtils.getStandardProperties(), new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()), null, initialDiscoveryShards);
    final DummyFlinkKinesisConsumer<String> consumerFunction = new DummyFlinkKinesisConsumer<>(fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));
    StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumerFunction);
    final AbstractStreamOperatorTestHarness<String> testHarness = new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);
    testHarness.setup();
    testHarness.initializeState(OperatorSnapshotUtil.getResourceFilename("kinesis-consumer-migration-test-flink" + testMigrateVersion + "-empty-snapshot"));
    testHarness.open();
    consumerFunction.run(new TestSourceContext<>());
    // assert that no state was restored
    assertTrue(consumerFunction.getRestoredState().isEmpty());
    // although the restore state is empty, the fetcher should still have been registered the
    // initial discovered shard;
    // furthermore, the discovered shard should be considered a newly created shard while the
    // job wasn't running,
    // and therefore should be consumed from the earliest sequence number
    KinesisStreamShardState restoredShardState = fetcher.getSubscribedShardsState().get(0);
    assertEquals(TEST_STREAM_NAME, restoredShardState.getStreamShardHandle().getStreamName());
    assertEquals(TEST_SHARD_ID, restoredShardState.getStreamShardHandle().getShard().getShardId());
    assertFalse(restoredShardState.getStreamShardHandle().isClosed());
    assertEquals(SentinelSequenceNumber.SENTINEL_EARLIEST_SEQUENCE_NUM.get(), restoredShardState.getLastProcessedSequenceNum());
    consumerOperator.close();
    consumerOperator.cancel();
}

Example 20 with StreamShardHandle

use of org.apache.flink.streaming.connectors.kinesis.model.StreamShardHandle in project flink by apache.

the class FlinkKinesisConsumerMigrationTest method testRestore.

@Test
public void testRestore() throws Exception {
    final List<StreamShardHandle> initialDiscoveryShards = new ArrayList<>(TEST_STATE.size());
    for (StreamShardMetadata shardMetadata : TEST_STATE.keySet()) {
        Shard shard = new Shard();
        shard.setShardId(shardMetadata.getShardId());
        SequenceNumberRange sequenceNumberRange = new SequenceNumberRange();
        sequenceNumberRange.withStartingSequenceNumber("1");
        shard.setSequenceNumberRange(sequenceNumberRange);
        initialDiscoveryShards.add(new StreamShardHandle(shardMetadata.getStreamName(), shard));
    }
    final TestFetcher<String> fetcher = new TestFetcher<>(Collections.singletonList(TEST_STREAM_NAME), new TestSourceContext<>(), new TestRuntimeContext(true, 1, 0), TestUtils.getStandardProperties(), new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()), null, initialDiscoveryShards);
    final DummyFlinkKinesisConsumer<String> consumerFunction = new DummyFlinkKinesisConsumer<>(fetcher, new KinesisDeserializationSchemaWrapper<>(new SimpleStringSchema()));
    StreamSource<String, DummyFlinkKinesisConsumer<String>> consumerOperator = new StreamSource<>(consumerFunction);
    final AbstractStreamOperatorTestHarness<String> testHarness = new AbstractStreamOperatorTestHarness<>(consumerOperator, 1, 1, 0);
    testHarness.setup();
    testHarness.initializeState(OperatorSnapshotUtil.getResourceFilename("kinesis-consumer-migration-test-flink" + testMigrateVersion + "-snapshot"));
    testHarness.open();
    consumerFunction.run(new TestSourceContext<>());
    // assert that state is correctly restored
    assertNotEquals(null, consumerFunction.getRestoredState());
    assertEquals(1, consumerFunction.getRestoredState().size());
    assertEquals(TEST_STATE, removeEquivalenceWrappers(consumerFunction.getRestoredState()));
    assertEquals(1, fetcher.getSubscribedShardsState().size());
    assertEquals(TEST_SEQUENCE_NUMBER, fetcher.getSubscribedShardsState().get(0).getLastProcessedSequenceNum());
    KinesisStreamShardState restoredShardState = fetcher.getSubscribedShardsState().get(0);
    assertEquals(TEST_STREAM_NAME, restoredShardState.getStreamShardHandle().getStreamName());
    assertEquals(TEST_SHARD_ID, restoredShardState.getStreamShardHandle().getShard().getShardId());
    assertFalse(restoredShardState.getStreamShardHandle().isClosed());
    assertEquals(TEST_SEQUENCE_NUMBER, restoredShardState.getLastProcessedSequenceNum());
    consumerOperator.close();
    consumerOperator.cancel();
}