From ef2f5f7aef6706379601dd44e0157cc64faa38bb Mon Sep 17 00:00:00 2001 From: Artem Bilan Date: Wed, 19 Oct 2016 17:47:55 -0400 Subject: [PATCH] GH-22: Add KinesisMessageDrivenChannelAdapter Fixes GH-22 (https://github.com/spring-projects/spring-integration-aws/issues/22) Rework logic to the `dispatching` and tasks * Add `LimitExceededException` and configurable retry logic for the `describeStream` * Skip `CLOSED` shards which has been read before according stored `checkpoint`. Will be useful for `after resharding` algorithm * Add `adapt to resharding` logic * Add `KinesisMessageDrivenChannelAdapterTests` based on mocks * Add Thread Affinity for `ShardConsumer`s via `concurrency` option * Introduce `concurrency` option and `ConsumerInvoker`. `ShardConsumer`s are now distributed between `ConsumerInvoker`s if `concurrency > 0` `ConsumerInvoker`s are scheduled to the `ConsumerExecutor` as a `isLongLived` tasks The concurrency is adjusted if there are no more `ShardConsumer`s to process (`STOP` state because of closed shard). At the same time newly populated `ShardConsumer`s (e.g. after resharding) is distributed evenly between existing `ConsumerExecutor`s if `maxConcurrency` is exceeded. Otherwise fresh `ConsumerExecutor` is started for new `ShardConsumer` --- build.gradle | 2 + .../aws/inbound/kinesis/CheckpointMode.java | 44 + .../aws/inbound/kinesis/Checkpointer.java | 31 + .../KinesisMessageDrivenChannelAdapter.java | 893 ++++++++++++++++++ .../inbound/kinesis/KinesisShardOffset.java | 215 +++++ .../aws/inbound/kinesis/ListenerMode.java | 40 + .../inbound/kinesis/ShardCheckpointer.java | 130 +++ .../aws/inbound/kinesis/package-info.java | 20 + .../aws/outbound/KinesisMessageHandler.java | 1 + .../integration/aws/support/AwsHeaders.java | 10 + ...nesisMessageDrivenChannelAdapterTests.java | 281 ++++++ src/test/resources/log4j.properties | 3 +- 12 files changed, 1669 insertions(+), 1 deletion(-) create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/CheckpointMode.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/Checkpointer.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisShardOffset.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/ListenerMode.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/ShardCheckpointer.java create mode 100644 src/main/java/org/springframework/integration/aws/inbound/kinesis/package-info.java create mode 100644 src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java diff --git a/build.gradle b/build.gradle index f0ce2fd..43021b0 100644 --- a/build.gradle +++ b/build.gradle @@ -91,6 +91,8 @@ dependencies { compile("com.amazonaws:aws-java-sdk-kinesis:$awsSdkVersion", optional) + compile("com.amazonaws:amazon-kinesis-client:1.7.0", optional) + compile("javax.servlet:javax.servlet-api:$servletApiVersion", provided) testCompile "org.springframework.integration:spring-integration-test:$springIntegrationVersion" diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/CheckpointMode.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/CheckpointMode.java new file mode 100644 index 0000000..c40ebf7 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/CheckpointMode.java @@ -0,0 +1,44 @@ +/* + * Copyright 2017 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; + +/** + * The listener mode, record or batch. + * + * @author Artem Bilan + * + * @since 1.1 + */ +public enum CheckpointMode { + + /** + * Checkpoint after each processed record. + * Makes sense only if {@link ListenerMode#record} is used. + */ + record, + + /** + * Checkpoint after each processed batch of records. + */ + batch, + + /** + * Checkpoint on demand via provided to the message {@link Checkpointer} callback. + */ + manual + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/Checkpointer.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/Checkpointer.java new file mode 100644 index 0000000..2d083b3 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/Checkpointer.java @@ -0,0 +1,31 @@ +/* + * Copyright 2017 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; + +/** + * A callback for target record process to perform checkpoint on the related shard. + * + * @author Artem Bilan + * @since 1.1 + */ +public interface Checkpointer { + + void checkpoint(); + + void checkpoint(String sequenceNumber); + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java new file mode 100644 index 0000000..fbb97a3 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisMessageDrivenChannelAdapter.java @@ -0,0 +1,893 @@ +/* + * Copyright 2017 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.math.BigInteger; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.Collections; +import java.util.Date; +import java.util.HashSet; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentSkipListSet; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; + +import org.springframework.beans.factory.DisposableBean; +import org.springframework.core.convert.converter.Converter; +import org.springframework.core.serializer.support.DeserializingConverter; +import org.springframework.integration.aws.support.AwsHeaders; +import org.springframework.integration.endpoint.MessageProducerSupport; +import org.springframework.integration.metadata.MetadataStore; +import org.springframework.integration.metadata.SimpleMetadataStore; +import org.springframework.integration.support.AbstractIntegrationMessageBuilder; +import org.springframework.integration.support.management.IntegrationManagedResource; +import org.springframework.jmx.export.annotation.ManagedOperation; +import org.springframework.jmx.export.annotation.ManagedResource; +import org.springframework.scheduling.SchedulingAwareRunnable; +import org.springframework.util.Assert; +import org.springframework.util.StringUtils; + +import com.amazonaws.services.kinesis.AmazonKinesis; +import com.amazonaws.services.kinesis.clientlibrary.utils.NamedThreadFactory; +import com.amazonaws.services.kinesis.model.DescribeStreamRequest; +import com.amazonaws.services.kinesis.model.DescribeStreamResult; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.GetShardIteratorRequest; +import com.amazonaws.services.kinesis.model.LimitExceededException; +import com.amazonaws.services.kinesis.model.Record; +import com.amazonaws.services.kinesis.model.ResourceNotFoundException; +import com.amazonaws.services.kinesis.model.Shard; +import com.amazonaws.services.kinesis.model.ShardIteratorType; +import com.amazonaws.services.kinesis.model.StreamStatus; + +/** + * The {@link MessageProducerSupport} implementation for receiving data from Amazon Kinesis stream(s). + * + * @author Artem Bilan + * + * @since 1.1 + */ +@ManagedResource +@IntegrationManagedResource +public class KinesisMessageDrivenChannelAdapter extends MessageProducerSupport implements DisposableBean { + + private final AmazonKinesis amazonKinesis; + + private final String[] streams; + + private final Set shardOffsets = new HashSet<>(); + + private final Map shardConsumers = new ConcurrentHashMap<>(); + + private final Set inResharding = new ConcurrentSkipListSet<>(); + + private final List consumerInvokers = new ArrayList<>(); + + private String consumerGroup = "SpringIntegration"; + + private MetadataStore checkpointStore = new SimpleMetadataStore(); + + private Executor dispatcherExecutor; + + private boolean dispatcherExecutorExplicitlySet; + + private Executor consumerExecutor; + + private boolean consumerExecutorExplicitlySet; + + private int maxConcurrency; + + private int concurrency; + + private KinesisShardOffset streamInitialSequence = KinesisShardOffset.latest(); + + private Converter converter = new DeserializingConverter(); + + private ListenerMode listenerMode = ListenerMode.record; + + private CheckpointMode checkpointMode = CheckpointMode.batch; + + private int recordsLimit = 25; + + private int consumerBackoff = 1000; + + private int startTimeout = 60 * 1000; + + private int describeStreamBackoff = 1000; + + private int describeStreamRetries = 50; + + private boolean resetCheckpoints; + + private volatile boolean active; + + private volatile int consumerInvokerMaxCapacity; + + public KinesisMessageDrivenChannelAdapter(AmazonKinesis amazonKinesis, String... streams) { + Assert.notNull(amazonKinesis, "'amazonKinesis' must not be null."); + Assert.notEmpty(streams, "'streams' must not be null."); + this.amazonKinesis = amazonKinesis; + this.streams = Arrays.copyOf(streams, streams.length); + } + + public KinesisMessageDrivenChannelAdapter(AmazonKinesis amazonKinesis, + KinesisShardOffset... shardOffsets) { + Assert.notNull(amazonKinesis, "'amazonKinesis' must not be null."); + Assert.notEmpty(shardOffsets, "'shardOffsets' must not be null."); + Assert.noNullElements(shardOffsets, "'shardOffsets' must not contain null elements."); + for (KinesisShardOffset shardOffset : shardOffsets) { + Assert.isTrue(StringUtils.hasText(shardOffset.getStream()) && StringUtils.hasText(shardOffset.getShard()), + "The 'shardOffsets' must be provided with particular 'stream' and 'shard' values."); + this.shardOffsets.add(new KinesisShardOffset(shardOffset)); + } + this.amazonKinesis = amazonKinesis; + this.streams = null; + } + + public void setConsumerGroup(String consumerGroup) { + Assert.hasText(consumerGroup, "'consumerGroup' must not be empty"); + this.consumerGroup = consumerGroup; + } + + public void setCheckpointStore(MetadataStore checkpointStore) { + Assert.notNull(checkpointStore, "'checkpointStore' must not be null"); + this.checkpointStore = checkpointStore; + } + + public void setConsumerExecutor(Executor executor) { + Assert.notNull(executor, "'executor' must not be null"); + this.consumerExecutor = executor; + this.consumerExecutorExplicitlySet = true; + } + + public void setDispatcherExecutor(Executor dispatcherExecutor) { + this.dispatcherExecutor = dispatcherExecutor; + this.dispatcherExecutorExplicitlySet = true; + } + + public void setStreamInitialSequence(KinesisShardOffset streamInitialSequence) { + Assert.notNull(streamInitialSequence, "'streamInitialSequence' must not be null"); + this.streamInitialSequence = streamInitialSequence; + } + + public void setConverter(Converter converter) { + Assert.notNull(converter, "'converter' must not be null"); + this.converter = converter; + } + + public void setListenerMode(ListenerMode listenerMode) { + Assert.notNull(listenerMode, "'listenerMode' must not be null"); + this.listenerMode = listenerMode; + } + + public void setCheckpointMode(CheckpointMode checkpointMode) { + Assert.notNull(checkpointMode, "'checkpointMode' must not be null"); + this.checkpointMode = checkpointMode; + } + + public void setRecordsLimit(int recordsLimit) { + Assert.isTrue(recordsLimit > 0, "'recordsLimit' must be more than 0"); + this.recordsLimit = recordsLimit; + } + + public void setConsumerBackoff(int consumerBackoff) { + this.consumerBackoff = Math.max(1000, consumerBackoff); + } + + public void setDescribeStreamBackoff(int describeStreamBackoff) { + this.describeStreamBackoff = Math.max(1000, describeStreamBackoff); + } + + public void setDescribeStreamRetries(int describeStreamRetries) { + Assert.isTrue(describeStreamRetries > 0, "'describeStreamRetries' must be more than 0"); + this.describeStreamRetries = describeStreamRetries; + } + + public void setStartTimeout(int startTimeout) { + Assert.isTrue(startTimeout > 0, "'startTimeout' must be more than 0"); + this.startTimeout = startTimeout; + } + + /** + * The maximum number of concurrent {@link ConsumerInvoker}s running. + * The {@link ShardConsumer}s are evenly distributed between {@link ConsumerInvoker}s. + * Messages from within the same shard will be processed sequentially. + * In other words each shard is tied with the particular thread. + * By default the concurrency is unlimited and shard + * is processed in the {@link #consumerExecutor} directly. + * @param concurrency the concurrency maximum number + */ + public void setConcurrency(int concurrency) { + this.maxConcurrency = concurrency; + } + + @Override + protected void onInit() { + super.onInit(); + + if (this.consumerExecutor == null) { + this.consumerExecutor = Executors.newCachedThreadPool( + new NamedThreadFactory((getComponentName() == null + ? "" + : getComponentName()) + + "-kinesis-consumer-")); + } + if (this.dispatcherExecutor == null) { + this.dispatcherExecutor = + Executors.newCachedThreadPool( + new NamedThreadFactory((getComponentName() == null + ? "" + : getComponentName()) + + "-kinesis-dispatcher-")); + } + } + + @Override + public void destroy() throws Exception { + if (!this.consumerExecutorExplicitlySet) { + ((ExecutorService) this.consumerExecutor).shutdownNow(); + } + if (!this.dispatcherExecutorExplicitlySet) { + ((ExecutorService) this.dispatcherExecutor).shutdownNow(); + } + } + + @ManagedOperation + public void stopConsumer(String stream, String shard) { + ShardConsumer shardConsumer = this.shardConsumers.remove(KinesisShardOffset.latest(stream, shard)); + if (shardConsumer != null) { + shardConsumer.stop(); + } + else { + if (this.logger.isDebugEnabled()) { + this.logger.debug("There is no ShardConsumer for shard [" + shard + "] in stream [" + shard + + "] to stop."); + } + } + } + + @ManagedOperation + public void startConsumer(String stream, String shard) { + KinesisShardOffset shardOffsetForSearch = KinesisShardOffset.latest(stream, shard); + ShardConsumer shardConsumer = this.shardConsumers.get(shardOffsetForSearch); + if (shardConsumer != null) { + if (this.logger.isDebugEnabled()) { + this.logger.debug("The [" + shardConsumer + "] has been started before."); + } + } + else { + synchronized (this.shardOffsets) { + for (KinesisShardOffset shardOffset : this.shardOffsets) { + if (shardOffsetForSearch.equals(shardOffset)) { + populateConsumer(shardOffset); + break; + } + } + } + } + } + + @ManagedOperation + public void resetCheckpointForShardToLatest(String stream, String shard) { + restartShardConsumerForOffset(KinesisShardOffset.latest(stream, shard)); + } + + @ManagedOperation + public void resetCheckpointForShardToTrimHorizon(String stream, String shard) { + restartShardConsumerForOffset(KinesisShardOffset.trimHorizon(stream, shard)); + } + + @ManagedOperation + public void resetCheckpointForShardToSequenceNumber(String stream, String shard, String sequenceNumber) { + restartShardConsumerForOffset(KinesisShardOffset.atSequenceNumber(stream, shard, sequenceNumber)); + } + + @ManagedOperation + public void resetCheckpointForShardAtTimestamp(String stream, String shard, long timestamp) { + restartShardConsumerForOffset(KinesisShardOffset.atTimestamp(stream, shard, new Date(timestamp))); + } + + private void restartShardConsumerForOffset(KinesisShardOffset shardOffset) { + Assert.isTrue(this.shardOffsets.contains(shardOffset), + "The [" + this + + "] doesn't operate shard [" + shardOffset.getShard() + + "] for stream [" + shardOffset.getStream() + "]"); + + if (logger.isDebugEnabled()) { + logger.debug("Resetting consumer for [" + shardOffset + "]..."); + } + shardOffset.reset(); + synchronized (this.shardOffsets) { + this.shardOffsets.remove(shardOffset); + this.shardOffsets.add(shardOffset); + } + if (this.active) { + ShardConsumer oldShardConsumer = this.shardConsumers.remove(shardOffset); + if (oldShardConsumer != null) { + oldShardConsumer.close(); + } + shardOffset.setReset(true); + populateConsumer(shardOffset); + } + } + + @ManagedOperation + public void resetCheckpoints() { + this.resetCheckpoints = true; + if (this.active) { + stopConsumers(); + populateConsumers(); + } + } + + @Override + protected void doStart() { + super.doStart(); + if (ListenerMode.batch.equals(this.listenerMode) && CheckpointMode.record.equals(this.checkpointMode)) { + this.checkpointMode = CheckpointMode.batch; + logger.warn("The 'checkpointMode' is overridden from [CheckpointMode.record] to [CheckpointMode.batch] " + + "because it does not make sense in case of [ListenerMode.batch]."); + } + if (this.streams != null) { + populateShardsForStreams(); + } + + populateConsumers(); + + this.active = true; + + this.concurrency = Math.min(this.maxConcurrency, this.shardOffsets.size()); + + for (int i = 0; i < this.concurrency; i++) { + Collection shardConsumers = shardConsumerSubset(i); + this.consumerInvokerMaxCapacity = Math.max(this.consumerInvokerMaxCapacity, shardConsumers.size()); + ConsumerInvoker consumerInvoker = new ConsumerInvoker(shardConsumers); + this.consumerInvokers.add(consumerInvoker); + this.consumerExecutor.execute(consumerInvoker); + } + + this.dispatcherExecutor.execute(new ConsumerDispatcher()); + } + + private Collection shardConsumerSubset(int i) { + List shardConsumers = new ArrayList<>(this.shardConsumers.values()); + if (this.concurrency == 1) { + return shardConsumers; + } + else { + int numConsumers = shardConsumers.size(); + if (numConsumers == this.concurrency) { + return Collections.singleton(shardConsumers.get(i)); + } + else { + int perInvoker = numConsumers / this.concurrency; + List subset; + if (i == this.concurrency - 1) { + subset = shardConsumers.subList(i * perInvoker, numConsumers); + } + else { + subset = shardConsumers.subList(i * perInvoker, (i + 1) * perInvoker); + } + return subset; + } + } + } + + private void populateShardsForStreams() { + this.shardOffsets.clear(); + final CountDownLatch shardsGatherLatch = new CountDownLatch(this.streams.length); + for (final String stream : this.streams) { + populateShardsForStream(stream, shardsGatherLatch); + } + try { + if (!shardsGatherLatch.await(this.startTimeout, TimeUnit.MILLISECONDS)) { + throw new IllegalStateException("The [ " + + this + + "] could not start during timeout: " + this.startTimeout); + } + } + catch (InterruptedException e) { + throw new IllegalStateException("The [ " + this + "] has been interrupted from start."); + } + } + + private void populateShardsForStream(final String stream, final CountDownLatch shardsGatherLatch) { + this.dispatcherExecutor.execute(new Runnable() { + + @Override + public void run() { + try { + int describeStreamRetries = 0; + List shards = new ArrayList<>(); + + String exclusiveStartShardId = null; + while (true) { + DescribeStreamRequest describeStreamRequest = new DescribeStreamRequest() + .withStreamName(stream) + .withExclusiveStartShardId(exclusiveStartShardId); + + DescribeStreamResult describeStreamResult = null; + // Call DescribeStream, with backoff and retries (if we get LimitExceededException). + try { + describeStreamResult = KinesisMessageDrivenChannelAdapter.this.amazonKinesis + .describeStream(describeStreamRequest); + } + catch (LimitExceededException e) { + logger.info("Got LimitExceededException when describing stream [" + stream + "]. " + + "Backing off for [" + + KinesisMessageDrivenChannelAdapter.this.describeStreamBackoff + "] millis."); + } + + if (describeStreamResult == null || !StreamStatus.ACTIVE.toString() + .equals(describeStreamResult.getStreamDescription().getStreamStatus())) { + if (describeStreamRetries++ > + KinesisMessageDrivenChannelAdapter.this.describeStreamRetries) { + ResourceNotFoundException resourceNotFoundException = + new ResourceNotFoundException("The stream [" + stream + + "] isn't ACTIVE or doesn't exist."); + resourceNotFoundException.setServiceName("Kinesis"); + throw resourceNotFoundException; + } + try { + Thread.sleep(KinesisMessageDrivenChannelAdapter.this.describeStreamBackoff); + continue; + } + catch (InterruptedException e) { + Thread.interrupted(); + throw new IllegalStateException("The [describeStream] thread for the stream [" + + stream + "] has been interrupted.", e); + } + } + + for (Shard shard : describeStreamResult.getStreamDescription().getShards()) { + String endingSequenceNumber = shard.getSequenceNumberRange().getEndingSequenceNumber(); + if (endingSequenceNumber != null) { + String key = KinesisMessageDrivenChannelAdapter.this.consumerGroup + + ":" + stream + + ":" + shard.getShardId(); + String checkpoint = + KinesisMessageDrivenChannelAdapter.this.checkpointStore.get(key); + + if (checkpoint != null && + new BigInteger(endingSequenceNumber) + .compareTo(new BigInteger(checkpoint)) <= 0) { + // Skip CLOSED shard which has been read before according a checkpoint + continue; + } + } + + shards.add(shard); + } + + if (describeStreamResult.getStreamDescription().getHasMoreShards()) { + exclusiveStartShardId = shards.get(shards.size() - 1).getShardId(); + } + else { + break; + } + } + + for (Shard shard : shards) { + KinesisShardOffset shardOffset = + new KinesisShardOffset(KinesisMessageDrivenChannelAdapter.this.streamInitialSequence); + shardOffset.setShard(shard.getShardId()); + shardOffset.setStream(stream); + boolean addedOffset; + synchronized (KinesisMessageDrivenChannelAdapter.this.shardOffsets) { + addedOffset = KinesisMessageDrivenChannelAdapter.this.shardOffsets.add(shardOffset); + } + if (addedOffset + && shardsGatherLatch == null + && KinesisMessageDrivenChannelAdapter.this.active) { + populateConsumer(shardOffset); + } + } + } + finally { + if (shardsGatherLatch != null) { + shardsGatherLatch.countDown(); + } + KinesisMessageDrivenChannelAdapter.this.inResharding.remove(stream); + } + } + + }); + } + + private void populateConsumers() { + synchronized (this.shardOffsets) { + for (KinesisShardOffset shardOffset : this.shardOffsets) { + shardOffset.setReset(this.resetCheckpoints); + populateConsumer(shardOffset); + } + } + + this.resetCheckpoints = false; + } + + + private void populateConsumer(KinesisShardOffset shardOffset) { + ShardConsumer shardConsumer = new ShardConsumer(shardOffset); + this.shardConsumers.put(shardOffset, shardConsumer); + + if (this.active) { + synchronized (this.consumerInvokers) { + if (this.consumerInvokers.size() < this.maxConcurrency) { + ConsumerInvoker consumerInvoker = new ConsumerInvoker(Collections.singleton(shardConsumer)); + this.consumerInvokers.add(consumerInvoker); + this.consumerExecutor.execute(consumerInvoker); + } + else { + for (ConsumerInvoker consumerInvoker : this.consumerInvokers) { + if (consumerInvoker.consumers.size() < this.consumerInvokerMaxCapacity) { + consumerInvoker.addConsumer(shardConsumer); + return; + } + } + + ConsumerInvoker firstConsumerInvoker = this.consumerInvokers.get(0); + firstConsumerInvoker.addConsumer(shardConsumer); + this.consumerInvokerMaxCapacity = firstConsumerInvoker.consumers.size(); + } + } + } + } + + @Override + protected void doStop() { + this.active = false; + super.doStop(); + stopConsumers(); + } + + private void stopConsumers() { + for (ShardConsumer shardConsumer : this.shardConsumers.values()) { + shardConsumer.stop(); + } + this.shardConsumers.clear(); + this.consumerInvokers.clear(); + } + + @Override + public String toString() { + return "KinesisMessageDrivenChannelAdapter{" + + "shardOffsets=" + this.shardOffsets + + ", consumerGroup='" + this.consumerGroup + '\'' + + '}'; + } + + private final class ConsumerDispatcher implements SchedulingAwareRunnable { + + private final Set inReshardingProcess = new HashSet<>(); + + @Override + public void run() { + // We can't rely on the 'isRunning()' because of race condition, + // when 'running' is set after submitting this task + while (KinesisMessageDrivenChannelAdapter.this.active) { + for (String stream : KinesisMessageDrivenChannelAdapter.this.inResharding) { + // Local store to avoid several tasks for the same 'stream' + if (this.inReshardingProcess.add(stream)) { + if (logger.isDebugEnabled()) { + logger.debug("Resharding has happened for stream [" + stream + "]. Rebalancing..."); + } + populateShardsForStream(stream, null); + } + } + + for (Iterator iterator = + KinesisMessageDrivenChannelAdapter.this.shardConsumers + .values() + .iterator(); + iterator.hasNext(); ) { + ShardConsumer shardConsumer = iterator.next(); + shardConsumer.execute(); + if (ConsumerState.STOP == shardConsumer.state) { + iterator.remove(); + if (KinesisMessageDrivenChannelAdapter.this.streams != null + && shardConsumer.shardIterator == null) { + // Shard is CLOSED and we are capable for resharding + KinesisShardOffset shardOffset = shardConsumer.shardOffset; + String stream = shardOffset.getStream(); + if (KinesisMessageDrivenChannelAdapter.this.inResharding.add(stream)) { + this.inReshardingProcess.remove(stream); + synchronized (KinesisMessageDrivenChannelAdapter.this.shardOffsets) { + KinesisMessageDrivenChannelAdapter.this.shardOffsets.remove(shardOffset); + } + } + } + } + } + } + } + + @Override + public boolean isLongLived() { + return true; + } + + } + + private final class ShardConsumer { + + private final KinesisShardOffset shardOffset; + + private final ShardCheckpointer checkpointer; + + private final Runnable processTask = processTask(); + + private volatile ConsumerState state = ConsumerState.NEW; + + private volatile Runnable task; + + private volatile String shardIterator; + + private volatile long sleepUntil; + + ShardConsumer(KinesisShardOffset shardOffset) { + this.shardOffset = new KinesisShardOffset(shardOffset); + String key = KinesisMessageDrivenChannelAdapter.this.consumerGroup + + ":" + shardOffset.getStream() + + ":" + shardOffset.getShard(); + this.checkpointer = new ShardCheckpointer(KinesisMessageDrivenChannelAdapter.this.checkpointStore, key); + } + + void stop() { + this.state = ConsumerState.STOP; + } + + void close() { + stop(); + this.checkpointer.close(); + } + + void execute() { + if (this.task == null) { + switch (this.state) { + + case NEW: + this.task = new Runnable() { + + @Override + public void run() { + if (logger.isInfoEnabled()) { + logger.info("The [" + this + "] has been started."); + } + if (ShardConsumer.this.shardOffset.isReset()) { + ShardConsumer.this.checkpointer.remove(); + } + else { + String checkpoint = ShardConsumer.this.checkpointer.getCheckpoint(); + if (checkpoint != null) { + ShardConsumer.this.shardOffset.setSequenceNumber(checkpoint); + ShardConsumer.this.shardOffset + .setIteratorType(ShardIteratorType.AFTER_SEQUENCE_NUMBER); + } + } + + GetShardIteratorRequest shardIteratorRequest = + ShardConsumer.this.shardOffset.toShardIteratorRequest(); + ShardConsumer.this.shardIterator = + KinesisMessageDrivenChannelAdapter.this.amazonKinesis + .getShardIterator(shardIteratorRequest) + .getShardIterator(); + if (ConsumerState.STOP != ShardConsumer.this.state) { + ShardConsumer.this.state = ConsumerState.CONSUME; + } + ShardConsumer.this.task = null; + } + + }; + break; + + case CONSUME: + this.task = this.processTask; + break; + + case SLEEP: + if (System.currentTimeMillis() >= this.sleepUntil) { + this.state = ConsumerState.CONSUME; + } + this.task = null; + break; + + case STOP: + if (this.shardIterator == null) { + if (logger.isInfoEnabled()) { + logger.info("Stopping the [" + this + + "] because the shard has been CLOSED and exhausted."); + } + } + else { + if (logger.isInfoEnabled()) { + logger.info("Stopping the [" + this + "]."); + } + } + this.task = null; + break; + } + + if (this.task != null && KinesisMessageDrivenChannelAdapter.this.concurrency == 0) { + KinesisMessageDrivenChannelAdapter.this.consumerExecutor.execute(this.task); + } + } + } + + private Runnable processTask() { + return new Runnable() { + + @Override + public void run() { + GetRecordsRequest getRecordsRequest = new GetRecordsRequest(); + getRecordsRequest.setShardIterator(ShardConsumer.this.shardIterator); + getRecordsRequest.setLimit(KinesisMessageDrivenChannelAdapter.this.recordsLimit); + GetRecordsResult result = KinesisMessageDrivenChannelAdapter.this.amazonKinesis + .getRecords(getRecordsRequest); + + List records = result.getRecords(); + + if (!records.isEmpty()) { + processRecords(records); + } + + ShardConsumer.this.shardIterator = result.getNextShardIterator(); + + if (ShardConsumer.this.shardIterator == null) { + // Shard is closed: nothing to consume any more. + // Resharding is possible. + ShardConsumer.this.state = ConsumerState.STOP; + } + + if (ConsumerState.STOP != ShardConsumer.this.state && records.isEmpty()) { + if (logger.isDebugEnabled()) { + logger.debug("No records for [" + ShardConsumer.this + + "] on sequenceNumber [" + + ShardConsumer.this.checkpointer.getLastCheckpointValue() + + "]. Suspend consuming for [" + + KinesisMessageDrivenChannelAdapter.this.consumerBackoff + "] milliseconds."); + } + ShardConsumer.this.sleepUntil = System.currentTimeMillis() + + KinesisMessageDrivenChannelAdapter.this.consumerBackoff; + ShardConsumer.this.state = ConsumerState.SLEEP; + } + + ShardConsumer.this.task = null; + } + + }; + } + + private void processRecords(List records) { + records = this.checkpointer.filterRecords(records); + if (!records.isEmpty()) { + if (logger.isDebugEnabled()) { + logger.debug("Processing records: " + records + " for [" + this + "]"); + } + switch (KinesisMessageDrivenChannelAdapter.this.listenerMode) { + case record: + for (Record record : records) { + Object payload = + KinesisMessageDrivenChannelAdapter.this.converter.convert(record.getData().array()); + AbstractIntegrationMessageBuilder messageBuilder = getMessageBuilderFactory() + .withPayload(payload) + .setHeader(AwsHeaders.STREAM, this.shardOffset.getStream()) + .setHeader(AwsHeaders.SHARD, this.shardOffset.getShard()) + .setHeader(AwsHeaders.PARTITION_KEY, record.getPartitionKey()) + .setHeader(AwsHeaders.SEQUENCE_NUMBER, record.getSequenceNumber()); + if (CheckpointMode.manual.equals(KinesisMessageDrivenChannelAdapter.this.checkpointMode)) { + messageBuilder.setHeader(AwsHeaders.CHECKPOINTER, this.checkpointer); + } + + sendMessage(messageBuilder.build()); + + if (CheckpointMode.record.equals(KinesisMessageDrivenChannelAdapter.this.checkpointMode)) { + this.checkpointer.checkpoint(record.getSequenceNumber()); + } + } + break; + + case batch: + AbstractIntegrationMessageBuilder messageBuilder = getMessageBuilderFactory() + .withPayload(records) + .setHeader(AwsHeaders.STREAM, this.shardOffset.getStream()) + .setHeader(AwsHeaders.SHARD, this.shardOffset.getShard()); + if (CheckpointMode.manual.equals(KinesisMessageDrivenChannelAdapter.this.checkpointMode)) { + messageBuilder.setHeader(AwsHeaders.CHECKPOINTER, this.checkpointer); + } + + sendMessage(messageBuilder.build()); + + break; + + } + + if (CheckpointMode.batch.equals(KinesisMessageDrivenChannelAdapter.this.checkpointMode)) { + this.checkpointer.checkpoint(); + } + } + } + + @Override + public String toString() { + return "ShardConsumer{" + + "shardOffset=" + this.shardOffset + + ", state=" + this.state + + '}'; + } + + } + + private enum ConsumerState { + + NEW, CONSUME, SLEEP, STOP + + } + + private final class ConsumerInvoker implements SchedulingAwareRunnable { + + private final List consumers = new ArrayList<>(); + + ConsumerInvoker(Collection shardConsumers) { + this.consumers.addAll(shardConsumers); + } + + void addConsumer(ShardConsumer shardConsumer) { + this.consumers.add(shardConsumer); + } + + @Override + public void run() { + while (KinesisMessageDrivenChannelAdapter.this.active && !this.consumers.isEmpty()) { + for (Iterator iterator = this.consumers.iterator(); iterator.hasNext(); ) { + ShardConsumer shardConsumer = iterator.next(); + if (ConsumerState.STOP == shardConsumer.state) { + iterator.remove(); + } + else { + if (shardConsumer.task != null) { + shardConsumer.task.run(); + } + } + } + synchronized (KinesisMessageDrivenChannelAdapter.this.consumerInvokers) { + // The attempt to survive if ShardConsumer has been added during synchronization + if (this.consumers.isEmpty()) { + KinesisMessageDrivenChannelAdapter.this.consumerInvokers.remove(this); + break; + } + } + } + } + + @Override + public boolean isLongLived() { + return true; + } + + } + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisShardOffset.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisShardOffset.java new file mode 100644 index 0000000..a4337dd --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/KinesisShardOffset.java @@ -0,0 +1,215 @@ +/* + * Copyright 2017 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.util.Date; +import java.util.Objects; + +import org.springframework.util.Assert; + +import com.amazonaws.services.kinesis.model.GetShardIteratorRequest; +import com.amazonaws.services.kinesis.model.ShardIteratorType; + +/** + * A model to represent a sequence in the shard for particular {@link ShardIteratorType}. + * + * @author Artem Bilan + * @since 1.1 + */ +public class KinesisShardOffset { + + private ShardIteratorType iteratorType; + + private String sequenceNumber; + + private Date timestamp; + + private String stream; + + private String shard; + + private boolean reset; + + public KinesisShardOffset(ShardIteratorType iteratorType) { + Assert.notNull(iteratorType, "'iteratorType' must not be null."); + this.iteratorType = iteratorType; + } + + public KinesisShardOffset(KinesisShardOffset other) { + this.iteratorType = other.getIteratorType(); + this.stream = other.getStream(); + this.shard = other.getShard(); + this.sequenceNumber = other.getSequenceNumber(); + this.timestamp = other.getTimestamp(); + this.reset = other.isReset(); + } + + public void setIteratorType(ShardIteratorType iteratorType) { + this.iteratorType = iteratorType; + } + + public ShardIteratorType getIteratorType() { + return this.iteratorType; + } + + public void setSequenceNumber(String sequenceNumber) { + this.sequenceNumber = sequenceNumber; + } + + public void setTimestamp(Date timestamp) { + this.timestamp = timestamp; + } + + public void setStream(String stream) { + this.stream = stream; + } + + public void setShard(String shard) { + this.shard = shard; + } + + public void setReset(boolean reset) { + this.reset = reset; + } + + public String getSequenceNumber() { + return this.sequenceNumber; + } + + public Date getTimestamp() { + return this.timestamp; + } + + public String getStream() { + return this.stream; + } + + public String getShard() { + return this.shard; + } + + public boolean isReset() { + return this.reset; + } + + public KinesisShardOffset reset() { + this.reset = true; + return this; + } + + public GetShardIteratorRequest toShardIteratorRequest() { + Assert.state(this.stream != null && this.shard != null, + "'stream' and 'shard' must not be null for conversion to the GetShardIteratorRequest."); + return new GetShardIteratorRequest() + .withStreamName(this.stream) + .withShardId(this.shard) + .withShardIteratorType(this.iteratorType) + .withStartingSequenceNumber(this.sequenceNumber) + .withTimestamp(this.timestamp); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + KinesisShardOffset that = (KinesisShardOffset) o; + return Objects.equals(this.stream, that.stream) && + Objects.equals(this.shard, that.shard); + } + + @Override + public int hashCode() { + return Objects.hash(this.stream, this.shard); + } + + + @Override + public String toString() { + return "KinesisShardOffset{" + + "iteratorType=" + this.iteratorType + + ", sequenceNumber='" + this.sequenceNumber + '\'' + + ", timestamp=" + this.timestamp + + ", stream='" + this.stream + '\'' + + ", shard='" + this.shard + '\'' + + ", reset=" + this.reset + + '}'; + } + + public static KinesisShardOffset latest() { + return latest(null, null); + } + + public static KinesisShardOffset latest(String stream, String shard) { + KinesisShardOffset kinesisShardOffset = new KinesisShardOffset(ShardIteratorType.LATEST); + kinesisShardOffset.stream = stream; + kinesisShardOffset.shard = shard; + return kinesisShardOffset; + } + + public static KinesisShardOffset trimHorizon() { + return trimHorizon(null, null); + } + + public static KinesisShardOffset trimHorizon(String stream, String shard) { + KinesisShardOffset kinesisShardOffset = new KinesisShardOffset(ShardIteratorType.TRIM_HORIZON); + kinesisShardOffset.stream = stream; + kinesisShardOffset.shard = shard; + return kinesisShardOffset; + } + + public static KinesisShardOffset atSequenceNumber(String sequenceNumber) { + return atSequenceNumber(null, null, sequenceNumber); + } + + public static KinesisShardOffset atSequenceNumber(String stream, String shard, String sequenceNumber) { + KinesisShardOffset kinesisShardOffset = new KinesisShardOffset(ShardIteratorType.AT_SEQUENCE_NUMBER); + kinesisShardOffset.stream = stream; + kinesisShardOffset.shard = shard; + kinesisShardOffset.sequenceNumber = sequenceNumber; + return kinesisShardOffset; + } + + public static KinesisShardOffset afterSequenceNumber(String sequenceNumber) { + return afterSequenceNumber(null, null, sequenceNumber); + } + + public static KinesisShardOffset afterSequenceNumber(String stream, String shard, String sequenceNumber) { + KinesisShardOffset kinesisShardOffset = new KinesisShardOffset(ShardIteratorType.AFTER_SEQUENCE_NUMBER); + kinesisShardOffset.stream = stream; + kinesisShardOffset.shard = shard; + kinesisShardOffset.sequenceNumber = sequenceNumber; + return kinesisShardOffset; + } + + public static KinesisShardOffset atTimestamp(Date timestamp) { + return atTimestamp(null, null, timestamp); + } + + public static KinesisShardOffset atTimestamp(String stream, String shard, Date timestamp) { + KinesisShardOffset kinesisShardOffset = new KinesisShardOffset(ShardIteratorType.AT_TIMESTAMP); + kinesisShardOffset.stream = stream; + kinesisShardOffset.shard = shard; + kinesisShardOffset.timestamp = timestamp; + return kinesisShardOffset; + } + + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/ListenerMode.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/ListenerMode.java new file mode 100644 index 0000000..ebc20b6 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/ListenerMode.java @@ -0,0 +1,40 @@ +/* + * Copyright 2017 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 org.springframework.messaging.Message; + +/** + * The listener mode, record or batch. + * + * @author Artem Bilan + * + * @since 1.1 + */ +public enum ListenerMode { + + /** + * Each {@link Message} will be converted from a single {@code Record}. + */ + record, + + /** + * Each {@link Message} will contains {@code List} if not empty. + */ + batch + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/ShardCheckpointer.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/ShardCheckpointer.java new file mode 100644 index 0000000..17ced33 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/ShardCheckpointer.java @@ -0,0 +1,130 @@ +/* + * Copyright 2017 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.math.BigInteger; +import java.util.Iterator; +import java.util.LinkedList; +import java.util.List; + +import org.apache.commons.logging.Log; +import org.apache.commons.logging.LogFactory; + +import org.springframework.integration.metadata.MetadataStore; + +import com.amazonaws.services.kinesis.model.Record; + +/** + * An internal {@link Checkpointer} implementation based on + * provided {@link MetadataStore} and unikey {@code key} for shard. + *

+ * The instances of this class is created by the {@link KinesisMessageDrivenChannelAdapter} + * for each {@code ShardConsumer}. + * + * @author Artem Bilan + * + * @since 1.1 + */ +class ShardCheckpointer implements Checkpointer { + + private static final Log logger = LogFactory.getLog(ShardCheckpointer.class); + + private final MetadataStore checkpointStore; + + private final String key; + + private volatile String lastCheckpointValue; + + private volatile boolean active = true; + + ShardCheckpointer(MetadataStore checkpointStore, String key) { + this.checkpointStore = checkpointStore; + this.key = key; + } + + @Override + public void checkpoint() { + checkpoint(this.lastCheckpointValue); + } + + @Override + public void checkpoint(String sequenceNumber) { + if (this.active) { + String existingSequence = this.checkpointStore.get(this.key); + if (existingSequence == null || + new BigInteger(existingSequence).compareTo(new BigInteger(sequenceNumber)) <= 0) { + this.checkpointStore.put(this.key, sequenceNumber); + } + } + else { + if (logger.isInfoEnabled()) { + logger.info("The [" + this + "] has been closed. Checkpoints aren't accepted anymore."); + } + } + } + + List filterRecords(List records) { + List recordsToProcess = new LinkedList<>(records); + this.lastCheckpointValue = this.checkpointStore.get(this.key); + if (this.lastCheckpointValue != null) { + for (Iterator iterator = recordsToProcess.iterator(); iterator.hasNext(); ) { + Record record = iterator.next(); + String sequenceNumber = record.getSequenceNumber(); + if (new BigInteger(sequenceNumber).compareTo(new BigInteger(this.lastCheckpointValue)) <= 0) { + if (logger.isDebugEnabled()) { + logger.debug("Removing record with sequenceNumber " + sequenceNumber + + " because the sequenceNumber is <= checkpoint(" + this.lastCheckpointValue + ")"); + } + iterator.remove(); + } + else { + this.lastCheckpointValue = sequenceNumber; + } + } + + } + else { + this.lastCheckpointValue = recordsToProcess.get(recordsToProcess.size() - 1).getSequenceNumber(); + } + return recordsToProcess; + } + + String getCheckpoint() { + return this.checkpointStore.get(this.key); + } + + String getLastCheckpointValue() { + return this.lastCheckpointValue; + } + + void remove() { + this.checkpointStore.remove(this.key); + } + + void close() { + this.active = false; + } + + @Override + public String toString() { + return "ShardCheckpointer{" + + "key='" + this.key + '\'' + + ", lastCheckpointValue='" + this.lastCheckpointValue + '\'' + + '}'; + } + +} diff --git a/src/main/java/org/springframework/integration/aws/inbound/kinesis/package-info.java b/src/main/java/org/springframework/integration/aws/inbound/kinesis/package-info.java new file mode 100644 index 0000000..9243d06 --- /dev/null +++ b/src/main/java/org/springframework/integration/aws/inbound/kinesis/package-info.java @@ -0,0 +1,20 @@ +/* + * Copyright 2016 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. + */ + +/** + * Provides components for AWS Kinesis. + */ +package org.springframework.integration.aws.inbound.kinesis; diff --git a/src/main/java/org/springframework/integration/aws/outbound/KinesisMessageHandler.java b/src/main/java/org/springframework/integration/aws/outbound/KinesisMessageHandler.java index 4b70b79..6255147 100644 --- a/src/main/java/org/springframework/integration/aws/outbound/KinesisMessageHandler.java +++ b/src/main/java/org/springframework/integration/aws/outbound/KinesisMessageHandler.java @@ -77,6 +77,7 @@ public class KinesisMessageHandler extends AbstractMessageHandler { private Expression sendTimeoutExpression = new ValueExpression<>(DEFAULT_SEND_TIMEOUT); public KinesisMessageHandler(AmazonKinesisAsync amazonKinesis) { + Assert.notNull(amazonKinesis, "'amazonKinesis' must not be null."); this.amazonKinesis = amazonKinesis; } diff --git a/src/main/java/org/springframework/integration/aws/support/AwsHeaders.java b/src/main/java/org/springframework/integration/aws/support/AwsHeaders.java index 61a90d8..6409b4e 100644 --- a/src/main/java/org/springframework/integration/aws/support/AwsHeaders.java +++ b/src/main/java/org/springframework/integration/aws/support/AwsHeaders.java @@ -70,6 +70,11 @@ public abstract class AwsHeaders { */ public static final String STREAM = PREFIX + "stream"; + /** + * The {@value SHARD} header to represent Kinesis shardId. + */ + public static final String SHARD = PREFIX + "shard"; + /** * The {@value PARTITION_KEY} header for sending/receiving data over Kinesis. */ @@ -80,4 +85,9 @@ public abstract class AwsHeaders { */ public static final String SEQUENCE_NUMBER = PREFIX + "sequenceNumber"; + /** + * The {@value CHECKPOINTER} header for checkpoint the shard sequenceNumber. + */ + public static final String CHECKPOINTER = PREFIX + "checkpointer"; + } diff --git a/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java b/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java new file mode 100644 index 0000000..c6e7801 --- /dev/null +++ b/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java @@ -0,0 +1,281 @@ +/* + * Copyright 2017 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; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.BDDMockito.given; +import static org.mockito.Mockito.mock; + +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import org.junit.Test; +import org.junit.runner.RunWith; + +import org.springframework.beans.DirectFieldAccessor; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.core.serializer.support.DeserializingConverter; +import org.springframework.core.serializer.support.SerializingConverter; +import org.springframework.integration.aws.inbound.kinesis.CheckpointMode; +import org.springframework.integration.aws.inbound.kinesis.Checkpointer; +import org.springframework.integration.aws.inbound.kinesis.KinesisMessageDrivenChannelAdapter; +import org.springframework.integration.aws.inbound.kinesis.KinesisShardOffset; +import org.springframework.integration.aws.inbound.kinesis.ListenerMode; +import org.springframework.integration.aws.support.AwsHeaders; +import org.springframework.integration.channel.QueueChannel; +import org.springframework.integration.config.EnableIntegration; +import org.springframework.integration.metadata.MetadataStore; +import org.springframework.integration.metadata.SimpleMetadataStore; +import org.springframework.integration.test.util.TestUtils; +import org.springframework.messaging.Message; +import org.springframework.messaging.MessageHeaders; +import org.springframework.messaging.PollableChannel; +import org.springframework.test.annotation.DirtiesContext; +import org.springframework.test.context.junit4.SpringRunner; + +import com.amazonaws.services.kinesis.AmazonKinesis; +import com.amazonaws.services.kinesis.model.DescribeStreamRequest; +import com.amazonaws.services.kinesis.model.DescribeStreamResult; +import com.amazonaws.services.kinesis.model.GetRecordsRequest; +import com.amazonaws.services.kinesis.model.GetRecordsResult; +import com.amazonaws.services.kinesis.model.GetShardIteratorResult; +import com.amazonaws.services.kinesis.model.Record; +import com.amazonaws.services.kinesis.model.SequenceNumberRange; +import com.amazonaws.services.kinesis.model.Shard; +import com.amazonaws.services.kinesis.model.StreamDescription; +import com.amazonaws.services.kinesis.model.StreamStatus; + +/** + * @author Artem Bilan + * @since 1.1 + */ +@RunWith(SpringRunner.class) +@DirtiesContext +public class KinesisMessageDrivenChannelAdapterTests { + + private static final String STREAM1 = "stream1"; + + @Autowired + private PollableChannel kinesisChannel; + + @Autowired + private KinesisMessageDrivenChannelAdapter kinesisMessageDrivenChannelAdapter; + + @Autowired + private MetadataStore checkpointStore; + + @Test + @SuppressWarnings({ "unchecked", "rawtypes" }) + public void testKinesisMessageDrivenChannelAdapter() throws InterruptedException { + final Set shardOffsets = + TestUtils.getPropertyValue(this.kinesisMessageDrivenChannelAdapter, "shardOffsets", Set.class); + + + KinesisShardOffset testOffset1 = KinesisShardOffset.latest(STREAM1, "1"); + KinesisShardOffset testOffset2 = KinesisShardOffset.latest(STREAM1, "2"); + synchronized (shardOffsets) { + assertThat(shardOffsets).contains(testOffset1, testOffset2); + assertThat(shardOffsets).doesNotContain(KinesisShardOffset.latest(STREAM1, "3")); + } + + Map shardConsumers = + TestUtils.getPropertyValue(this.kinesisMessageDrivenChannelAdapter, "shardConsumers", Map.class); + + assertThat(shardConsumers).containsKeys(testOffset1, testOffset2); + + Message message = this.kinesisChannel.receive(10000); + assertThat(message).isNotNull(); + assertThat(message.getPayload()).isEqualTo("foo"); + MessageHeaders headers = message.getHeaders(); + assertThat(headers.get(AwsHeaders.PARTITION_KEY)).isEqualTo("partition1"); + assertThat(headers.get(AwsHeaders.SHARD)).isEqualTo("1"); + assertThat(headers.get(AwsHeaders.SEQUENCE_NUMBER)).isEqualTo("1"); + assertThat(headers.get(AwsHeaders.STREAM)).isEqualTo(STREAM1); + Checkpointer checkpointer = headers.get(AwsHeaders.CHECKPOINTER, Checkpointer.class); + assertThat(checkpointer).isNotNull(); + + checkpointer.checkpoint(); + + message = this.kinesisChannel.receive(10000); + assertThat(message).isNotNull(); + assertThat(message.getPayload()).isEqualTo("bar"); + headers = message.getHeaders(); + assertThat(headers.get(AwsHeaders.PARTITION_KEY)).isEqualTo("partition1"); + assertThat(headers.get(AwsHeaders.SHARD)).isEqualTo("1"); + assertThat(headers.get(AwsHeaders.SEQUENCE_NUMBER)).isEqualTo("2"); + assertThat(headers.get(AwsHeaders.STREAM)).isEqualTo(STREAM1); + + assertThat(this.kinesisChannel.receive(10)).isNull(); + + assertThat(this.checkpointStore.get("SpringIntegration" + ":" + STREAM1 + ":" + "1")).isEqualTo("2"); + + this.kinesisMessageDrivenChannelAdapter.stop(); + this.kinesisMessageDrivenChannelAdapter.setListenerMode(ListenerMode.batch); + this.kinesisMessageDrivenChannelAdapter + .setCheckpointMode(CheckpointMode.record); + this.checkpointStore.put("SpringIntegration" + ":" + STREAM1 + ":" + "1", "1"); + this.kinesisMessageDrivenChannelAdapter.start(); + + message = this.kinesisChannel.receive(10000); + assertThat(message).isNotNull(); + assertThat(message.getPayload()).isInstanceOf(List.class); + List payload = (List) message.getPayload(); + assertThat(payload).size().isEqualTo(1); + Record record = payload.get(0); + assertThat(record.getPartitionKey()).isEqualTo("partition1"); + assertThat(record.getSequenceNumber()).isEqualTo("2"); + + DeserializingConverter deserializingConverter = new DeserializingConverter(); + assertThat(deserializingConverter.convert(record.getData().array())).isEqualTo("bar"); + + int n = 0; + + while (n++ < 100) { + if (!this.checkpointStore.get("SpringIntegration" + ":" + STREAM1 + ":" + "1").equals("2")) { + Thread.sleep(100); + } + else { + break; + } + } + + assertThat(n).isLessThan(100); + assertThat(this.checkpointStore.get("SpringIntegration" + ":" + STREAM1 + ":" + "1")).isEqualTo("2"); + + List consumerInvoker = + TestUtils.getPropertyValue(this.kinesisMessageDrivenChannelAdapter, "consumerInvokers", List.class); + assertThat(consumerInvoker.size()).isEqualTo(2); + } + + + @Configuration + @EnableIntegration + public static class Config { + + @Bean + public AmazonKinesis amazonKinesis() { + AmazonKinesis amazonKinesis = mock(AmazonKinesis.class); + + given(amazonKinesis.describeStream(new DescribeStreamRequest().withStreamName(STREAM1))) + .willReturn(new DescribeStreamResult() + .withStreamDescription(new StreamDescription() + .withStreamName(STREAM1) + .withStreamStatus(StreamStatus.UPDATING)), + new DescribeStreamResult() + .withStreamDescription(new StreamDescription() + .withStreamName(STREAM1) + .withStreamStatus(StreamStatus.ACTIVE) + .withHasMoreShards(false) + .withShards(new Shard() + .withShardId("1") + .withSequenceNumberRange(new SequenceNumberRange()), + new Shard() + .withShardId("2") + .withSequenceNumberRange(new SequenceNumberRange()), + new Shard() + .withShardId("3") + .withSequenceNumberRange(new SequenceNumberRange() + .withEndingSequenceNumber("1"))))); + + String shardIterator1 = "shardIterator1"; + + given(amazonKinesis.getShardIterator(KinesisShardOffset.latest(STREAM1, "1").toShardIteratorRequest())) + .willReturn(new GetShardIteratorResult() + .withShardIterator(shardIterator1)); + + String shardIterator2 = "shardIterator2"; + + given(amazonKinesis.getShardIterator(KinesisShardOffset.latest(STREAM1, "2").toShardIteratorRequest())) + .willReturn(new GetShardIteratorResult() + .withShardIterator(shardIterator2)); + + SerializingConverter serializingConverter = new SerializingConverter(); + + String shardIterator11 = "shardIterator11"; + + given(amazonKinesis.getRecords(new GetRecordsRequest() + .withShardIterator(shardIterator1) + .withLimit(25))) + .willReturn(new GetRecordsResult() + .withNextShardIterator(shardIterator11) + .withRecords(new Record() + .withPartitionKey("partition1") + .withSequenceNumber("1") + .withData(ByteBuffer.wrap(serializingConverter.convert("foo"))), + new Record() + .withPartitionKey("partition1") + .withSequenceNumber("2") + .withData(ByteBuffer.wrap(serializingConverter.convert("bar"))))); + + given(amazonKinesis.getRecords(new GetRecordsRequest() + .withShardIterator(shardIterator2) + .withLimit(25))) + .willReturn(new GetRecordsResult() + .withNextShardIterator(shardIterator2)); + + given(amazonKinesis.getRecords(new GetRecordsRequest() + .withShardIterator(shardIterator11) + .withLimit(25))) + .willReturn(new GetRecordsResult() + .withNextShardIterator(shardIterator11)); + + given(amazonKinesis.getShardIterator(KinesisShardOffset.afterSequenceNumber(STREAM1, "1", "1") + .toShardIteratorRequest())) + .willReturn(new GetShardIteratorResult() + .withShardIterator(shardIterator1)); + + return amazonKinesis; + } + + @Bean + public MetadataStore checkpointStore() { + SimpleMetadataStore simpleMetadataStore = new SimpleMetadataStore(); + String testKey = "SpringIntegration" + ":" + STREAM1 + ":" + "3"; + simpleMetadataStore.put(testKey, "1"); + return simpleMetadataStore; + } + + @Bean + public KinesisMessageDrivenChannelAdapter kinesisMessageDrivenChannelAdapter() { + KinesisMessageDrivenChannelAdapter adapter = + new KinesisMessageDrivenChannelAdapter(amazonKinesis(), STREAM1); + adapter.setOutputChannel(kinesisChannel()); + adapter.setCheckpointStore(checkpointStore()); + adapter.setCheckpointMode(CheckpointMode.manual); + adapter.setStartTimeout(10000); + adapter.setDescribeStreamRetries(1); + adapter.setConcurrency(10); + + DirectFieldAccessor dfa = new DirectFieldAccessor(adapter); + dfa.setPropertyValue("describeStreamBackoff", 10); + + return adapter; + } + + @Bean + public PollableChannel kinesisChannel() { + return new QueueChannel(); + } + + + } + +} diff --git a/src/test/resources/log4j.properties b/src/test/resources/log4j.properties index bf146e4..d1ba7b6 100644 --- a/src/test/resources/log4j.properties +++ b/src/test/resources/log4j.properties @@ -2,7 +2,8 @@ log4j.rootCategory=WARN, stdout log4j.appender.stdout=org.apache.log4j.ConsoleAppender log4j.appender.stdout.layout=org.apache.log4j.PatternLayout -log4j.appender.stdout.layout.ConversionPattern=%d %p [%c] - <%m>%n +log4j.appender.stdout.layout.ConversionPattern=%d %5p %c{1} [%t] : %m%n log4j.category.org.springframework=WARN log4j.category.org.springframework.integration.aws=WARN +log4j.category.org.springframework.integration=WARN