GH-249 Added properties for handling backpressure in KplMessageHandler

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,14 @@ public class KplMessageHandler extends AbstractAwsMessageHandler<Void> implement
 
 	private volatile ScheduledFuture<?> flushFuture;
 
+	private long maxInFlightRecords = 0;
+
+	private int maxInFlightRecordsInitialBackoffDuration = 100;
+
+	private int maxInFlightRecordsBackoffRate = 2;
+
+	private int maxInFlightRecordsBackoffMaxAttempts = 3;
+
 	public KplMessageHandler(KinesisProducer kinesisProducer) {
 		Assert.notNull(kinesisProducer, "'kinesisProducer' must not be null.");
 		this.kinesisProducer = kinesisProducer;
@@ -115,6 +124,62 @@ public void setConverter(Converter<Object, byte[]> converter) {
 		setMessageConverter(new ConvertingFromMessageConverter(converter));
 	}
 
+	/**
+	 * Configure maximum records in flight for handling backpressure. Used together with
+	 * {@link KplMessageHandler#maxInFlightRecordsInitialBackoffDuration}
+	 * @param maxInFlightRecords Defaulted to 0. Value of 0 indicates that Backpressure handling is not enabled.
+	 * @since 3.0.9
+	 * @see KplMessageHandler#setMaxInFlightRecordsInitialBackoffDuration
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffRate
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffMaxAttempts
+	 */
+	public void setMaxInFlightRecords(long maxInFlightRecords) {
+		Assert.isTrue(maxInFlightRecords > 0, "'maxInFlightRecords must be greater than 0.");
+		this.maxInFlightRecords = maxInFlightRecords;
+	}
+
+	/**
+	 * Configure initial backoff duration period in milliseconds when the number of records in flight
+	 * is greater than or equal to {@link  KplMessageHandler#maxInFlightRecords}.
+	 * The configuration helps in handling backpressure by sleeping the Thread using exponential backoff.
+	 * Enabled when {@link KplMessageHandler#maxInFlightRecords} is greater than 0.
+	 * @param maxInFlightRecordsInitialBackoffDuration  Initial backoff duration in milliseconds. Default is 100ms.
+	 * @since 3.0.9
+	 * @see KplMessageHandler#setMaxInFlightRecords
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffRate
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffMaxAttempts
+	 */
+	public void setMaxInFlightRecordsInitialBackoffDuration(int maxInFlightRecordsInitialBackoffDuration) {
+		Assert.isTrue(maxInFlightRecordsInitialBackoffDuration > 0,
+				"'maxInFlightRecordsBackoffDuration must be greater than 0.");
+		this.maxInFlightRecordsInitialBackoffDuration = maxInFlightRecordsInitialBackoffDuration;
+	}
+
+	/**
+	 * Configure exponential backoff rate when handling backpressure using {@link KplMessageHandler#maxInFlightRecords}.
+	 * @param maxInFlightRecordsBackoffRate Exponential back off rate. Default is 2
+	 * @since 3.0.9
+	 * @see KplMessageHandler#setMaxInFlightRecords
+	 * @see KplMessageHandler#setMaxInFlightRecordsInitialBackoffDuration
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffMaxAttempts
+	 */
+	public void setMaxInFlightRecordsBackoffRate(int maxInFlightRecordsBackoffRate) {
+		this.maxInFlightRecordsBackoffRate = maxInFlightRecordsBackoffRate;
+	}
+
+	/**
+	 * Configure maximum number of retry attempts with exponential backoff until there is available capacity. On
+	 * maximum attempts exhausted, RunTimeException is thrown.
+	 * @param maxInFlightRecordsBackoffMaxAttempts maximum of exponential retry attempts to waiting for capacity.
+	 * @since 3.0.9
+	 * @see KplMessageHandler#setMaxInFlightRecords
+	 * @see KplMessageHandler#setMaxInFlightRecordsInitialBackoffDuration
+	 * @see KplMessageHandler#setMaxInFlightRecordsBackoffRate
+	 */
+	public void setMaxInFlightRecordsBackoffMaxAttempts(int maxInFlightRecordsBackoffMaxAttempts) {
+		this.maxInFlightRecordsBackoffMaxAttempts = maxInFlightRecordsBackoffMaxAttempts;
+	}
+
 	/**
 	 * Configure a {@link MessageConverter} for converting payload to {@code byte[]} for Kinesis record.
 	 * @param messageConverter the {@link MessageConverter} to use.
@@ -368,11 +433,50 @@ private void setGlueSchemaIntoUserRecordIfAny(UserRecord userRecord, Message<?>
 	}
 
 	private CompletableFuture<UserRecordResponse> handleUserRecord(UserRecord userRecord) {
+		if (this.maxInFlightRecords > 0) {
+			waitForCapacityInUserRecordsBuffer();
+		}
+
 		ListenableFuture<UserRecordResult> recordResult = this.kinesisProducer.addUserRecord(userRecord);
 		return listenableFutureToCompletableFuture(recordResult)
 				.thenApply(UserRecordResponse::new);
 	}
 
+	private void waitForCapacityInUserRecordsBuffer() {
+		var kplOutstandingRecordsCount = this.kinesisProducer.getOutstandingRecordsCount();
+		var attempts = 1;
+		var sleepDuration = this.maxInFlightRecordsInitialBackoffDuration;
+		while (kplOutstandingRecordsCount >= this.maxInFlightRecords &&
+				attempts <= this.maxInFlightRecordsBackoffMaxAttempts) {
+			try {
+				logger.info("Buffer size: {} has reached the max records limit of {}. Attempt# {}".formatted(
+						kplOutstandingRecordsCount, this.maxInFlightRecords));
+				logger.info("Buffer sleeping for {} ms".formatted(
+						sleepDuration));
+				Thread.sleep(sleepDuration);
+			}
+			catch (InterruptedException ex) {
+				Thread.currentThread().interrupt();
+			}
+			finally {
+				sleepDuration = sleepDuration * this.maxInFlightRecordsBackoffRate;
+				attempts++;
+				kplOutstandingRecordsCount = this.kinesisProducer.getOutstandingRecordsCount();
+			}
+		}
+
+		if (kplOutstandingRecordsCount < this.maxInFlightRecords) {
+			logger.info("Buffer cleared on number of attempts: {}".formatted(attempts));
+			return;
+		}
+
+		if (attempts > this.maxInFlightRecordsBackoffMaxAttempts) {
+			logger.error("Buffer not cleared after maximum {} number of attempts & {} sleepDuration".formatted(attempts,
+					sleepDuration));
+			throw new RuntimeException("Buffer already at max capacity.");
+		}
+	}
+
 	private PutRecordRequest buildPutRecordRequest(Message<?> message) {
 		Object payload = message.getPayload();
 
@@ -403,7 +507,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/test/java/org/springframework/integration/aws/outbound/KplMessageHandlerTests.java

@@ -0,0 +1,212 @@
+/*
+ * 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.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.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(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());
+		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(ListenableFuture.class));
+		this.kplMessageHandler.setMaxInFlightRecords(2);
+		this.kplMessageHandler.setMaxInFlightRecordsInitialBackoffDuration(100);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffMaxAttempts(2);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffRate(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(ListenableFuture.class));
+		this.kplMessageHandler.setMaxInFlightRecords(2);
+		this.kplMessageHandler.setMaxInFlightRecordsInitialBackoffDuration(100);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffMaxAttempts(2);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffRate(2);
+		given(this.kinesisProducer.getOutstandingRecordsCount())
+				.willReturn(3)
+				.willReturn(2)
+				.willReturn(1)
+				.willReturn(0);
+		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(3)).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_failure_test() {
+		given(this.kinesisProducer.addUserRecord(any(UserRecord.class)))
+				.willReturn(mock(ListenableFuture.class));
+		this.kplMessageHandler.setMaxInFlightRecords(2);
+		this.kplMessageHandler.setMaxInFlightRecordsInitialBackoffDuration(100);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffMaxAttempts(2);
+		this.kplMessageHandler.setMaxInFlightRecordsBackoffRate(2);
+		given(this.kinesisProducer.getOutstandingRecordsCount())
+				.willReturn(5)
+				.willReturn(4)
+				.willReturn(3);
+		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(RuntimeException.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(3)).getOutstandingRecordsCount();
+	}
+
+	@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");
+			return kplMessageHandler;
+		}
+	}
+
+}