GH-249 Added properties for handling backpressure in KplMessageHandle…

…r. Added Unit Tests.

src/main/java/org/springframework/integration/aws/outbound/KplMessageHandler.java

@@ -68,6 +68,7 @@
  *
  * @author Arnaud Lecollaire
  * @author Artem Bilan
+ * @author Siddharth Jain
  *
  * @since 2.2
  *
@@ -99,6 +100,10 @@ public class KplMessageHandler extends AbstractAwsMessageHandler<Void> implement
 
 	private volatile ScheduledFuture<?> flushFuture;
 
+	private long maxRecordsInFlight = 0;
+
+	private int maxRecordInFlightsSleepDurationInMillis = 100;
+
 	public KplMessageHandler(KinesisProducer kinesisProducer) {
 		Assert.notNull(kinesisProducer, "'kinesisProducer' must not be null.");
 		this.kinesisProducer = kinesisProducer;
@@ -115,8 +120,29 @@ public void setConverter(Converter<Object, byte[]> converter) {
 		setMessageConverter(new ConvertingFromMessageConverter(converter));
 	}
 
+	/**
+	 * When in KPL mode, the setting allows handling backpressure on the KPL native process. Setting this value would enable a sleep on the KPL Thread for the specified number of milliseconds defined in maxRecordInFlightsSleepDurationInMillis.
+	 *
+	 * @param maxRecordsInFlight Defaulted to 0. Value of 0 indicates that Backpressure handling is not enabled. Specify a positive value to enable back pressure.
+	 */
+	public void setMaxOutstandingRecordsInFlight(long maxRecordsInFlight) {
+		Assert.isTrue(maxRecordsInFlight > 0, "'maxRecordsInFlight must be greater than 0.");
+		this.maxRecordsInFlight = maxRecordsInFlight;
+	}
+
+	/**
+	 * The setting allows handling backpressure on the KPL native process. Enabled when maxOutstandingRecordsCount is greater than 0. The configurations puts the KPL Thread to sleep for the specified number of milliseconds.
+	 *
+	 * @param maxRecordInFlightsSleepDurationInMillis Default is 100ms.
+	 */
+	public void setMaxRecordInFlightsSleepDurationInMillis(int maxRecordInFlightsSleepDurationInMillis) {
+		Assert.isTrue(maxRecordInFlightsSleepDurationInMillis > 0, "'maxRecordInFlightsSleepDurationInMillis must be greater than 0.");
+		this.maxRecordInFlightsSleepDurationInMillis = maxRecordInFlightsSleepDurationInMillis;
+	}
+
 	/**
 	 * Configure a {@link MessageConverter} for converting payload to {@code byte[]} for Kinesis record.
+	 *
 	 * @param messageConverter the {@link MessageConverter} to use.
 	 * @since 2.3
 	 */
@@ -368,6 +394,16 @@ private void setGlueSchemaIntoUserRecordIfAny(UserRecord userRecord, Message<?>
 	}
 
 	private CompletableFuture<UserRecordResponse> handleUserRecord(UserRecord userRecord) {
+		if (this.maxRecordsInFlight != -1 && this.kinesisProducer.getOutstandingRecordsCount() > this.maxRecordsInFlight) {
+			try {
+				Thread.sleep(this.maxRecordInFlightsSleepDurationInMillis);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+				throw new RuntimeException(e);
+			}
+		}
+
 		ListenableFuture<UserRecordResult> recordResult = this.kinesisProducer.addUserRecord(userRecord);
 		return listenableFutureToCompletableFuture(recordResult)
 				.thenApply(UserRecordResponse::new);

src/test/java/org/springframework/integration/aws/outbound/KplMessageHandlerTest.java

@@ -0,0 +1,142 @@
+/*
+ * Copyright 2019-2024 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.integration.aws.outbound;
+
+import com.amazonaws.services.kinesis.producer.KinesisProducer;
+import com.amazonaws.services.kinesis.producer.UserRecord;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.integration.annotation.ServiceActivator;
+import org.springframework.integration.aws.support.AwsHeaders;
+import org.springframework.integration.config.EnableIntegration;
+import org.springframework.integration.support.json.EmbeddedJsonHeadersMessageMapper;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.MessageChannel;
+import org.springframework.messaging.MessageHandler;
+import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit.jupiter.SpringJUnitConfig;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.clearInvocations;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+/**
+ * @author Siddharth Jain
+ *
+ * @since 4.0
+ */
+@SpringJUnitConfig
+@DirtiesContext
+public class KplMessageHandlerTest {
+
+	@Autowired
+	protected KinesisProducer kinesisProducer;
+
+	@Autowired
+	protected MessageChannel kinesisSendChannel;
+
+	@Autowired
+	protected KplMessageHandler kplMessageHandler;
+
+	@Test
+	@SuppressWarnings("unchecked")
+	void testKPLMessageHandler_raw_payload_success() {
+		given(this.kinesisProducer.addUserRecord(any(UserRecord.class)))
+				.willReturn(mock(ListenableFuture.class));
+		final Message<?> message = MessageBuilder
+				.withPayload("message1")
+				.setHeader(AwsHeaders.PARTITION_KEY, "fooKey")
+				.setHeader(AwsHeaders.SEQUENCE_NUMBER, "10")
+				.setHeader("foo", "bar")
+				.build();
+
+
+		ArgumentCaptor<UserRecord> userRecordRequestArgumentCaptor = ArgumentCaptor
+				.forClass(UserRecord.class);
+
+		this.kinesisSendChannel.send(message);
+		verify(this.kinesisProducer).addUserRecord(userRecordRequestArgumentCaptor.capture());
+
+		UserRecord userRecord = userRecordRequestArgumentCaptor.getValue();
+		assertThat(userRecord.getStreamName()).isEqualTo("foo");
+		assertThat(userRecord.getPartitionKey()).isEqualTo("fooKey");
+		assertThat(userRecord.getExplicitHashKey()).isNull();
+	}
+
+	@Test
+	@SuppressWarnings("unchecked")
+	void testKPLMessageHandler_raw_payload_success_backpressure_test() {
+		given(this.kinesisProducer.addUserRecord(any(UserRecord.class)))
+				.willReturn(mock(ListenableFuture.class));
+		this.kplMessageHandler.setMaxOutstandingRecordsInFlight(1);
+		this.kplMessageHandler.setMaxRecordInFlightsSleepDurationInMillis(100);
+		given(this.kinesisProducer.getOutstandingRecordsCount()).willReturn(2);
+		final Message<?> message = MessageBuilder
+				.withPayload("message1")
+				.setHeader(AwsHeaders.PARTITION_KEY, "fooKey")
+				.setHeader(AwsHeaders.SEQUENCE_NUMBER, "10")
+				.setHeader("foo", "bar")
+				.build();
+
+
+		ArgumentCaptor<UserRecord> userRecordRequestArgumentCaptor = ArgumentCaptor
+				.forClass(UserRecord.class);
+
+		this.kinesisSendChannel.send(message);
+		verify(this.kinesisProducer).addUserRecord(userRecordRequestArgumentCaptor.capture());
+
+		UserRecord userRecord = userRecordRequestArgumentCaptor.getValue();
+		assertThat(userRecord.getStreamName()).isEqualTo("foo");
+		assertThat(userRecord.getPartitionKey()).isEqualTo("fooKey");
+		assertThat(userRecord.getExplicitHashKey()).isNull();
+	}
+
+	@AfterEach
+	public void tearDown() {
+		clearInvocations(this.kinesisProducer);
+	}
+
+	@Configuration
+	@EnableIntegration
+	public static class ContextConfiguration {
+
+		@Bean
+		public KinesisProducer kinesisProducer() {
+			return mock(KinesisProducer.class);
+		}
+
+		@Bean
+		@ServiceActivator(inputChannel = "kinesisSendChannel")
+		public MessageHandler kplMessageHandler(KinesisProducer kinesisProducer) {
+			KplMessageHandler kplMessageHandler = new KplMessageHandler(kinesisProducer);
+			kplMessageHandler.setAsync(true);
+			kplMessageHandler.setStream("foo");
+			kplMessageHandler.setEmbeddedHeadersMapper(new EmbeddedJsonHeadersMessageMapper("foo"));
+			return kplMessageHandler;
+		}
+	}
+}