software.amazon.awssdk.services.kinesis.model.GetRecordsResponse Java Examples

@Test
public void testSingleSuccessfulRequestFuture() throws Exception {
    AsynchronousGetRecordsRetrievalStrategy strategy = new AsynchronousGetRecordsRetrievalStrategy(dataFetcher,
            executorService, (int) RETRY_GET_RECORDS_IN_SECONDS, completionServiceSupplier, SHARD_ID);

    when(executorService.isShutdown()).thenReturn(false);
    when(completionService.submit(any())).thenReturn(successfulFuture);
    when(completionService.poll(anyLong(), any())).thenReturn(successfulFuture);
    when(successfulFuture.get()).thenReturn(dataFetcherResult);

    GetRecordsResponse result = strategy.getRecords(10);

    verify(executorService).isShutdown();
    verify(completionService).submit(any());
    verify(completionService).poll(eq(RETRY_GET_RECORDS_IN_SECONDS), eq(TimeUnit.SECONDS));
    verify(successfulFuture).get();
    verify(successfulFuture).cancel(eq(true));

    assertThat(result, equalTo(expectedResponses));
}
 

private GetRecordsResponse tryNext() {
    GetRecordsResponse response = kinesisClient.getRecords(GetRecordsRequest.builder()
            .shardIterator(id)
            .limit(fetchRecordLimit)
            .build())
            .join();
    this.id = response.nextShardIterator();
    LOG.debug("next() with id " + this.id + " returned " + response.records().size() + " records");
    if (!response.records().isEmpty()) {
        this.shardPosition = fromPosition(
                shardPosition.shardName(),
                response.records().get(response.records().size()-1).sequenceNumber()
        );
    }
    return response;
}
 

@Test
public void shouldImplementEqualsAndHashCode() {
    final Instant now = now();
    final GetRecordsResponse response = GetRecordsResponse
            .builder()
            .records(Record.builder()
                    .sequenceNumber("1")
                    .partitionKey("first")
                    .approximateArrivalTimestamp(now)
                    .build())
            .nextShardIterator("nextIter")
            .millisBehindLatest(0L)
            .build();
    final ShardResponse first = kinesisShardResponse(fromPosition("shard", "42"), response);
    final ShardResponse second = kinesisShardResponse(fromPosition("shard", "42"), response);

    assertThat(first.equals(second), is(true));
    assertThat(first.hashCode(), is(second.hashCode()));
}
 

@Test
public void testBlockedAndSuccessfulFuture() throws Exception {
    AsynchronousGetRecordsRetrievalStrategy strategy = new AsynchronousGetRecordsRetrievalStrategy(dataFetcher,
            executorService, (int) RETRY_GET_RECORDS_IN_SECONDS, completionServiceSupplier, SHARD_ID);

    when(executorService.isShutdown()).thenReturn(false);
    when(completionService.submit(any())).thenReturn(blockedFuture).thenReturn(successfulFuture);
    when(completionService.poll(anyLong(), any())).thenReturn(null).thenReturn(successfulFuture);
    when(successfulFuture.get()).thenReturn(dataFetcherResult);
    when(successfulFuture.cancel(anyBoolean())).thenReturn(false);
    when(blockedFuture.cancel(anyBoolean())).thenReturn(true);
    when(successfulFuture.isCancelled()).thenReturn(false);
    when(blockedFuture.isCancelled()).thenReturn(true);

    GetRecordsResponse actualResults = strategy.getRecords(10);

    verify(completionService, times(2)).submit(any());
    verify(completionService, times(2)).poll(eq(RETRY_GET_RECORDS_IN_SECONDS), eq(TimeUnit.SECONDS));
    verify(successfulFuture).get();
    verify(blockedFuture, never()).get();
    verify(successfulFuture).cancel(eq(true));
    verify(blockedFuture).cancel(eq(true));

    assertThat(actualResults, equalTo(expectedResponses));
}
 

