diff --git a/build.gradle b/build.gradle index 7d8031e..2f60b54 100644 --- a/build.gradle +++ b/build.gradle @@ -37,6 +37,8 @@ ext { log4jVersion = '2.11.1' springCloudAwsVersion = '2.1.0.BUILD-SNAPSHOT' springIntegrationVersion = '5.1.0.RELEASE' + kinesisClientVersion = '2.0.5' + kinesisProducerVersion = '0.12.11' idPrefix = 'aws' @@ -93,7 +95,8 @@ dependencies { compile('org.springframework.integration:spring-integration-file', optional) compile('org.springframework.integration:spring-integration-http', optional) -// compile('com.amazonaws:amazon-kinesis-client:1.9.1', optional) + compile("software.amazon.kinesis:amazon-kinesis-client:$kinesisClientVersion", optional) + compile("com.amazonaws:amazon-kinesis-producer:$kinesisProducerVersion", optional) compile('com.amazonaws:aws-java-sdk-kinesis', optional) compile('com.amazonaws:aws-java-sdk-dynamodb', optional) 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 new file mode 100644 index 0000000..d21bd80 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KclMessageDrivenChannelAdapter.java @@ -0,0 +1,407 @@ +/* + * Copyright 2017-2018 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 + * + * http://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.inbound.kinesis; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.UUID; +import java.util.concurrent.Executor; + +import org.springframework.beans.factory.DisposableBean; +import org.springframework.core.AttributeAccessor; +import org.springframework.integration.aws.support.AwsHeaders; +import org.springframework.integration.endpoint.MessageProducerSupport; +import org.springframework.integration.mapping.InboundMessageMapper; +import org.springframework.integration.support.AbstractIntegrationMessageBuilder; +import org.springframework.integration.support.ErrorMessageStrategy; +import org.springframework.integration.support.ErrorMessageUtils; +import org.springframework.integration.support.management.IntegrationManagedResource; +import org.springframework.jmx.export.annotation.ManagedResource; +import org.springframework.messaging.Message; +import org.springframework.util.Assert; + +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.cloudwatch.CloudWatchAsyncClient; +import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient; +import software.amazon.awssdk.services.kinesis.KinesisAsyncClient; +import software.amazon.kinesis.common.ConfigsBuilder; +import software.amazon.kinesis.common.InitialPositionInStream; +import software.amazon.kinesis.common.InitialPositionInStreamExtended; +import software.amazon.kinesis.coordinator.Scheduler; +import software.amazon.kinesis.exceptions.InvalidStateException; +import software.amazon.kinesis.exceptions.ShutdownException; +import software.amazon.kinesis.exceptions.ThrottlingException; +import software.amazon.kinesis.lifecycle.events.InitializationInput; +import software.amazon.kinesis.lifecycle.events.LeaseLostInput; +import software.amazon.kinesis.lifecycle.events.ProcessRecordsInput; +import software.amazon.kinesis.lifecycle.events.ShardEndedInput; +import software.amazon.kinesis.lifecycle.events.ShutdownRequestedInput; +import software.amazon.kinesis.processor.RecordProcessorCheckpointer; +import software.amazon.kinesis.processor.ShardRecordProcessor; +import software.amazon.kinesis.processor.ShardRecordProcessorFactory; +import software.amazon.kinesis.retrieval.KinesisClientRecord; + +/** + * The {@link MessageProducerSupport} implementation for receiving data from Amazon + * Kinesis stream(s) using AWS KCL. + * + * @author Hervé Fortin + * + * @since 2.1.0 + */ +@ManagedResource +@IntegrationManagedResource +public class KclMessageDrivenChannelAdapter extends MessageProducerSupport implements DisposableBean { + + private static final ThreadLocal attributesHolder = new ThreadLocal<>(); + + private final String stream; + + private String consumerGroup = "SpringIntegration"; + + private InboundMessageMapper embeddedHeadersMapper; + + private Scheduler scheduler; + + private final Executor executor; + + private final KinesisAsyncClient kinesisClient; + + private final CloudWatchAsyncClient cloudWatchClient; + + private final DynamoDbAsyncClient dynamoDBClient; + + private InitialPositionInStreamExtended streamInitialSequence = + InitialPositionInStreamExtended.newInitialPosition(InitialPositionInStream.LATEST); + + private int idleBetweenPolls; + + private int consumerBackoff; + + private long checkpointsInterval = 60_000L; + + public KclMessageDrivenChannelAdapter(String streams, Executor executor) { + this(streams, executor, KinesisAsyncClient.builder().build(), + CloudWatchAsyncClient.builder().build(), DynamoDbAsyncClient.builder().build()); + } + + public KclMessageDrivenChannelAdapter(String streams, Executor executor, Region region) { + this(streams, executor, KinesisAsyncClient.builder().region(region).build(), + CloudWatchAsyncClient.builder().region(region).build(), DynamoDbAsyncClient.builder().region(region).build()); + } + + public KclMessageDrivenChannelAdapter(String stream, Executor executor, + KinesisAsyncClient kinesisClient, CloudWatchAsyncClient cloudWatchClient, DynamoDbAsyncClient dynamoDBClient) { + Assert.notNull(stream, "'stream' must not be null."); + Assert.notNull(executor, "'executor' must not be null."); + Assert.notNull(kinesisClient, "'kinesisClient' must not be null."); + Assert.notNull(cloudWatchClient, "'cloudWatchClient' must not be null."); + Assert.notNull(dynamoDBClient, "'dynamoDBClient' must not be null."); + this.stream = stream; + this.executor = executor; + this.kinesisClient = kinesisClient; + this.cloudWatchClient = cloudWatchClient; + this.dynamoDBClient = dynamoDBClient; + } + + public void setConsumerGroup(String consumerGroup) { + Assert.hasText(consumerGroup, "'consumerGroup' must not be empty"); + this.consumerGroup = consumerGroup; + } + + /** + * Specify an {@link InboundMessageMapper} to extract message headers embedded + * into the record data. + * + * @param embeddedHeadersMapper the {@link InboundMessageMapper} to use. + * @since 2.0 + */ + public void setEmbeddedHeadersMapper(InboundMessageMapper embeddedHeadersMapper) { + this.embeddedHeadersMapper = embeddedHeadersMapper; + } + + @Override + protected void onInit() { + super.onInit(); + + String workerId = UUID.randomUUID().toString(); + RecordProcessorFactory recordProcessorFactory = new RecordProcessorFactory(); + + ConfigsBuilder configsBuilder = new ConfigsBuilder(this.stream, this.consumerGroup, + this.kinesisClient, this.dynamoDBClient, this.cloudWatchClient, workerId, recordProcessorFactory); + configsBuilder.retrievalConfig().initialPositionInStreamExtended(this.streamInitialSequence); + configsBuilder.retrievalConfig().listShardsBackoffTimeInMillis(this.consumerBackoff); + configsBuilder.coordinatorConfig().parentShardPollIntervalMillis(this.idleBetweenPolls); + + this.scheduler = new Scheduler(configsBuilder.checkpointConfig(), + configsBuilder.coordinatorConfig(), + configsBuilder.leaseManagementConfig(), + configsBuilder.lifecycleConfig(), + configsBuilder.metricsConfig(), + configsBuilder.processorConfig(), + configsBuilder.retrievalConfig()); + } + + @Override + protected void doStart() { + super.doStart(); + this.executor.execute(this.scheduler); + } + + /** + * Takes no action by default. Subclasses may override this if they need + * lifecycle-managed behavior. + */ + @Override + protected void doStop() { + + super.doStop(); + this.scheduler.shutdown(); + + } + + @Override + protected AttributeAccessor getErrorMessageAttributes(org.springframework.messaging.Message message) { + AttributeAccessor attributes = attributesHolder.get(); + if (attributes == null) { + return super.getErrorMessageAttributes(message); + } + else { + return attributes; + } + } + + public void setStreamInitialSequence(InitialPositionInStream streamInitialSequence) { + setStreamInitialSequenceExtended(InitialPositionInStreamExtended.newInitialPosition(streamInitialSequence)); + } + + public void setStreamInitialSequenceExtended(InitialPositionInStreamExtended streamInitialSequence) { + Assert.notNull(streamInitialSequence, "'streamInitialSequence' must not be null"); + this.streamInitialSequence = streamInitialSequence; + } + + public void setIdleBetweenPolls(int idleBetweenPolls) { + this.idleBetweenPolls = Math.max(250, idleBetweenPolls); + } + + public void setConsumerBackoff(int consumerBackoff) { + this.consumerBackoff = Math.max(1000, consumerBackoff); + } + + /** + * Sets the interval between 2 checkpoints. + * + * @param checkpointsInterval interval between 2 checkpoints (in milliseconds) + */ + public void setCheckpointsInterval(long checkpointsInterval) { + this.checkpointsInterval = checkpointsInterval; + } + + @Override + public String toString() { + return "KclMessageDrivenChannelAdapter{consumerGroup='" + this.consumerGroup + '\'' + ", stream='" + this.stream + "'}"; + } + + private class RecordProcessorFactory implements ShardRecordProcessorFactory { + @Override + public ShardRecordProcessor shardRecordProcessor() { + return new RecordProcessor(); + } + } + + /** + * Processes records and checkpoints progress. + */ + private class RecordProcessor implements ShardRecordProcessor { + + private String shardId; + + private long nextCheckpointTimeInMillis; + + /** {@inheritDoc} */ + @Override + public void initialize(InitializationInput initializationInput) { + this.shardId = initializationInput.shardId(); + if (logger.isInfoEnabled()) { + logger.info("Initializing record processor for shard: " + this.shardId); + } + } + + /** {@inheritDoc} */ + @Override + public void leaseLost(LeaseLostInput leaseLostInput) { + logger.info("Lost lease, so terminating."); + } + + /** {@inheritDoc} */ + @Override + public void shardEnded(ShardEndedInput shardEndedInput) { + try { + logger.info("Reached shard end checkpointing."); + shardEndedInput.checkpointer().checkpoint(); + } + catch (ShutdownException | InvalidStateException e) { + logger.error("Exception while checkpointing at shard end. Giving up", e); + } + } + + /** {@inheritDoc} */ + @Override + public void shutdownRequested(ShutdownRequestedInput shutdownRequestedInput) { + try { + logger.info("Scheduler is shutting down, checkpointing."); + shutdownRequestedInput.checkpointer().checkpoint(); + } + catch (ShutdownException | InvalidStateException e) { + logger.error("Exception while checkpointing at requested shutdown. Giving up", e); + } + } + + /** + * Process records. Skip "poison pill" records. + * + * @param records Data records to be processed. + */ + private void processRecords(List records) { + for (KinesisClientRecord record : records) { + try { + processSingleRecord(record); + } + catch (Throwable t) { + logger.warn("Caught throwable while processing record " + record, t); + } + finally { + attributesHolder.remove(); + } + } + } + + /** + * Process a single record. + * + * @param record The record to be processed. + */ + private void processSingleRecord(KinesisClientRecord record) { + + // Convert AWS Record in Spring Message. + performSend(prepareMessageForRecord(record), record); + } + + private void performSend(AbstractIntegrationMessageBuilder messageBuilder, Object rawRecord) { + Message messageToSend = messageBuilder.build(); + setAttributesIfNecessary(rawRecord, messageToSend); + try { + sendMessage(messageToSend); + } + catch (Exception e) { + logger.error("Got an exception during sending a '" + messageToSend + "'" + "\nfor the '" + rawRecord + + "'.\n" + "Consider to use 'errorChannel' flow for the compensation logic.", e); + } + } + + /** + * If there's an error channel, we create a new attributes holder here. + * Then set the attributes for use by the {@link ErrorMessageStrategy}. + * @param record the Kinesis record to use. + * @param message the Spring Messaging message to use. + */ + private void setAttributesIfNecessary(Object record, Message message) { + if (getErrorChannel() != null) { + AttributeAccessor attributes = ErrorMessageUtils.getAttributeAccessor(message, null); + attributesHolder.set(attributes); + attributes.setAttribute(AwsHeaders.RAW_RECORD, record); + } + } + + private AbstractIntegrationMessageBuilder prepareMessageForRecord(KinesisClientRecord record) { + ByteBuffer data = record.data(); + byte[] dataArray = new byte[data.remaining()]; + Object payload = dataArray; + data.get(dataArray); + Message messageToUse = null; + + if (KclMessageDrivenChannelAdapter.this.embeddedHeadersMapper != null) { + try { + messageToUse = KclMessageDrivenChannelAdapter.this.embeddedHeadersMapper.toMessage((byte[]) payload); + + payload = messageToUse.getPayload(); + } + catch (Exception e) { + logger.warn("Could not parse embedded headers. Remain payload untouched.", e); + } + } + + AbstractIntegrationMessageBuilder messageBuilder = getMessageBuilderFactory().withPayload(payload) + .setHeader(AwsHeaders.RECEIVED_PARTITION_KEY, record.partitionKey()) + .setHeader(AwsHeaders.RECEIVED_SEQUENCE_NUMBER, record.sequenceNumber()) + .setHeader(AwsHeaders.RECEIVED_STREAM, KclMessageDrivenChannelAdapter.this.stream) + .setHeader(AwsHeaders.SHARD, this.shardId); + + if (messageToUse != null) { + messageBuilder.copyHeadersIfAbsent(messageToUse.getHeaders()); + } + + return messageBuilder; + } + + /** + * Checkpoint with retries. + * + * @param checkpointer checkpointer + */ + private void checkpoint(RecordProcessorCheckpointer checkpointer) { + if (logger.isInfoEnabled()) { + logger.info("Checkpointing shard " + shardId); + } + try { + checkpointer.checkpoint(); + } + catch (ShutdownException se) { + // Ignore checkpoint if the processor instance has been shutdown (fail over). + logger.info("Caught shutdown exception, skipping checkpoint.", se); + } + catch (ThrottlingException e) { + if (logger.isInfoEnabled()) { + logger.info("Transient issue when checkpointing", e); + } + } + catch (InvalidStateException e) { + // This indicates an issue with the DynamoDB table (check for table, provisioned + // IOPS). + logger.error("Cannot save checkpoint to the DynamoDB table used by the Amazon Kinesis Client Library.", + e); + } + } + + @Override + public void processRecords(ProcessRecordsInput processRecordsInput) { + List records = processRecordsInput.records(); + if (logger.isDebugEnabled()) { + logger.debug("Processing " + records.size() + " records from " + this.shardId); + } + + // Process records and perform all exception handling. + processRecords(records); + + // Checkpoint once every checkpoint interval. + if (System.currentTimeMillis() > nextCheckpointTimeInMillis) { + checkpoint(processRecordsInput.checkpointer()); + this.nextCheckpointTimeInMillis = System.currentTimeMillis() + checkpointsInterval; + } + } + } +} diff --git a/src/main/java/org/springframework/integration/aws/outbound/KplMessageHandler.java b/src/main/java/org/springframework/integration/aws/outbound/KplMessageHandler.java new file mode 100644 index 0000000..7e6e390 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/outbound/KplMessageHandler.java @@ -0,0 +1,282 @@ +/* + * Copyright 2017-2018 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 + * + * http://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 java.nio.ByteBuffer; +import java.util.concurrent.Future; + +import org.springframework.core.convert.converter.Converter; +import org.springframework.core.serializer.support.SerializingConverter; +import org.springframework.expression.Expression; +import org.springframework.expression.common.LiteralExpression; +import org.springframework.integration.aws.support.AwsHeaders; +import org.springframework.integration.aws.support.AwsRequestFailureException; +import org.springframework.integration.handler.AbstractMessageHandler; +import org.springframework.integration.mapping.HeaderMapper; +import org.springframework.integration.mapping.OutboundMessageMapper; +import org.springframework.integration.support.AbstractIntegrationMessageBuilder; +import org.springframework.integration.support.MutableMessage; +import org.springframework.messaging.Message; +import org.springframework.messaging.MessageHeaders; +import org.springframework.util.Assert; +import org.springframework.util.StringUtils; + +import com.amazonaws.AmazonWebServiceRequest; +import com.amazonaws.handlers.AsyncHandler; +import com.amazonaws.services.kinesis.AmazonKinesisAsync; +import com.amazonaws.services.kinesis.model.PutRecordRequest; +import com.amazonaws.services.kinesis.model.PutRecordResult; +import com.amazonaws.services.kinesis.model.PutRecordsRequest; +import com.amazonaws.services.kinesis.producer.KinesisProducer; +import com.amazonaws.services.kinesis.producer.UserRecord; +import com.amazonaws.services.kinesis.producer.UserRecordResult; +import com.google.common.util.concurrent.FutureCallback; +import com.google.common.util.concurrent.Futures; +import com.google.common.util.concurrent.ListenableFuture; +import com.google.common.util.concurrent.MoreExecutors; + +/** + * The {@link AbstractMessageHandler} implementation for the Amazon Kinesis Producer Library {@code putRecord(s)}. + * + * @author Arnaud Lecollaire + * + * @since 2.1.0 + * + * @see AmazonKinesisAsync#putRecord(PutRecordRequest) + * @see AmazonKinesisAsync#putRecords(PutRecordsRequest) + * @see com.amazonaws.handlers.AsyncHandler + */ +public class KplMessageHandler extends AbstractAwsMessageHandler { + + private final KinesisProducer kinesisProducer; + + private Converter converter = new SerializingConverter(); + + private Expression streamExpression; + + private Expression partitionKeyExpression; + + private Expression explicitHashKeyExpression; + + private Expression sequenceNumberExpression; + + private OutboundMessageMapper embeddedHeadersMapper; + + public KplMessageHandler(KinesisProducer kinesisProducer) { + Assert.notNull(kinesisProducer, "'kinesisProducer' must not be null."); + this.kinesisProducer = kinesisProducer; + } + + /** + * Specify a {@link Converter} to serialize {@code payload} to the {@code byte[]} + * if that isn't {@code byte[]} already. + * @param converter the {@link Converter} to use; cannot be null. + */ + public void setConverter(Converter converter) { + Assert.notNull(converter, "'converter' must not be null."); + this.converter = converter; + } + + public void setStream(String stream) { + setStreamExpression(new LiteralExpression(stream)); + } + + public void setStreamExpressionString(String streamExpression) { + setStreamExpression(EXPRESSION_PARSER.parseExpression(streamExpression)); + } + + public void setStreamExpression(Expression streamExpression) { + this.streamExpression = streamExpression; + } + + public void setPartitionKey(String partitionKey) { + setPartitionKeyExpression(new LiteralExpression(partitionKey)); + } + + public void setPartitionKeyExpressionString(String partitionKeyExpression) { + setPartitionKeyExpression(EXPRESSION_PARSER.parseExpression(partitionKeyExpression)); + } + + public void setPartitionKeyExpression(Expression partitionKeyExpression) { + this.partitionKeyExpression = partitionKeyExpression; + } + + public void setExplicitHashKey(String explicitHashKey) { + setExplicitHashKeyExpression(new LiteralExpression(explicitHashKey)); + } + + public void setExplicitHashKeyExpressionString(String explicitHashKeyExpression) { + setExplicitHashKeyExpression(EXPRESSION_PARSER.parseExpression(explicitHashKeyExpression)); + } + + public void setExplicitHashKeyExpression(Expression explicitHashKeyExpression) { + this.explicitHashKeyExpression = explicitHashKeyExpression; + } + + public void setSequenceNumberExpressionString(String sequenceNumberExpression) { + setSequenceNumberExpression(EXPRESSION_PARSER.parseExpression(sequenceNumberExpression)); + } + + public void setSequenceNumberExpression(Expression sequenceNumberExpression) { + this.sequenceNumberExpression = sequenceNumberExpression; + } + + /** + * Specify a {@link OutboundMessageMapper} for embedding message headers into the record data + * together with payload. + * @param embeddedHeadersMapper the {@link OutboundMessageMapper} to embed headers into the record data. + * @since 2.0 + * @see org.springframework.integration.support.json.EmbeddedJsonHeadersMessageMapper + */ + public void setEmbeddedHeadersMapper(OutboundMessageMapper embeddedHeadersMapper) { + this.embeddedHeadersMapper = embeddedHeadersMapper; + } + + /** + * Unsupported operation. Use {@link #setEmbeddedHeadersMapper} instead. + * @param headerMapper is not used. + * @see #setEmbeddedHeadersMapper + */ + @Override + public void setHeaderMapper(HeaderMapper headerMapper) { + throw new UnsupportedOperationException("Kinesis doesn't support headers.\n" + + "Consider to use 'OutboundMessageMapper' for embedding headers into the record data."); + } + + @Override + protected Future handleMessageToAws(Message message) throws Exception { + if (message.getPayload() instanceof PutRecordsRequest) { + throw new UnsupportedOperationException("not implemented"); + } + else if (message.getPayload() instanceof UserRecord) { + return handleUserRecord(message, buildPutRecordRequest(message), (UserRecord) message.getPayload()); + } + else { + final PutRecordRequest putRecordRequest = + (message.getPayload() instanceof PutRecordRequest) + ? (PutRecordRequest) message.getPayload() + : buildPutRecordRequest(message); + + // convert the PutRecordRequest to a UserRecord + UserRecord userRecord = new UserRecord(); + userRecord.setExplicitHashKey(putRecordRequest.getExplicitHashKey()); + userRecord.setData(putRecordRequest.getData()); + userRecord.setPartitionKey(putRecordRequest.getPartitionKey()); + userRecord.setStreamName(putRecordRequest.getStreamName()); + return handleUserRecord(message, putRecordRequest, userRecord); + } + } + + private Future handleUserRecord(Message message, final PutRecordRequest putRecordRequest, + UserRecord userRecord) { + ListenableFuture recordResult = this.kinesisProducer.addUserRecord(userRecord); + + final AsyncHandler asyncHandler = + obtainAsyncHandler(message, putRecordRequest); + final FutureCallback callback = new FutureCallback() { + @Override + public void onFailure(Throwable t) { + asyncHandler.onError((t instanceof Exception) ? ((Exception) t) : new AwsRequestFailureException(message, putRecordRequest, t)); + } + + @Override + public void onSuccess(UserRecordResult result) { + asyncHandler.onSuccess(putRecordRequest, result); + } + }; + Futures.addCallback(recordResult, callback, MoreExecutors.directExecutor()); + + return recordResult; + } + + private PutRecordRequest buildPutRecordRequest(Message message) throws Exception { + MessageHeaders messageHeaders = message.getHeaders(); + String stream = messageHeaders.get(AwsHeaders.STREAM, String.class); + if (!StringUtils.hasText(stream) && this.streamExpression != null) { + stream = this.streamExpression.getValue(getEvaluationContext(), message, String.class); + } + Assert.state(stream != null, "'stream' must not be null for sending a Kinesis record. " + + "Consider configuring this handler with a 'stream'( or 'streamExpression') or supply an " + + "'aws_stream' message header."); + + String partitionKey = messageHeaders.get(AwsHeaders.PARTITION_KEY, String.class); + 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. " + + "Consider configuring this handler with a 'partitionKey'( or 'partitionKeyExpression') or supply an " + + "'aws_partitionKey' message header."); + + String explicitHashKey = + (this.explicitHashKeyExpression != null + ? this.explicitHashKeyExpression.getValue(getEvaluationContext(), message, String.class) + : null); + + String sequenceNumber = messageHeaders.get(AwsHeaders.SEQUENCE_NUMBER, String.class); + if (!StringUtils.hasText(sequenceNumber) && this.sequenceNumberExpression != null) { + sequenceNumber = this.sequenceNumberExpression.getValue(getEvaluationContext(), message, String.class); + } + + Object payload = message.getPayload(); + + ByteBuffer data = null; + + Message messageToEmbed = null; + + if (payload instanceof ByteBuffer) { + data = (ByteBuffer) payload; + if (this.embeddedHeadersMapper != null) { + messageToEmbed = new MutableMessage<>(data.array(), messageHeaders); + } + } + else { + byte[] bytes = + payload instanceof byte[] + ? (byte[]) payload + : this.converter.convert(payload); + if (this.embeddedHeadersMapper != null) { + messageToEmbed = new MutableMessage<>(bytes, messageHeaders); + } + else { + data = ByteBuffer.wrap(bytes); + } + } + + if (messageToEmbed != null) { + data = ByteBuffer.wrap(this.embeddedHeadersMapper.fromMessage(messageToEmbed)); + } + + return new PutRecordRequest() + .withStreamName(stream) + .withPartitionKey(partitionKey) + .withExplicitHashKey(explicitHashKey) + .withSequenceNumberForOrdering(sequenceNumber) + .withData(data); + } + + @Override + protected void additionalOnSuccessHeaders(AbstractIntegrationMessageBuilder messageBuilder, + AmazonWebServiceRequest request, Object result) { + + if (result instanceof PutRecordResult) { + messageBuilder + .setHeader(AwsHeaders.SHARD, ((PutRecordResult) result).getShardId()) + .setHeader(AwsHeaders.SEQUENCE_NUMBER, ((PutRecordResult) result).getSequenceNumber()); + } + } + +}