From ab2f9b89a6c127e013199dffce48d6e33d58b7c5 Mon Sep 17 00:00:00 2001
From: Artem Bilan <artem.bilan@broadcom.com>
Date: Wed, 25 Sep 2024 16:04:32 -0400
Subject: [PATCH] Expose KCL `pollingMaxRecords` & `pollingIdleTime` options

---
 .../KclMessageDrivenChannelAdapter.java       | 28 ++++++++++++++++++-
 .../KclMessageDrivenChannelAdapterTests.java  | 10 +++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
index 22bdfba..82a184f 100644
--- a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
+++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java
@@ -170,6 +170,10 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 
 	private boolean emptyRecordList;
 
+	private int pollingMaxRecords = PollingConfig.DEFAULT_MAX_RECORDS;
+
+	private long pollingIdleTime = 1500L;
+
 	public KclMessageDrivenChannelAdapter(String... streams) {
 		this(KinesisAsyncClient.create(), CloudWatchAsyncClient.create(), DynamoDbAsyncClient.create(), streams);
 	}
@@ -369,6 +373,26 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 		this.emptyRecordList = emptyRecordList;
 	}
 
+	/**
+	 * The number of records to poll from Kinesis when using {@link PollingConfig}.
+	 * @param pollingMaxRecords the number of records to poll from Kinesis.
+	 * @since 3.0.8
+	 * @see PollingConfig#maxRecords(int)
+	 */
+	public void setPollingMaxRecords(int pollingMaxRecords) {
+		this.pollingMaxRecords = pollingMaxRecords;
+	}
+
+	/**
+	 * The idle timeout between polls when using {@link PollingConfig}.
+	 * @param pollingIdleTime idle timeout between polls.
+	 * @since 3.0.8
+	 * @see PollingConfig#idleTimeBetweenReadsInMillis(long)
+	 */
+	public void setPollingIdleTime(long pollingIdleTime) {
+		this.pollingIdleTime = pollingIdleTime;
+	}
+
 	@Override
 	protected void onInit() {
 		super.onInit();
@@ -425,7 +449,9 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport
 		else {
 			retrievalSpecificConfig =
 					new PollingConfig(this.kinesisClient)
-							.streamName(singleStreamName);
+							.streamName(singleStreamName)
+							.maxRecords(this.pollingMaxRecords)
+							.idleTimeBetweenReadsInMillis(this.pollingIdleTime);
 		}
 
 		RetrievalConfig retrievalConfig = this.config.retrievalConfig()
diff --git a/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java b/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
index d5b757b..8571300 100644
--- a/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
+++ b/src/test/java/org/springframework/integration/aws/kinesis/KclMessageDrivenChannelAdapterTests.java
@@ -162,6 +162,15 @@ public class KclMessageDrivenChannelAdapterTests implements LocalstackContainerT
 		assertThat(shardConsumerDispatchPollIntervalMillis).isEqualTo(500L);
 	}
 
+	@Test
+	public void pollingMaxRecordsIsPropagated() {
+		Integer maxRecords =
+				TestUtils.getPropertyValue(this.kclMessageDrivenChannelAdapter,
+						"scheduler.retrievalConfig.retrievalSpecificConfig.maxRecords",
+						Integer.class);
+		assertThat(maxRecords).isEqualTo(99);
+	}
+
 	@Configuration
 	@EnableIntegration
 	public static class TestConfiguration {
@@ -184,6 +193,7 @@ public class KclMessageDrivenChannelAdapterTests implements LocalstackContainerT
 					coordinatorConfig.shardConsumerDispatchPollIntervalMillis(500L));
 			adapter.setBindSourceRecord(true);
 			adapter.setEmptyRecordList(true);
+			adapter.setPollingMaxRecords(99);
 			return adapter;
 		}