@Test
public void testPoolOutOfResources() throws Exception {
    AsynchronousGetRecordsRetrievalStrategy strategy = new AsynchronousGetRecordsRetrievalStrategy(dataFetcher,
            executorService, (int) RETRY_GET_RECORDS_IN_SECONDS, completionServiceSupplier, SHARD_ID);

    when(executorService.isShutdown()).thenReturn(false);
    when(completionService.submit(any())).thenReturn(blockedFuture).thenThrow(new RejectedExecutionException("Rejected!")).thenReturn(successfulFuture);
    when(completionService.poll(anyLong(), any())).thenReturn(null).thenReturn(null).thenReturn(successfulFuture);
    when(successfulFuture.get()).thenReturn(dataFetcherResult);
    when(successfulFuture.cancel(anyBoolean())).thenReturn(false);
    when(blockedFuture.cancel(anyBoolean())).thenReturn(true);
    when(successfulFuture.isCancelled()).thenReturn(false);
    when(blockedFuture.isCancelled()).thenReturn(true);

    GetRecordsResponse actualResult = strategy.getRecords(10);

    verify(completionService, times(3)).submit(any());
    verify(completionService, times(3)).poll(eq(RETRY_GET_RECORDS_IN_SECONDS), eq(TimeUnit.SECONDS));
    verify(successfulFuture).cancel(eq(true));
    verify(blockedFuture).cancel(eq(true));


    assertThat(actualResult, equalTo(expectedResponses));
}
 

private DataFetcherResult assertAdvanced(GetRecordsResponse expectedResult, String previousValue,
        String nextValue) {
    DataFetcherResult acceptResult = kinesisDataFetcher.getRecords();
    assertEquals(expectedResult, acceptResult.getResult());

    assertEquals(previousValue, kinesisDataFetcher.getNextIterator());
    assertFalse(kinesisDataFetcher.isShardEndReached());

    assertEquals(expectedResult, acceptResult.accept());
    assertEquals(nextValue, kinesisDataFetcher.getNextIterator());
    if (nextValue == null) {
        assertTrue(kinesisDataFetcher.isShardEndReached());
    }

    return acceptResult;
}
 

@Test
@Ignore
public void testRestartIterator() throws Exception {
    GetRecordsResponse getRecordsResult = mock(GetRecordsResponse.class);
    GetRecordsResponse restartGetRecordsResponse = makeGetRecordsResponse(null, null).get(anyLong(),
            any(TimeUnit.class));
    Record record = mock(Record.class);
    final String nextShardIterator = "NextShardIterator";
    final String sequenceNumber = "SequenceNumber";

    when(getRecordsResult.records()).thenReturn(Collections.singletonList(record));
    when(getRecordsResult.nextShardIterator()).thenReturn(nextShardIterator);
    when(record.sequenceNumber()).thenReturn(sequenceNumber);

    kinesisDataFetcher.initialize(InitialPositionInStream.LATEST.toString(), INITIAL_POSITION_LATEST);
    assertEquals(getRecordsResult, kinesisDataFetcher.getRecords().accept());

    kinesisDataFetcher.restartIterator();
    assertEquals(restartGetRecordsResponse, kinesisDataFetcher.getRecords().accept());
}
 

@Before
public void setup() {
    dataFetcher = spy(new KinesisDataFetcherForTests(kinesisClient, streamName, shardId, numberOfRecords));
    rejectedExecutionHandler = spy(new ThreadPoolExecutor.AbortPolicy());
    executorService = spy(new ThreadPoolExecutor(
            CORE_POOL_SIZE,
            MAX_POOL_SIZE,
            TIME_TO_LIVE,
            TimeUnit.SECONDS,
            new LinkedBlockingQueue<>(1),
            new ThreadFactoryBuilder().setDaemon(true).setNameFormat("getrecords-worker-%d").build(),
            rejectedExecutionHandler));
    completionService = spy(new ExecutorCompletionService<DataFetcherResult>(executorService));
    getRecordsRetrivalStrategy = new AsynchronousGetRecordsRetrievalStrategy(dataFetcher, executorService,
            RETRY_GET_RECORDS_IN_SECONDS, completionServiceSupplier, "shardId-0001");
    getRecordsResponse = GetRecordsResponse.builder().build();

    when(completionServiceSupplier.get()).thenReturn(completionService);
    when(result.accept()).thenReturn(getRecordsResponse);
}
 

