GH-249 Added properties for handling backpressure in KplMessageHandler

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

@@ -49,6 +49,7 @@
 import org.springframework.expression.Expression;
 import org.springframework.expression.common.LiteralExpression;
 import org.springframework.integration.aws.support.AwsHeaders;
+import org.springframework.integration.aws.support.KPLBackpressureException;
 import org.springframework.integration.aws.support.UserRecordResponse;
 import org.springframework.integration.expression.ValueExpression;
 import org.springframework.integration.handler.AbstractMessageHandler;
@@ -68,6 +69,7 @@
  *
  * @author Arnaud Lecollaire
  * @author Artem Bilan
+ * @author Siddharth Jain
  *
  * @since 2.2
  *
@@ -99,6 +101,8 @@ public class KplMessageHandler extends AbstractAwsMessageHandler<Void> implement
 
 	private volatile ScheduledFuture<?> flushFuture;
 
+	private long backPressureThreshold = 0;
+
 	public KplMessageHandler(KinesisProducer kinesisProducer) {
 		Assert.notNull(kinesisProducer, "'kinesisProducer' must not be null.");
 		this.kinesisProducer = kinesisProducer;
@@ -115,6 +119,18 @@ public void setConverter(Converter<Object, byte[]> converter) {
 		setMessageConverter(new ConvertingFromMessageConverter(converter));
 	}
 
+	/**
+	 * Configure maximum records in flight for handling backpressure. By Default, backpressure handling is not enabled.
+	 * On number of records in flight exceeding the threshold, {@link KPLBackpressureException} would be thrown.
+	 * If Backpressure handling is enabled, {@link KPLBackpressureException} must be handled.
+	 * @param backPressureThreshold Defaulted to 0. Set a value greater than 0 to enable backpressure handling.
+	 * @since 3.0.9
+	 */
+	public void setBackPressureThreshold(long backPressureThreshold) {
+		Assert.isTrue(backPressureThreshold > 0, "'backPressureThreshold must be greater than 0.");
+		this.backPressureThreshold = backPressureThreshold;
+	}
+
 	/**
 	 * Configure a {@link MessageConverter} for converting payload to {@code byte[]} for Kinesis record.
 	 * @param messageConverter the {@link MessageConverter} to use.
@@ -367,7 +383,19 @@ private void setGlueSchemaIntoUserRecordIfAny(UserRecord userRecord, Message<?>
 		}
 	}
 
-	private CompletableFuture<UserRecordResponse> handleUserRecord(UserRecord userRecord) {
+	private CompletableFuture<UserRecordResponse> handleUserRecord(UserRecord userRecord)
+			throws KPLBackpressureException {
+
+		if (this.backPressureThreshold > 0) {
+			var numberOfRecordsInFlight = this.kinesisProducer.getOutstandingRecordsCount();
+			if (numberOfRecordsInFlight > this.backPressureThreshold) {
+				logger.error(String.format("Backpressure handling is enabled, Number of records in flight: %s is " +
+						"greater than backpressure threshold: %s" +
+						".", numberOfRecordsInFlight, this.backPressureThreshold));
+				throw new KPLBackpressureException("Buffer already at max capacity.");
+			}
+		}
+
 		ListenableFuture<UserRecordResult> recordResult = this.kinesisProducer.addUserRecord(userRecord);
 		return listenableFutureToCompletableFuture(recordResult)
 				.thenApply(UserRecordResponse::new);
@@ -403,7 +431,8 @@ private PutRecordRequest buildPutRecordRequest(Message<?> message) {
 			if (!StringUtils.hasText(partitionKey) && this.partitionKeyExpression != null) {
 				partitionKey = this.partitionKeyExpression.getValue(getEvaluationContext(), message, String.class);
 			}
-			Assert.state(partitionKey != null, "'partitionKey' must not be null for sending a Kinesis record. "
+			Assert.state(partitionKey != null,
+					"'partitionKey' must not be null for sending a Kinesis record. "
 					+ "Consider configuring this handler with a 'partitionKey'( or 'partitionKeyExpression') " +
 					"or supply an 'aws_partitionKey' message header.");
 

src/main/java/org/springframework/integration/aws/support/KPLBackpressureException.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2017-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.support;
+
+import org.springframework.messaging.MessagingException;
+
+/**
+ * An exception triggered from {@link org.springframework.integration.aws.outbound.KplMessageHandler} while sending
+ * records to kinesis when maximum number of records in flight exceeds the backpressure threshold.
+ * @author Siddharth Jain
+ * @since 3.0.9
+ */
+public class KPLBackpressureException extends MessagingException {
+	private static final long serialVersionUID = 1L;
+
+	public KPLBackpressureException(String message) {
+		super(message);
+	}
+}

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

@@ -0,0 +1,170 @@
+/*
+ * 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 org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+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.aws.support.KPLBackpressureException;
+import org.springframework.integration.config.EnableIntegration;
+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 3.0.9
+ */
+@SpringJUnitConfig
+@DirtiesContext
+public class KplMessageHandlerTests {
+
+	@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());
+		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());
+		verify(this.kinesisProducer, Mockito.times(0)).getOutstandingRecordsCount();
+		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_backpressure_capacity_available_test() {
+		given(this.kinesisProducer.addUserRecord(any(UserRecord.class)))
+				.willReturn(mock());
+		this.kplMessageHandler.setBackPressureThreshold(2);
+		given(this.kinesisProducer.getOutstandingRecordsCount())
+				.willReturn(1);
+		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());
+		verify(this.kinesisProducer, Mockito.times(1)).getOutstandingRecordsCount();
+		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_backpressure_insufficient_capacity_test() {
+		given(this.kinesisProducer.addUserRecord(any(UserRecord.class)))
+				.willReturn(mock());
+		this.kplMessageHandler.setBackPressureThreshold(2);
+		given(this.kinesisProducer.getOutstandingRecordsCount())
+				.willReturn(5);
+		final Message<?> message = MessageBuilder
+				.withPayload("message1")
+				.setHeader(AwsHeaders.PARTITION_KEY, "fooKey")
+				.setHeader(AwsHeaders.SEQUENCE_NUMBER, "10")
+				.setHeader("foo", "bar")
+				.build();
+
+		try {
+			this.kinesisSendChannel.send(message);
+		}
+		catch (Exception ex) {
+			assertThat(ex).isNotNull();
+			assertThat(ex.getCause()).isNotNull();
+			assertThat(ex.getCause().getClass()).isEqualTo(KPLBackpressureException.class);
+			assertThat(ex.getCause().getMessage()).isEqualTo("Buffer already at max capacity.");
+		}
+
+		verify(this.kinesisProducer, Mockito.times(0)).addUserRecord(any(UserRecord.class));
+		verify(this.kinesisProducer, Mockito.times(1)).getOutstandingRecordsCount();
+	}
+
+	@AfterEach
+	public void tearDown() {
+		clearInvocations(this.kinesisProducer);
+	}
+
+	@Configuration
+	@EnableIntegration
+	public static class ContextConfiguration {
+
+		@Bean
+		public KinesisProducer kinesisProducer() {
+			return mock();
+		}
+
+		@Bean
+		@ServiceActivator(inputChannel = "kinesisSendChannel")
+		public MessageHandler kplMessageHandler(KinesisProducer kinesisProducer) {
+			KplMessageHandler kplMessageHandler = new KplMessageHandler(kinesisProducer);
+			kplMessageHandler.setAsync(true);
+			kplMessageHandler.setStream("foo");
+			return kplMessageHandler;
+		}
+	}
+
+}