Add bindSourceRecord for Kinesis Channel Adapters

Author: artembilan

build.gradle

@@ -12,7 +12,7 @@ plugins {
 	id 'eclipse'
 	id 'idea'
 	id 'jacoco'
-	id 'org.sonarqube' version '2.7'
+	id 'org.sonarqube' version '2.7.1'
 	id 'checkstyle'
 	id 'org.ajoberstar.grgit' version '3.1.1'
 }
@@ -34,11 +34,11 @@ ext {
 	assertjVersion = '3.12.2'
 	awaitilityVersion = '3.1.6'
 	dynamodbLockClientVersion = '1.1.0'
-	jacksonVersion = '2.9.8'
+	jacksonVersion = '2.9.9'
 	servletApiVersion = '4.0.1'
 	log4jVersion = '2.11.2'
 	springCloudAwsVersion = '2.1.1.RELEASE'
-	springIntegrationVersion = '5.1.4.RELEASE'
+	springIntegrationVersion = '5.1.6.RELEASE'
 	kinesisClientVersion = '1.10.0'
 	kinesisProducerVersion = '0.12.11'
 
@@ -87,7 +87,7 @@ jacoco {
 
 checkstyle {
 	configFile = file("${rootDir}/src/checkstyle/checkstyle.xml")
-	toolVersion = "8.20"
+	toolVersion = "8.21"
 }
 
 dependencies {

src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java

@@ -24,6 +24,7 @@
 import org.springframework.core.task.SimpleAsyncTaskExecutor;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.core.task.support.ExecutorServiceAdapter;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.aws.support.AwsHeaders;
 import org.springframework.integration.endpoint.MessageProducerSupport;
 import org.springframework.integration.mapping.InboundMessageMapper;
@@ -105,6 +106,8 @@ public class KclMessageDrivenChannelAdapter extends MessageProducerSupport {
 
 	private String workerId = UUID.randomUUID().toString();
 
+	private boolean bindSourceRecord;
+
 	public KclMessageDrivenChannelAdapter(String streams) {
 		this(streams, AmazonKinesisClientBuilder.defaultClient(),
 				AmazonCloudWatchClientBuilder.defaultClient(), AmazonDynamoDBClientBuilder.defaultClient(),
@@ -188,6 +191,17 @@ public void setWorkerId(String workerId) {
 		this.workerId = workerId;
 	}
 
+	/**
+	 * Set to true to bind the source consumer record in the header named
+	 * {@link IntegrationMessageHeaderAccessor#SOURCE_DATA}.
+	 * Does not apply to batch listeners.
+	 * @param bindSourceRecord true to bind.
+	 * @since 2.2
+	 */
+	public void setBindSourceRecord(boolean bindSourceRecord) {
+		this.bindSourceRecord = bindSourceRecord;
+	}
+
 	@Override
 	protected void onInit() {
 		super.onInit();
@@ -364,6 +378,10 @@ private AbstractIntegrationMessageBuilder<Object> prepareMessageForRecord(Record
 					.setHeader(AwsHeaders.RECEIVED_STREAM, KclMessageDrivenChannelAdapter.this.stream)
 					.setHeader(AwsHeaders.SHARD, this.shardId);
 
+			if (KclMessageDrivenChannelAdapter.this.bindSourceRecord) {
+				messageBuilder.setHeader(IntegrationMessageHeaderAccessor.SOURCE_DATA, record);
+			}
+
 			if (messageToUse != null) {
 				messageBuilder.copyHeadersIfAbsent(messageToUse.getHeaders());
 			}

src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java

@@ -46,6 +46,7 @@
 import org.springframework.core.AttributeAccessor;
 import org.springframework.core.convert.converter.Converter;
 import org.springframework.core.serializer.support.DeserializingConverter;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.aws.support.AwsHeaders;
 import org.springframework.integration.endpoint.MessageProducerSupport;
 import org.springframework.integration.mapping.InboundMessageMapper;
@@ -159,6 +160,8 @@ public class KinesisMessageDrivenChannelAdapter extends MessageProducerSupport i
 
 	private LockRegistry lockRegistry;
 
+	private boolean bindSourceRecord;
+
 	private volatile boolean active;
 
 	private volatile int consumerInvokerMaxCapacity;
@@ -235,7 +238,7 @@ public void setCheckpointMode(CheckpointMode checkpointMode) {
 	/**
 	 * Sets the interval between 2 checkpoints. Only used when checkpointMode is periodic.
 	 * @param checkpointsInterval interval between 2 checkpoints (in milliseconds)
-	 * @since 2.2.0
+	 * @since 2.2
 	 */
 	public void setCheckpointsInterval(long checkpointsInterval) {
 		this.checkpointsInterval = checkpointsInterval;
@@ -311,6 +314,17 @@ public void setLockRegistry(LockRegistry lockRegistry) {
 		this.lockRegistry = lockRegistry;
 	}
 
+	/**
+	 * Set to true to bind the source consumer record in the header named
+	 * {@link IntegrationMessageHeaderAccessor#SOURCE_DATA}.
+	 * Does not apply to batch listeners.
+	 * @param bindSourceRecord true to bind.
+	 * @since 2.2
+	 */
+	public void setBindSourceRecord(boolean bindSourceRecord) {
+		this.bindSourceRecord = bindSourceRecord;
+	}
+
 	@Override
 	protected void onInit() {
 		super.onInit();
@@ -527,7 +541,8 @@ private void populateShardsForStream(final String stream, final CountDownLatch s
 					}
 
 					if (describeStreamResult == null ||
-							!StreamStatus.ACTIVE.toString().equals(describeStreamResult.getStreamDescription().getStreamStatus())) {
+							!StreamStatus.ACTIVE.toString().equals(
+									describeStreamResult.getStreamDescription().getStreamStatus())) {
 
 						if (describeStreamRetries++ > this.describeStreamRetries) {
 							ResourceNotFoundException resourceNotFoundException =
@@ -541,7 +556,7 @@ private void populateShardsForStream(final String stream, final CountDownLatch s
 							continue;
 						}
 						catch (InterruptedException e) {
-							Thread.interrupted();
+							Thread.currentThread().interrupt();
 							throw new IllegalStateException("The [describeStream] thread for the stream ["
 									+ stream + "] has been interrupted.", e);
 						}
@@ -729,10 +744,9 @@ public void run() {
 					}
 				}
 
-				for (Iterator<ShardConsumer> iterator =
+				Iterator<ShardConsumer> iterator =
 						KinesisMessageDrivenChannelAdapter.this.shardConsumers.values().iterator();
-						iterator.hasNext(); ) {
-
+				while (iterator.hasNext()) {
 					ShardConsumer shardConsumer = iterator.next();
 					shardConsumer.execute();
 					if (ConsumerState.STOP == shardConsumer.state) {
@@ -795,7 +809,8 @@ private final class ShardConsumer {
 		ShardConsumer(KinesisShardOffset shardOffset) {
 			this.shardOffset = new KinesisShardOffset(shardOffset);
 			this.key = buildCheckpointKeyForShard(shardOffset.getStream(), shardOffset.getShard());
-			this.checkpointer = new ShardCheckpointer(KinesisMessageDrivenChannelAdapter.this.checkpointStore, this.key);
+			this.checkpointer = new ShardCheckpointer(KinesisMessageDrivenChannelAdapter.this.checkpointStore,
+					this.key);
 		}
 
 		void setNotifier(Runnable notifier) {
@@ -838,7 +853,8 @@ void execute() {
 								if (logger.isInfoEnabled() && this.state == ConsumerState.NEW) {
 									logger.info("The [" + this + "] has been started.");
 								}
-								GetShardIteratorRequest shardIteratorRequest = this.shardOffset.toShardIteratorRequest();
+								GetShardIteratorRequest shardIteratorRequest =
+										this.shardOffset.toShardIteratorRequest();
 								this.shardIterator =
 										KinesisMessageDrivenChannelAdapter.this.amazonKinesis
 												.getShardIterator(shardIteratorRequest)
@@ -1034,8 +1050,8 @@ private void processRecords(List<Record> records) {
 			}
 			else if (CheckpointMode.periodic.equals(KinesisMessageDrivenChannelAdapter.this.checkpointMode) &&
 					System.currentTimeMillis() > nextCheckpointTimeInMillis) {
-					this.checkpointer.checkpoint();
-					this.nextCheckpointTimeInMillis = System.currentTimeMillis() + checkpointsInterval;
+				this.checkpointer.checkpoint();
+				this.nextCheckpointTimeInMillis = System.currentTimeMillis() + checkpointsInterval;
 			}
 		}
 
@@ -1068,6 +1084,10 @@ private AbstractIntegrationMessageBuilder<Object> prepareMessageForRecord(Record
 							.setHeader(AwsHeaders.RECEIVED_PARTITION_KEY, record.getPartitionKey())
 							.setHeader(AwsHeaders.RECEIVED_SEQUENCE_NUMBER, record.getSequenceNumber());
 
+			if (KinesisMessageDrivenChannelAdapter.this.bindSourceRecord) {
+				messageBuilder.setHeader(IntegrationMessageHeaderAccessor.SOURCE_DATA, record);
+			}
+
 			if (messageToUse != null) {
 				messageBuilder.copyHeadersIfAbsent(messageToUse.getHeaders());
 			}
@@ -1157,8 +1177,8 @@ public void run() {
 								shardConsumer.task.run();
 							}
 							catch (Exception e) {
-								logger.info("Got an exception " + e + " during [" + shardConsumer + "] task invocation.\n" +
-										"Process will be retried on the next iteration.");
+								logger.info("Got an exception " + e + " during [" + shardConsumer + "] task invocation"
+										+ ".\nProcess will be retried on the next iteration.");
 							}
 						}
 					}

src/test/java/org/springframework/integration/aws/kinesis/KinesisIntegrationTests.java

@@ -33,6 +33,7 @@
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.integration.IntegrationMessageHeaderAccessor;
 import org.springframework.integration.annotation.ServiceActivator;
 import org.springframework.integration.aws.KinesisLocalRunning;
 import org.springframework.integration.aws.inbound.kinesis.KinesisMessageDrivenChannelAdapter;
@@ -90,7 +91,7 @@ public static void tearDown() {
 	}
 
 	@Test
-	public void testKinesisInboundOutbound() throws InterruptedException {
+	public void testKinesisInboundOutbound() {
 		this.kinesisSendChannel.send(
 				MessageBuilder.withPayload("foo")
 						.setHeader(AwsHeaders.STREAM, TEST_STREAM)
@@ -107,12 +108,13 @@ public void testKinesisInboundOutbound() throws InterruptedException {
 		assertThat(receive).isNotNull();
 		assertThat(receive.getPayload()).isEqualTo(now);
 		assertThat(receive.getHeaders()).contains(entry("foo", "BAR"));
+		assertThat(receive.getHeaders()).containsKey(IntegrationMessageHeaderAccessor.SOURCE_DATA);
 
 		Message<?> errorMessage = this.errorChannel.receive(10_000);
 		assertThat(errorMessage).isNotNull();
 		assertThat(errorMessage.getHeaders().get(AwsHeaders.RAW_RECORD)).isNotNull();
 		assertThat(((Exception) errorMessage.getPayload()).getMessage())
-				.contains("Channel 'kinesisReceiveChannel' expected one of the following datataypes " +
+				.contains("Channel 'kinesisReceiveChannel' expected one of the following data types " +
 						"[class java.util.Date], but received [class java.lang.String]");
 
 
@@ -171,6 +173,7 @@ private KinesisMessageDrivenChannelAdapter kinesisMessageDrivenChannelAdapter()
 			adapter.setCheckpointStore(checkpointStore());
 			adapter.setLockRegistry(lockRegistry());
 			adapter.setEmbeddedHeadersMapper(new EmbeddedJsonHeadersMessageMapper("foo"));
+			adapter.setBindSourceRecord(true);
 
 			DirectFieldAccessor dfa = new DirectFieldAccessor(adapter);
 			dfa.setPropertyValue("describeStreamBackoff", 10);