@Before
public void setup() {
    when(getRecordsRetrievalStrategy.getDataFetcher()).thenReturn(dataFetcher);

    executorService = spy(Executors.newFixedThreadPool(1));
    getRecordsCache = new PrefetchRecordsPublisher(
            MAX_SIZE,
            3 * SIZE_1_MB,
            MAX_RECORDS_COUNT,
            MAX_RECORDS_PER_CALL,
            getRecordsRetrievalStrategy,
            executorService,
            IDLE_MILLIS_BETWEEN_CALLS,
            new NullMetricsFactory(),
            operation,
            "shardId");
    spyQueue = spy(getRecordsCache.getPublisherSession().prefetchRecordsQueue());
    records = spy(new ArrayList<>());
    getRecordsResponse = GetRecordsResponse.builder().records(records).build();

    when(getRecordsRetrievalStrategy.getRecords(eq(MAX_RECORDS_PER_CALL))).thenReturn(getRecordsResponse);
}
 

@Test
public void testNonNullGetRecords() throws Exception {
    final String nextIterator = "TestIterator";
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    DataFetcherResult dataFetcherResult = kinesisDataFetcher.getRecords();

    assertNotNull(dataFetcherResult);
    assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
    assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
}
 

public ShardResponse next() {
    if (!stopSignal.get()) {
        GetRecordsResponse recordsResponse = tryNextWithRetry();
        if (recordsResponse.records() == null || recordsResponse.records().size() == 0) {
            LOG.warn("GetRecordsResponse contains no records.");
        }
        if (recordsResponse.millisBehindLatest() == null || recordsResponse.millisBehindLatest() == Long.MAX_VALUE) {
            LOG.warn("GetRecordsResponse millisBehindLatest was not set.");
        }
        return KinesisShardResponse.kinesisShardResponse(shardPosition, recordsResponse);
    } else {
        throw new IllegalStateException(format("Cannot iterate on shard '%s' after stop signal was received", shardPosition.shardName()));
    }
}
 

@Before
public void before() {
    expectedResponses = GetRecordsResponse.builder().build();

    when(completionServiceSupplier.get()).thenReturn(completionService);
    when(dataFetcherResult.getResult()).thenReturn(expectedResponses);
    when(dataFetcherResult.accept()).thenReturn(expectedResponses);
}
 

@Override
public GetRecordsResponse answer(InvocationOnMock invocation) throws Throwable {
    GetRecordsResponse response = iterator.next();
    if (!iterator.hasNext()) {
        iterator = responses.iterator();
    }
    return response;
}
 

public void resetIteratorTo(String nextIterator) {
    Iterator<GetRecordsResponse> newIterator = responses.iterator();
    while(newIterator.hasNext()) {
        GetRecordsResponse current = newIterator.next();
        if (StringUtils.equals(nextIterator, current.nextShardIterator())) {
            if (!newIterator.hasNext()) {
                iterator = responses.iterator();
            } else {
                newIterator.next();
                iterator = newIterator;
            }
            break;
        }
    }
}
 

@Test
public void oneRequestMultithreadTest() {
    when(result.accept()).thenReturn(null);
    GetRecordsResponse getRecordsResult = getRecordsRetrivalStrategy.getRecords(numberOfRecords);
    verify(dataFetcher, atLeast(getLeastNumberOfCalls())).getRecords();
    verify(executorService, atLeast(getLeastNumberOfCalls())).execute(any());
    assertNull(getRecordsResult);
}
 

@Override
public GetRecordsResponse accept() {
    nextIterator = result.nextShardIterator();
    if (result.records() != null && !result.records().isEmpty()) {
        lastKnownSequenceNumber = Iterables.getLast(result.records()).sequenceNumber();
    }
    if (nextIterator == null) {
        isShardEndReached = true;
    }
    return getResult();
}
 

@Test
public void multiRequestTest() {
    ExecutorCompletionService<DataFetcherResult> completionService1 = spy(new ExecutorCompletionService<DataFetcherResult>(executorService));
    when(completionServiceSupplier.get()).thenReturn(completionService1);
    GetRecordsResponse getRecordsResult = getRecordsRetrivalStrategy.getRecords(numberOfRecords);
    verify(dataFetcher, atLeast(getLeastNumberOfCalls())).getRecords();
    verify(executorService, atLeast(getLeastNumberOfCalls())).execute(any());
    assertThat(getRecordsResult, equalTo(getRecordsResponse));

    when(result.accept()).thenReturn(null);
    ExecutorCompletionService<DataFetcherResult> completionService2 = spy(new ExecutorCompletionService<DataFetcherResult>(executorService));
    when(completionServiceSupplier.get()).thenReturn(completionService2);
    getRecordsResult = getRecordsRetrivalStrategy.getRecords(numberOfRecords);
    assertThat(getRecordsResult, nullValue(GetRecordsResponse.class));
}
 

