From c53bdfbf4c2394da9095849dc6c49d3bdaf7ccc1 Mon Sep 17 00:00:00 2001
From: Greg Eales <gregory.eales@ocado.com>
Date: Mon, 19 Oct 2020 11:47:57 +0100
Subject: [PATCH] GH-175: Iterate `listShards()` pages

Fixes https://github.com/spring-projects/spring-integration-aws/issues/175

`KinesisMessageDrivenChannelAdapter` doesn't read from all shards

* Add loop to get all pages of shards from Kinesis
* Upgrade to the latest SC-AWS
---
 build.gradle                                     |  2 +-
 .../KinesisMessageDrivenChannelAdapter.java      | 16 ++++++++++++----
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/build.gradle b/build.gradle
index b49c44e..00838ed 100644
--- a/build.gradle
+++ b/build.gradle
@@ -31,7 +31,7 @@ ext {
 	servletApiVersion = '4.0.1'
     localstackVersion = '0.1.22'
 	log4jVersion = '2.13.3'
-	springCloudAwsVersion = '2.2.3.RELEASE'
+	springCloudAwsVersion = '2.2.4.RELEASE'
 	springIntegrationVersion = '5.2.8.RELEASE'
 	kinesisClientVersion = '1.13.3'
 	kinesisProducerVersion = '0.14.1'
diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java
index 5929414..4038096 100644
--- a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java
+++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java
@@ -42,8 +42,6 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.stream.Collectors;
 
-import javax.annotation.Nullable;
-
 import org.springframework.beans.factory.DisposableBean;
 import org.springframework.context.ApplicationEventPublisher;
 import org.springframework.context.ApplicationEventPublisherAware;
@@ -64,6 +62,7 @@ import org.springframework.integration.support.locks.LockRegistry;
 import org.springframework.integration.support.management.IntegrationManagedResource;
 import org.springframework.jmx.export.annotation.ManagedOperation;
 import org.springframework.jmx.export.annotation.ManagedResource;
+import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.scheduling.SchedulingAwareRunnable;
 import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
@@ -542,8 +541,17 @@ public class KinesisMessageDrivenChannelAdapter extends MessageProducerSupport
 
 		try {
 			ListShardsResult listShardsResult = this.amazonKinesis.listShards(listShardsRequest);
-
-			shardList.addAll(listShardsResult.getShards());
+			while (true) {
+				shardList.addAll(listShardsResult.getShards());
+				if (listShardsResult.getNextToken() == null) {
+					break;
+				}
+				else {
+					listShardsResult =
+							this.amazonKinesis.listShards(new ListShardsRequest()
+									.withNextToken(listShardsResult.getNextToken()));
+				}
+			}
 
 		}
 		catch (LimitExceededException limitExceededException) {