@Test
public void testGetRecordsWithResourceNotFoundException() throws Exception {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    // Set up arguments used by proxy
    final String nextIterator = "TestShardIterator";

    final GetShardIteratorRequest expectedIteratorRequest = makeGetShardIteratorRequest(
            ShardIteratorType.LATEST.name());
    final GetRecordsRequest expectedRecordsRequest = makeGetRecordsRequest(nextIterator);

    final CompletableFuture<GetRecordsResponse> future = mock(CompletableFuture.class);

    // Set up proxy mock methods
    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(nextIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(future);
    when(future.get(anyLong(), any(TimeUnit.class))).thenThrow(
            new ExecutionException(ResourceNotFoundException.builder().message("Test Exception").build()));

    // Create data fectcher and initialize it with latest type checkpoint
    kinesisDataFetcher.initialize(SentinelCheckpoint.LATEST.toString(), INITIAL_POSITION_LATEST);
    final GetRecordsRetrievalStrategy getRecordsRetrievalStrategy = new SynchronousGetRecordsRetrievalStrategy(
            kinesisDataFetcher);
    try {
        // Call records of dataFetcher which will throw an exception
        getRecordsRetrievalStrategy.getRecords(MAX_RECORDS);
    } finally {
        // Test shard has reached the end
        assertTrue("Shard should reach the end", kinesisDataFetcher.isShardEndReached());
        assertEquals(expectedIteratorRequest.startingSequenceNumber(), iteratorCaptor.getValue().startingSequenceNumber());
        assertEquals(expectedRecordsRequest.shardIterator(), recordsCaptor.getValue().shardIterator());
    }
}
 

public ProcessRecordsInput getNextResult() {
    GetRecordsResponse getRecordsResult = getRecordsRetrievalStrategy.getRecords(maxRecordsPerCall);
    final RequestDetails getRecordsRequestDetails = new RequestDetails(getRecordsResult.responseMetadata().requestId(), Instant.now().toString());
    setLastSuccessfulRequestDetails(getRecordsRequestDetails);
    List<KinesisClientRecord> records = getRecordsResult.records().stream()
            .map(KinesisClientRecord::fromRecord).collect(Collectors.toList());
    return ProcessRecordsInput.builder()
            .records(records)
            .millisBehindLatest(getRecordsResult.millisBehindLatest())
            .build();
}
 

private DataFetcherResult assertNoAdvance(final GetRecordsResponse expectedResult, final String previousValue) {
    assertEquals(previousValue, kinesisDataFetcher.getNextIterator());
    DataFetcherResult noAcceptResult = kinesisDataFetcher.getRecords();
    assertEquals(expectedResult, noAcceptResult.getResult());

    assertEquals(previousValue, kinesisDataFetcher.getNextIterator());

    return noAcceptResult;
}
 

@Test
public void shouldConvertRecordsToMessages() {
    final Instant firstArrival = now().minus(1, SECONDS);
    final Instant secondArrival = now();
    final GetRecordsResponse recordsResponse = GetRecordsResponse
            .builder()
            .records(
                    Record.builder()
                            .sequenceNumber("1")
                            .approximateArrivalTimestamp(firstArrival)
                            .partitionKey("first")
                            .data(SdkBytes.fromByteArray("content".getBytes(UTF_8)))
                            .build(),
                    Record.builder()
                            .sequenceNumber("2")
                            .approximateArrivalTimestamp(secondArrival)
                            .partitionKey("second")
                            .data(SdkBytes.fromByteArray("content".getBytes(UTF_8)))
                            .build()
                    )
            .nextShardIterator("nextIter")
            .millisBehindLatest(1L)
            .build();
    final ShardResponse response = kinesisShardResponse(fromPosition("shard", "42"), recordsResponse);

    assertThat(response.getShardName(), is("shard"));
    assertThat(response.getDurationBehind(), is(Duration.ofMillis(1L)));
    assertThat(response.getShardPosition(), is(fromPosition("shard", "42")));
    assertThat(response.getMessages(), contains(
            message("first", Header.of(fromPosition("shard", "1"), ImmutableMap.of(MSG_ARRIVAL_TS.key(), firstArrival.toString())), "content"),
            message("second", Header.of(fromPosition("shard", "2"), ImmutableMap.of(MSG_ARRIVAL_TS.key(), secondArrival.toString())), "content")
    ));
}
 

public static ShardResponse kinesisShardResponse(final ShardPosition shardPosition,
                                                 final GetRecordsResponse recordsResponse) {
    final KinesisDecoder kinesisDecoder = new KinesisDecoder();
    return shardResponse(
            shardPosition,
            ofMillis(recordsResponse.millisBehindLatest() == null ? Long.MAX_VALUE : recordsResponse.millisBehindLatest()),
            recordsResponse.records()
                    .stream()
                    .map(record -> kinesisDecoder.apply(new RecordWithShard(shardPosition.shardName(), record)))
                    .collect(toImmutableList())
    );
}
 

private GetRecordsResponse tryNextWithRetry() {
    RuntimeException exception = null;
    int retry = 0;
    while (retry++ < MAX_RETRIES) {
        try {
            return tryNext();
        } catch (RuntimeException e) {
            exception = e;
            LOG.warn("Failed to iterate on kinesis shard. Try to reset iterator and retry ({}/{}).", retry, MAX_RETRIES);
            id = createShardIteratorId();
        }
    }
    throw exception;
}
 

private CompletableFuture<GetRecordsResponse> makeGetRecordsResponse(String nextIterator, List<Record> records)
        throws InterruptedException, ExecutionException {
    return CompletableFuture.completedFuture(GetRecordsResponse.builder().nextShardIterator(nextIterator)
            .records(CollectionUtils.isNullOrEmpty(records) ? Collections.emptyList() : records).build());
}
 

@Test
public void testFetcherDoesNotAdvanceWithoutAccept() throws InterruptedException, ExecutionException {
    final ArgumentCaptor<GetShardIteratorRequest> iteratorCaptor = ArgumentCaptor
            .forClass(GetShardIteratorRequest.class);
    final ArgumentCaptor<GetRecordsRequest> recordsCaptor = ArgumentCaptor.forClass(GetRecordsRequest.class);
    final String initialIterator = "InitialIterator";
    final String nextIterator1 = "NextIteratorOne";
    final String nextIterator2 = "NextIteratorTwo";
    final CompletableFuture<GetRecordsResponse> nonAdvancingResult1 = makeGetRecordsResponse(initialIterator, null);
    final CompletableFuture<GetRecordsResponse> nonAdvancingResult2 = makeGetRecordsResponse(nextIterator1, null);
    final CompletableFuture<GetRecordsResponse> finalNonAdvancingResult = makeGetRecordsResponse(nextIterator2,
            null);
    final CompletableFuture<GetRecordsResponse> advancingResult1 = makeGetRecordsResponse(nextIterator1, null);
    final CompletableFuture<GetRecordsResponse> advancingResult2 = makeGetRecordsResponse(nextIterator2, null);
    final CompletableFuture<GetRecordsResponse> finalAdvancingResult = makeGetRecordsResponse(null, null);

    when(kinesisClient.getShardIterator(iteratorCaptor.capture()))
            .thenReturn(makeGetShardIteratorResonse(initialIterator));
    when(kinesisClient.getRecords(recordsCaptor.capture())).thenReturn(nonAdvancingResult1, advancingResult1,
            nonAdvancingResult2, advancingResult2, finalNonAdvancingResult, finalAdvancingResult);

    kinesisDataFetcher.initialize("TRIM_HORIZON",
            InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.TRIM_HORIZON));

    assertNoAdvance(nonAdvancingResult1.get(), initialIterator);
    assertAdvanced(advancingResult1.get(), initialIterator, nextIterator1);
    verify(kinesisClient, times(2)).getRecords(any(GetRecordsRequest.class));

    assertNoAdvance(nonAdvancingResult2.get(), nextIterator1);
    assertAdvanced(advancingResult2.get(), nextIterator1, nextIterator2);
    verify(kinesisClient, times(4)).getRecords(any(GetRecordsRequest.class));

    assertNoAdvance(finalNonAdvancingResult.get(), nextIterator2);
    assertAdvanced(finalAdvancingResult.get(), nextIterator2, null);
    verify(kinesisClient, times(6)).getRecords(any(GetRecordsRequest.class));



    reset(kinesisClient);

    DataFetcherResult terminal = kinesisDataFetcher.getRecords();
    assertTrue(terminal.isShardEnd());
    assertNotNull(terminal.getResult());

    final GetRecordsResponse terminalResult = terminal.getResult();
    assertNotNull(terminalResult.records());
    assertEquals(0, terminalResult.records().size());
    assertNull(terminalResult.nextShardIterator());
    assertEquals(kinesisDataFetcher.TERMINAL_RESULT, terminal);

    verify(kinesisClient, never()).getRecords(any(GetRecordsRequest.class));
}
 

private void testGets(final String streamName, final Shard shard) {
    // Get an iterator for the first shard.
    GetShardIteratorResponse iteratorResult = client.getShardIterator(
            GetShardIteratorRequest.builder()
                                   .streamName(streamName)
                                   .shardId(shard.shardId())
                                   .shardIteratorType(ShardIteratorType.AT_SEQUENCE_NUMBER)
                                   .startingSequenceNumber(shard.sequenceNumberRange().startingSequenceNumber())
                                   .build());
    Assert.assertNotNull(iteratorResult);

    String iterator = iteratorResult.shardIterator();
    Assert.assertNotNull(iterator);

    int tries = 0;
    GetRecordsResponse result;
    List<Record> records;

    // Read the first record from the first shard (looping until it's
    // available).
    while (true) {
        tries += 1;
        if (tries > 100) {
            Assert.fail("Failed to read any records after 100 seconds");
        }

        result = client.getRecords(GetRecordsRequest.builder()
                                                    .shardIterator(iterator)
                                                    .limit(1)
                                                    .build());
        Assert.assertNotNull(result);
        Assert.assertNotNull(result.records());
        Assert.assertNotNull(result.nextShardIterator());

        records = result.records();
        if (records.size() > 0) {
            long arrivalTime = records.get(0).approximateArrivalTimestamp().toEpochMilli();
            Long delta = Math.abs(Instant.now().minusMillis(arrivalTime).toEpochMilli());
            // Assert that the arrival date is within 5 minutes of the current date to make sure it unmarshalled correctly.
            assertThat(delta, Matchers.lessThan(60 * 5000L));
            break;
        }

        try {
            Thread.sleep(1000);
        } catch (InterruptedException exception) {
            throw new RuntimeException(exception);
        }

        iterator = result.nextShardIterator();
    }

    System.out.println("  [Succeeded after " + tries + " tries]");
    Assert.assertEquals(1, records.size());
    validateRecord(records.get(0), "See No Evil");

    // Read the second record from the first shard.
    result = client.getRecords(GetRecordsRequest.builder()
                                                .shardIterator(result.nextShardIterator())
                                                .build());
    Assert.assertNotNull(result);
    Assert.assertNotNull(result.records());
    Assert.assertNotNull(result.nextShardIterator());

    records = result.records();
    Assert.assertEquals(1, records.size());
    validateRecord(records.get(0), "See No Evil");

    // Try to read some more, get EOF.
    result = client.getRecords(GetRecordsRequest.builder()
                                                .shardIterator(result.nextShardIterator())
                                                .build());
    Assert.assertNotNull(result);
    Assert.assertNotNull(result.records());
    Assert.assertTrue(result.records().isEmpty());
    Assert.assertNull(result.nextShardIterator());
}
 

@Override
public DataFetcherResult getRecords() {
    GetRecordsResponse getRecordsResult = GetRecordsResponse.builder().records(new ArrayList<>(records)).millisBehindLatest(1000L).build();

    return new AdvancingResult(getRecordsResult);
}
 

public RetrieverAnswer(List<GetRecordsResponse> responses) {
    this.responses = responses;
    this.iterator = responses.iterator();
}
 

@Test
public void testResetClearsRemainingData() {
    List<GetRecordsResponse> responses = Stream.iterate(0, i -> i + 1).limit(10).map(i -> {
        Record record = Record.builder().partitionKey("record-" + i).sequenceNumber("seq-" + i)
                .data(SdkBytes.fromByteArray(new byte[] { 1, 2, 3 })).approximateArrivalTimestamp(Instant.now())
                .build();
        String nextIterator = "shard-iter-" + (i + 1);
        return GetRecordsResponse.builder().records(record).nextShardIterator(nextIterator).build();
    }).collect(Collectors.toList());

    RetrieverAnswer retrieverAnswer = new RetrieverAnswer(responses);

    when(getRecordsRetrievalStrategy.getRecords(anyInt())).thenAnswer(retrieverAnswer);
    doAnswer(a -> {
        String resetTo = a.getArgumentAt(0, String.class);
        retrieverAnswer.resetIteratorTo(resetTo);
        return null;
    }).when(dataFetcher).resetIterator(anyString(), anyString(), any());

    getRecordsCache.start(sequenceNumber, initialPosition);

    RecordsRetrieved lastProcessed = blockUntilRecordsAvailable(() -> evictPublishedEvent(getRecordsCache, "shardId"), 1000);
    RecordsRetrieved expected = blockUntilRecordsAvailable(() -> evictPublishedEvent(getRecordsCache, "shardId"), 1000);

    //
    // Skip some of the records the cache
    //
    blockUntilRecordsAvailable(() -> evictPublishedEvent(getRecordsCache, "shardId"), 1000);
    blockUntilRecordsAvailable(() -> evictPublishedEvent(getRecordsCache, "shardId"), 1000);

    verify(getRecordsRetrievalStrategy, atLeast(2)).getRecords(anyInt());

    while(getRecordsCache.getPublisherSession().prefetchRecordsQueue().remainingCapacity() > 0) {
        Thread.yield();
    }

    getRecordsCache.restartFrom(lastProcessed);
    RecordsRetrieved postRestart = blockUntilRecordsAvailable(() -> evictPublishedEvent(getRecordsCache, "shardId"), 1000);

    assertThat(postRestart.processRecordsInput(), eqProcessRecordsInput(expected.processRecordsInput()));
    verify(dataFetcher).resetIterator(eq(responses.get(0).nextShardIterator()),
            eq(responses.get(0).records().get(0).sequenceNumber()), any());

}
 

@Test(timeout = 10000L)
public void testNoDeadlockOnFullQueueAndLossOfNotification() {
    //
    // Fixes https://github.com/awslabs/amazon-kinesis-client/issues/602
    //
    // This test is to verify that the data consumption is not stuck in the case of an failed event delivery
    // to the subscriber.
    GetRecordsResponse response = GetRecordsResponse.builder().records(
            Record.builder().data(SdkBytes.fromByteArray(new byte[] { 1, 2, 3 })).sequenceNumber("123").build())
            .build();
    when(getRecordsRetrievalStrategy.getRecords(anyInt())).thenReturn(response);

    getRecordsCache.start(sequenceNumber, initialPosition);

    //
    // Wait for the queue to fill up, and the publisher to block on adding items to the queue.
    //
    log.info("Waiting for queue to fill up");
    while (getRecordsCache.getPublisherSession().prefetchRecordsQueue().size() < MAX_SIZE) {
        Thread.yield();
    }

    log.info("Queue is currently at {} starting subscriber", getRecordsCache.getPublisherSession().prefetchRecordsQueue().size());
    AtomicInteger receivedItems = new AtomicInteger(0);

    final int expectedItems = MAX_SIZE * 20;

    Object lock = new Object();

    Subscriber<RecordsRetrieved> delegateSubscriber = new Subscriber<RecordsRetrieved>() {
        Subscription sub;

        @Override
        public void onSubscribe(Subscription s) {
            sub = s;
            s.request(1);
        }

        @Override
        public void onNext(RecordsRetrieved recordsRetrieved) {
            receivedItems.incrementAndGet();
            if (receivedItems.get() >= expectedItems) {
                synchronized (lock) {
                    log.info("Notifying waiters");
                    lock.notifyAll();
                }
                sub.cancel();
            } else {
                sub.request(1);
            }
        }

        @Override
        public void onError(Throwable t) {
            log.error("Caught error", t);
            throw new RuntimeException(t);
        }

        @Override
        public void onComplete() {
            fail("onComplete not expected in this test");
        }
    };

    Subscriber<RecordsRetrieved> subscriber = new LossyNotificationSubscriber(delegateSubscriber, getRecordsCache);

    synchronized (lock) {
        log.info("Awaiting notification");
        Flowable.fromPublisher(getRecordsCache).subscribeOn(Schedulers.computation())
                .observeOn(Schedulers.computation(), true, 8).subscribe(subscriber);
        try {
            lock.wait();
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }
    }
    verify(getRecordsRetrievalStrategy, atLeast(expectedItems)).getRecords(anyInt());
    assertThat(receivedItems.get(), equalTo(expectedItems));
}