diff --git a/README.md b/README.md
index 3de3bab..fc11f72 100644
--- a/README.md
+++ b/README.md
@@ -617,6 +617,15 @@ this.amazonKinesis = AmazonKinesisAsyncClientBuilder.standard()
Where you should specify the port on which you have ran the Kinesalite service.
Also you can use for you testing purpose a copy of `org.springframework.integration.aws.KinesisLocalRunning` in the `/test` directory of this project.
+## Lock Registry for Amazon DynamoDB
+
+Starting with _version 2.0_, the `DynamoDbLockRegistry` implementation is available.
+Certain components (for example aggregator and resequencer) use a lock obtained from a `LockRegistry` instance to ensure that only one thread is manipulating a group at a time.
+The `DefaultLockRegistry` performs this function within a single component; you can now configure an external lock registry on these components.
+When used with a shared `MessageGroupStore`, the `DynamoDbLockRegistry` can be use to provide this functionality across multiple application instances, such that only one instance can manipulate the group at a time.
+This implementation can also be used for the distributed leader elections using a [LockRegistryLeaderInitiator][].
+The `com.amazonaws:dynamodb-lock-client` dependency must be present to make a `DynamoDbLockRegistry` working.
+
[Spring Cloud AWS]: https://github.com/spring-cloud/spring-cloud-aws
[AWS SDK for Java]: http://aws.amazon.com/sdkforjava/
[Amazon Web Services]: http://aws.amazon.com/
@@ -632,3 +641,4 @@ Also you can use for you testing purpose a copy of `org.springframework.integrat
[Kinesalite]: https://github.com/mhart/kinesalite
[Amazon SQS Message Attributes]: https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-attributes.html
[Amazon SNS Message Attributes]: https://docs.aws.amazon.com/sns/latest/dg/SNSMessageAttributes.html
+[LockRegistryLeaderInitiator]: https://docs.spring.io/spring-integration/docs/current/reference/html/messaging-endpoints-chapter.html#leadership-event-handling
diff --git a/build.gradle b/build.gradle
index f940376..02a447d 100644
--- a/build.gradle
+++ b/build.gradle
@@ -32,10 +32,11 @@ repositories {
ext {
assertjVersion = '3.9.1'
+ dynamodbLockClientVersion = '1.0.0'
servletApiVersion = '3.1.0'
log4jVersion = '2.11.0'
- springCloudAwsVersion = '2.0.0.M4'
- springIntegrationVersion = '5.0.4.RELEASE'
+ springCloudAwsVersion = '2.0.0.RC2'
+ springIntegrationVersion = '5.0.6.BUILD-SNAPSHOT'
idPrefix = 'aws'
@@ -106,6 +107,7 @@ dependencies {
compile('com.amazonaws:aws-java-sdk-kinesis', optional)
compile('com.amazonaws:aws-java-sdk-dynamodb', optional)
+ compile("com.amazonaws:dynamodb-lock-client:$dynamodbLockClientVersion", optional)
compile("javax.servlet:javax.servlet-api:$servletApiVersion", provided)
diff --git a/gradle/wrapper/gradle-wrapper.jar b/gradle/wrapper/gradle-wrapper.jar
index f6b961f..91ca28c 100644
Binary files a/gradle/wrapper/gradle-wrapper.jar and b/gradle/wrapper/gradle-wrapper.jar differ
diff --git a/gradle/wrapper/gradle-wrapper.properties b/gradle/wrapper/gradle-wrapper.properties
index bf3de21..16d2805 100644
--- a/gradle/wrapper/gradle-wrapper.properties
+++ b/gradle/wrapper/gradle-wrapper.properties
@@ -1,5 +1,5 @@
distributionBase=GRADLE_USER_HOME
distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-4.6-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-4.7-bin.zip
zipStoreBase=GRADLE_USER_HOME
zipStorePath=wrapper/dists
diff --git a/src/main/java/org/springframework/integration/aws/lock/DynamoDbLockRegistry.java b/src/main/java/org/springframework/integration/aws/lock/DynamoDbLockRegistry.java
new file mode 100644
index 0000000..009ebde
--- /dev/null
+++ b/src/main/java/org/springframework/integration/aws/lock/DynamoDbLockRegistry.java
@@ -0,0 +1,538 @@
+/*
+ * Copyright 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.lock;
+
+import java.text.SimpleDateFormat;
+import java.util.Date;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+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 java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.beans.DirectFieldAccessor;
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.dao.CannotAcquireLockException;
+import org.springframework.dao.DataAccessResourceFailureException;
+import org.springframework.integration.support.locks.ExpirableLockRegistry;
+import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
+import org.springframework.util.Assert;
+import org.springframework.util.ReflectionUtils;
+
+import com.amazonaws.services.dynamodbv2.AcquireLockOptions;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClient;
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBLockClientOptions;
+import com.amazonaws.services.dynamodbv2.CreateDynamoDBTableOptions;
+import com.amazonaws.services.dynamodbv2.LockItem;
+import com.amazonaws.services.dynamodbv2.model.LockTableDoesNotExistException;
+import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;
+
+/**
+ * An {@link ExpirableLockRegistry} implementation for the AWS DynamoDB.
+ * The algorithm is based on the {@link AmazonDynamoDBLockClient}.
+ *
+ * Can create table in DynamoDB if an external {@link AmazonDynamoDBLockClient} is not provided.
+ *
+ * @author Artem Bilan
+ *
+ * @since 2.0
+ */
+public class DynamoDbLockRegistry implements ExpirableLockRegistry, InitializingBean, DisposableBean {
+
+ /**
+ * The {@value DEFAULT_TABLE_NAME} default name for the locks table in the DynamoDB.
+ */
+ public static final String DEFAULT_TABLE_NAME = "SpringIntegrationLockRegistry";
+
+ /**
+ * The {@value DEFAULT_PARTITION_KEY_NAME} default name for the partition key in the table.
+ */
+ public static final String DEFAULT_PARTITION_KEY_NAME = "lockKey";
+
+ /**
+ * The {@value DEFAULT_SORT_KEY_NAME} default name for the sort key in the table.
+ */
+ public static final String DEFAULT_SORT_KEY_NAME = "sortKey";
+
+ /**
+ * The {@value DEFAULT_SORT_KEY} default value for the sort key in the table.
+ */
+ public static final String DEFAULT_SORT_KEY = "SpringIntegrationLocks";
+
+ /**
+ * The {@value DEFAULT_REFRESH_PERIOD_MS} default period in milliseconds between DB polling requests.
+ */
+ public static final long DEFAULT_REFRESH_PERIOD_MS = 1000L;
+
+ private static final Log logger = LogFactory.getLog(DynamoDbLockRegistry.class);
+
+ private final Map locks = new ConcurrentHashMap<>();
+
+ private final CountDownLatch createTableLatch = new CountDownLatch(1);
+
+ private final AtomicBoolean running = new AtomicBoolean();
+
+ private final AmazonDynamoDB dynamoDB;
+
+ private final String tableName;
+
+ private AmazonDynamoDBLockClient dynamoDBLockClient;
+
+ private boolean dynamoDBLockClientExplicitlySet;
+
+ private long readCapacity = 1L;
+
+ private long writeCapacity = 1L;
+
+ private String partitionKey = DEFAULT_PARTITION_KEY_NAME;
+
+ private String sortKeyName = DEFAULT_SORT_KEY_NAME;
+
+ private String sortKey = DEFAULT_SORT_KEY;
+
+ private long refreshPeriod = DEFAULT_REFRESH_PERIOD_MS;
+
+ private long leaseDuration = 20L;
+
+ private long heartbeatPeriod = 5L;
+
+ /**
+ * An {@link ExecutorService} to call {@link AmazonDynamoDBLockClient#releaseLock(LockItem)}
+ * in the separate thread when the current one is interrupted.
+ */
+ private Executor executor =
+ Executors.newCachedThreadPool(new CustomizableThreadFactory("dynamodb-lock-registry-"));
+
+ /**
+ * Flag to denote whether the {@link ExecutorService} was provided via the setter and
+ * thus should not be shutdown when {@link #destroy()} is called.
+ */
+ private boolean executorExplicitlySet;
+
+
+ public DynamoDbLockRegistry(AmazonDynamoDB dynamoDB) {
+ this(dynamoDB, DEFAULT_TABLE_NAME);
+ }
+
+ public DynamoDbLockRegistry(AmazonDynamoDB dynamoDB, String tableName) {
+ Assert.notNull(dynamoDB, "'dynamoDB' must not be null");
+ Assert.hasText(tableName, "'tableName' must not be empty");
+
+ this.dynamoDB = dynamoDB;
+ this.tableName = tableName;
+ }
+
+ public DynamoDbLockRegistry(AmazonDynamoDBLockClient dynamoDBLockClient) {
+ Assert.notNull(dynamoDBLockClient, "'dynamoDBLockClient' must not be null");
+
+ this.dynamoDBLockClient = dynamoDBLockClient;
+ this.dynamoDBLockClientExplicitlySet = true;
+ this.dynamoDB = null;
+ this.tableName = null;
+ }
+
+ public void setReadCapacity(long readCapacity) {
+ this.readCapacity = readCapacity;
+ }
+
+ public void setWriteCapacity(long writeCapacity) {
+ this.writeCapacity = writeCapacity;
+ }
+
+ public void setPartitionKey(String partitionKey) {
+ Assert.hasText(partitionKey, "'partitionKey' must not be empty");
+ this.partitionKey = partitionKey;
+ }
+
+ /**
+ * Specify a name of the table attribute which is used as a sort key.
+ * @param sortKeyName the sort key attribute name to use.
+ */
+ public void setSortKeyName(String sortKeyName) {
+ this.sortKeyName = sortKeyName;
+ }
+
+ /**
+ * Specify a value for the sort key attribute of the lock item.
+ * @param sortKey the sort key value to use.
+ */
+ public void setSortKey(String sortKey) {
+ this.sortKey = sortKey;
+ }
+
+ public void setLeaseDuration(long leaseDuration) {
+ this.leaseDuration = leaseDuration;
+ }
+
+ public void setHeartbeatPeriod(long heartbeatPeriod) {
+ this.heartbeatPeriod = heartbeatPeriod;
+ }
+
+ public void setRefreshPeriod(long refreshPeriod) {
+ this.refreshPeriod = refreshPeriod;
+ }
+
+ /**
+ * Set the {@link Executor}, where is not provided then a default of
+ * cached thread pool Executor will be used.
+ * @param executor the executor service
+ */
+ public void setExecutor(Executor executor) {
+ this.executor = executor;
+ this.executorExplicitlySet = true;
+ }
+
+ @Override
+ public void afterPropertiesSet() {
+ if (!this.dynamoDBLockClientExplicitlySet) {
+ AmazonDynamoDBLockClientOptions dynamoDBLockClientOptions =
+ AmazonDynamoDBLockClientOptions
+ .builder(this.dynamoDB, this.tableName)
+ .withPartitionKeyName(this.partitionKey)
+ .withSortKeyName(this.sortKeyName)
+ .withHeartbeatPeriod(this.heartbeatPeriod)
+ .withLeaseDuration(this.leaseDuration)
+ .build();
+
+ this.dynamoDBLockClient = new AmazonDynamoDBLockClient(dynamoDBLockClientOptions);
+ }
+
+ this.leaseDuration =
+ (long) new DirectFieldAccessor(this.dynamoDBLockClient)
+ .getPropertyValue("leaseDurationInMilliseconds");
+
+
+ this.executor.execute(() -> {
+ try {
+ if (!this.dynamoDBLockClientExplicitlySet) {
+ try {
+ this.dynamoDBLockClient.assertLockTableExists();
+ return;
+ }
+ catch (LockTableDoesNotExistException e) {
+ if (logger.isInfoEnabled()) {
+ logger.info("No table '" + this.tableName + "'. Creating one...");
+ }
+ }
+
+ CreateDynamoDBTableOptions createDynamoDBTableOptions =
+ CreateDynamoDBTableOptions
+ .builder(this.dynamoDB,
+ new ProvisionedThroughput(this.readCapacity, this.writeCapacity),
+ this.tableName)
+ .withPartitionKeyName(this.partitionKey)
+ .withSortKeyName(this.sortKeyName)
+ .build();
+
+ AmazonDynamoDBLockClient.createLockTableInDynamoDB(createDynamoDBTableOptions);
+ }
+
+ int i = 0;
+ // We need up to one minute to wait until table is created on AWS.
+ while (i++ < 60) {
+ if (this.dynamoDBLockClient.lockTableExists()) {
+ return;
+ }
+ else {
+ try {
+ // This is allowed minimum for constant AWS requests.
+ Thread.sleep(1000);
+ }
+ catch (InterruptedException e) {
+ ReflectionUtils.rethrowRuntimeException(e);
+ }
+ }
+ }
+
+ logger.error("Cannot describe DynamoDb table: " + this.tableName);
+ }
+ finally {
+ // Release create table barrier either way.
+ // If there is an error during creation/description,
+ // we deffer the actual ResourceNotFoundException to the end-user active calls.
+ this.createTableLatch.countDown();
+ }
+ });
+ }
+
+ private void awaitForActive() {
+ IllegalStateException illegalStateException =
+ new IllegalStateException(
+ "The DynamoDb table " + this.tableName + " has not been created during " + 60 + " seconds");
+ try {
+ if (!this.createTableLatch.await(60, TimeUnit.SECONDS)) {
+ throw illegalStateException;
+ }
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw illegalStateException;
+ }
+ }
+
+ @Override
+ public void destroy() throws Exception {
+ if (!this.executorExplicitlySet) {
+ ((ExecutorService) this.executor).shutdown();
+ }
+
+ if (!this.dynamoDBLockClientExplicitlySet) {
+ this.dynamoDBLockClient.close();
+ }
+ }
+
+ @Override
+ public Lock obtain(Object lockKey) {
+ Assert.isInstanceOf(String.class, lockKey, "'lockKey' must of String type");
+ return this.locks.computeIfAbsent((String) lockKey, DynamoDbLock::new);
+ }
+
+ @Override
+ public void expireUnusedOlderThan(long age) {
+ Iterator> iterator = this.locks.entrySet().iterator();
+ long now = System.currentTimeMillis();
+ while (iterator.hasNext()) {
+ Map.Entry entry = iterator.next();
+ DynamoDbLock lock = entry.getValue();
+ if (now - lock.lastUsed > age && !lock.delegate.isHeldByCurrentThread()) {
+ iterator.remove();
+ }
+ }
+ }
+
+ private final class DynamoDbLock implements Lock {
+
+ private final ReentrantLock delegate = new ReentrantLock();
+
+ private final String key;
+
+ // It is safe to use a shared instance - access is guaranteed by the delegate lock.
+ private final AcquireLockOptions.AcquireLockOptionsBuilder acquireLockOptionsBuilder;
+
+ private LockItem lockItem;
+
+ private volatile long lastUsed = System.currentTimeMillis();
+
+ private DynamoDbLock(String key) {
+ this.key = key;
+ this.acquireLockOptionsBuilder =
+ AcquireLockOptions.builder(this.key)
+ .withReplaceData(false)
+ .withSortKey(DynamoDbLockRegistry.this.sortKey)
+ .withTimeUnit(TimeUnit.MILLISECONDS)
+ .withRefreshPeriod(DynamoDbLockRegistry.this.refreshPeriod);
+ }
+
+ private void rethrowAsLockException(Exception e) {
+ throw new CannotAcquireLockException("Failed to lock at " + this.key, e);
+ }
+
+ @Override
+ public void lock() {
+ awaitForActive();
+
+ this.delegate.lock();
+
+ this.acquireLockOptionsBuilder
+ .withAdditionalTimeToWaitForLock(Long.MAX_VALUE - DynamoDbLockRegistry.this.leaseDuration);
+
+ boolean wasInterruptedWhileUninterruptible = false;
+
+ try {
+ while (true) {
+ try {
+ while (!doLock()) {
+ Thread.sleep(100); //NOSONAR
+ }
+ break;
+ }
+ catch (InterruptedException e) {
+ /*
+ * This method must be uninterruptible so catch and ignore
+ * interrupts and only break out of the while loop when
+ * we get the lock.
+ */
+ wasInterruptedWhileUninterruptible = true;
+ }
+ catch (Exception e) {
+ this.delegate.unlock();
+ rethrowAsLockException(e);
+ }
+ }
+ }
+ finally {
+ if (wasInterruptedWhileUninterruptible) {
+ Thread.currentThread().interrupt();
+ }
+ }
+
+ }
+
+ @Override
+ public void lockInterruptibly() throws InterruptedException {
+ awaitForActive();
+
+ this.delegate.lockInterruptibly();
+
+ this.acquireLockOptionsBuilder
+ .withAdditionalTimeToWaitForLock(Long.MAX_VALUE - DynamoDbLockRegistry.this.leaseDuration);
+
+ try {
+ while (!doLock()) {
+ Thread.sleep(100); //NOSONAR
+ if (Thread.currentThread().isInterrupted()) {
+ throw new InterruptedException();
+ }
+ }
+ }
+ catch (InterruptedException ie) {
+ this.delegate.unlock();
+ Thread.currentThread().interrupt();
+ throw ie;
+ }
+ catch (Exception e) {
+ this.delegate.unlock();
+ rethrowAsLockException(e);
+ }
+ }
+
+ @Override
+ public boolean tryLock() {
+ awaitForActive();
+
+ try {
+ return tryLock(0, TimeUnit.MILLISECONDS);
+ }
+ catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ return false;
+ }
+ }
+
+ @Override
+ public boolean tryLock(long time, TimeUnit unit) throws InterruptedException {
+ awaitForActive();
+
+ long start = System.currentTimeMillis();
+
+ if (!this.delegate.tryLock(time, unit)) {
+ return false;
+ }
+
+ this.acquireLockOptionsBuilder
+ .withAdditionalTimeToWaitForLock(
+ System.currentTimeMillis() - start + TimeUnit.MILLISECONDS.convert(time, unit));
+
+ boolean acquired = false;
+ try {
+ acquired = doLock();
+
+ if (!acquired) {
+ this.delegate.unlock();
+ }
+ else {
+ this.lastUsed = System.currentTimeMillis();
+ }
+ }
+ catch (Exception e) {
+ this.delegate.unlock();
+ rethrowAsLockException(e);
+ }
+
+ return acquired;
+ }
+
+ private boolean doLock() throws InterruptedException {
+ boolean acquired;
+ if (this.lockItem != null) {
+ this.lockItem.sendHeartBeat();
+ acquired = true;
+ }
+ else {
+ this.lockItem =
+ DynamoDbLockRegistry.this.dynamoDBLockClient
+ .tryAcquireLock(this.acquireLockOptionsBuilder.build())
+ .orElse(null);
+
+ acquired = this.lockItem != null;
+ }
+
+ if (acquired) {
+ this.lastUsed = System.currentTimeMillis();
+ }
+
+ return acquired;
+ }
+
+ @Override
+ public void unlock() {
+ if (!this.delegate.isHeldByCurrentThread()) {
+ throw new IllegalMonitorStateException("You do not own lock at " + this.key);
+ }
+ if (this.delegate.getHoldCount() > 1) {
+ this.delegate.unlock();
+ return;
+ }
+ try {
+ if (Thread.currentThread().isInterrupted()) {
+ LockItem lockItemToRelease = this.lockItem;
+ DynamoDbLockRegistry.this.executor.execute(() ->
+ DynamoDbLockRegistry.this.dynamoDBLockClient.releaseLock(lockItemToRelease)
+ );
+ }
+ else {
+ DynamoDbLockRegistry.this.dynamoDBLockClient.releaseLock(this.lockItem);
+ }
+ }
+ catch (Exception e) {
+ throw new DataAccessResourceFailureException("Failed to release lock at " + this.key, e);
+ }
+ finally {
+ this.lockItem = null;
+ this.delegate.unlock();
+ }
+ }
+
+ @Override
+ public Condition newCondition() {
+ throw new UnsupportedOperationException("DynamoDb locks don't support conditions.");
+ }
+
+ @Override
+ public String toString() {
+ SimpleDateFormat dateFormat = new SimpleDateFormat("YYYY-MM-dd@HH:mm:ss.SSS");
+ return "DynamoDbLock [lockKey=" + this.key
+ + ",lockedAt=" + dateFormat.format(new Date(this.lastUsed))
+ + ", lockItem=" + this.lockItem
+ + "]";
+ }
+
+ }
+
+}
diff --git a/src/main/java/org/springframework/integration/aws/lock/package-info.java b/src/main/java/org/springframework/integration/aws/lock/package-info.java
new file mode 100644
index 0000000..34693f1
--- /dev/null
+++ b/src/main/java/org/springframework/integration/aws/lock/package-info.java
@@ -0,0 +1,4 @@
+/**
+ * Provides classes supporting lock registry.
+ */
+package org.springframework.integration.aws.lock;
diff --git a/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java b/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java
index 44ef72d..511e7d1 100644
--- a/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java
+++ b/src/test/java/org/springframework/integration/aws/inbound/KinesisMessageDrivenChannelAdapterTests.java
@@ -207,7 +207,7 @@ public class KinesisMessageDrivenChannelAdapterTests {
assertThat(n).isLessThan(100);
// When resharding happens the describeStream() is performed again
- verify(this.amazonKinesisForResharding, atLeast(2)).describeStream(any(DescribeStreamRequest.class));
+ verify(this.amazonKinesisForResharding, atLeast(1)).describeStream(any(DescribeStreamRequest.class));
this.reshardingChannelAdapter.stop();
}
diff --git a/src/test/java/org/springframework/integration/aws/leader/DynamoDbLockRegistryLeaderInitiatorTests.java b/src/test/java/org/springframework/integration/aws/leader/DynamoDbLockRegistryLeaderInitiatorTests.java
new file mode 100644
index 0000000..dd6b5ed
--- /dev/null
+++ b/src/test/java/org/springframework/integration/aws/leader/DynamoDbLockRegistryLeaderInitiatorTests.java
@@ -0,0 +1,249 @@
+/*
+ * Copyright 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.leader;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+
+import org.springframework.integration.aws.DynamoDbLocalRunning;
+import org.springframework.integration.aws.lock.DynamoDbLockRegistry;
+import org.springframework.integration.leader.Context;
+import org.springframework.integration.leader.DefaultCandidate;
+import org.springframework.integration.leader.event.LeaderEventPublisher;
+import org.springframework.integration.support.leader.LockRegistryLeaderInitiator;
+import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
+import com.amazonaws.services.dynamodbv2.model.DescribeTableRequest;
+import com.amazonaws.waiters.FixedDelayStrategy;
+import com.amazonaws.waiters.MaxAttemptsRetryStrategy;
+import com.amazonaws.waiters.PollingStrategy;
+import com.amazonaws.waiters.Waiter;
+import com.amazonaws.waiters.WaiterParameters;
+
+/**
+ * @author Artem Bilan
+ *
+ * @since 2.0
+ */
+public class DynamoDbLockRegistryLeaderInitiatorTests {
+
+ @ClassRule
+ public static final DynamoDbLocalRunning DYNAMO_DB_RUNNING = DynamoDbLocalRunning.isRunning(4567);
+
+ private static AmazonDynamoDBAsync dynamoDB;
+
+ @BeforeClass
+ public static void init() {
+ dynamoDB = DYNAMO_DB_RUNNING.getDynamoDB();
+
+ try {
+ dynamoDB.deleteTableAsync(DynamoDbLockRegistry.DEFAULT_TABLE_NAME);
+
+ Waiter waiter =
+ dynamoDB.waiters()
+ .tableNotExists();
+
+ waiter.run(new WaiterParameters<>(new DescribeTableRequest(DynamoDbLockRegistry.DEFAULT_TABLE_NAME))
+ .withPollingStrategy(new PollingStrategy(new MaxAttemptsRetryStrategy(25),
+ new FixedDelayStrategy(1))));
+ }
+ catch (Exception e) {
+
+ }
+ }
+
+ @AfterClass
+ public static void destroy() {
+ dynamoDB.deleteTable(DynamoDbLockRegistry.DEFAULT_TABLE_NAME);
+ }
+
+ @Test
+ public void testDistributedLeaderElection() throws Exception {
+ CountDownLatch granted = new CountDownLatch(1);
+ CountingPublisher countingPublisher = new CountingPublisher(granted);
+ List registries = new ArrayList<>();
+ List initiators = new ArrayList<>();
+ for (int i = 0; i < 2; i++) {
+ DynamoDbLockRegistry lockRepository = new DynamoDbLockRegistry(dynamoDB);
+ lockRepository.afterPropertiesSet();
+ registries.add(lockRepository);
+
+ LockRegistryLeaderInitiator initiator =
+ new LockRegistryLeaderInitiator(
+ lockRepository,
+ new DefaultCandidate("foo#" + i, "bar"));
+ initiator.setExecutorService(
+ Executors.newSingleThreadExecutor(new CustomizableThreadFactory("lock-leadership-" + i + "-")));
+ initiator.setLeaderEventPublisher(countingPublisher);
+ initiators.add(initiator);
+ }
+
+ for (LockRegistryLeaderInitiator initiator : initiators) {
+ initiator.start();
+ }
+
+ assertThat(granted.await(10, TimeUnit.SECONDS)).isTrue();
+
+ LockRegistryLeaderInitiator initiator1 = countingPublisher.initiator;
+
+ LockRegistryLeaderInitiator initiator2 = null;
+
+ for (LockRegistryLeaderInitiator initiator : initiators) {
+ if (initiator != initiator1) {
+ initiator2 = initiator;
+ break;
+ }
+ }
+
+ assertThat(initiator2).isNotNull();
+
+ assertThat(initiator1.getContext().isLeader()).isTrue();
+ assertThat(initiator2.getContext().isLeader()).isFalse();
+
+ final CountDownLatch granted1 = new CountDownLatch(1);
+ final CountDownLatch granted2 = new CountDownLatch(1);
+ CountDownLatch revoked1 = new CountDownLatch(1);
+ CountDownLatch revoked2 = new CountDownLatch(1);
+ CountDownLatch acquireLockFailed1 = new CountDownLatch(1);
+ CountDownLatch acquireLockFailed2 = new CountDownLatch(1);
+
+ initiator1.setLeaderEventPublisher(new CountingPublisher(granted1, revoked1, acquireLockFailed1));
+
+ initiator2.setLeaderEventPublisher(new CountingPublisher(granted2, revoked2, acquireLockFailed2));
+
+ // It's hard to see round-robin election, so let's make the yielding initiator to sleep long before restarting
+ initiator1.setBusyWaitMillis(1000);
+
+ initiator1.getContext().yield();
+
+ assertThat(revoked1.await(20, TimeUnit.SECONDS)).isTrue();
+ assertThat(granted2.await(20, TimeUnit.SECONDS)).isTrue();
+
+ assertThat(initiator2.getContext().isLeader()).isTrue();
+ assertThat(initiator1.getContext().isLeader()).isFalse();
+
+ initiator1.setBusyWaitMillis(LockRegistryLeaderInitiator.DEFAULT_BUSY_WAIT_TIME);
+ initiator2.setBusyWaitMillis(10000);
+
+ initiator2.getContext().yield();
+
+ assertThat(revoked2.await(20, TimeUnit.SECONDS)).isTrue();
+ assertThat(granted1.await(20, TimeUnit.SECONDS)).isTrue();
+
+ assertThat(initiator1.getContext().isLeader()).isTrue();
+ assertThat(initiator2.getContext().isLeader()).isFalse();
+
+ initiator2.stop();
+
+ CountDownLatch revoked11 = new CountDownLatch(1);
+ initiator1.setLeaderEventPublisher(new CountingPublisher(new CountDownLatch(1), revoked11,
+ new CountDownLatch(1)));
+
+ initiator1.getContext().yield();
+
+ assertThat(revoked11.await(10, TimeUnit.SECONDS)).isTrue();
+ assertThat(initiator1.getContext().isLeader()).isFalse();
+
+ for (DynamoDbLockRegistry registry : registries) {
+ registry.destroy();
+ }
+ }
+
+ @Test
+ public void testLostConnection() throws Exception {
+ CountDownLatch granted = new CountDownLatch(1);
+ CountingPublisher countingPublisher = new CountingPublisher(granted);
+
+ DynamoDbLockRegistry lockRepository = new DynamoDbLockRegistry(dynamoDB);
+ lockRepository.afterPropertiesSet();
+
+ LockRegistryLeaderInitiator initiator = new LockRegistryLeaderInitiator(lockRepository);
+ initiator.setLeaderEventPublisher(countingPublisher);
+
+ initiator.start();
+
+ assertThat(granted.await(10, TimeUnit.SECONDS)).isTrue();
+
+ destroy();
+
+ assertThat(countingPublisher.revoked.await(10, TimeUnit.SECONDS)).isTrue();
+
+ granted = new CountDownLatch(1);
+ countingPublisher = new CountingPublisher(granted);
+ initiator.setLeaderEventPublisher(countingPublisher);
+
+ init();
+
+ lockRepository.afterPropertiesSet();
+
+ assertThat(granted.await(10, TimeUnit.SECONDS)).isTrue();
+
+ initiator.stop();
+
+ lockRepository.destroy();
+ }
+
+ private static class CountingPublisher implements LeaderEventPublisher {
+
+ private final CountDownLatch granted;
+
+ private final CountDownLatch revoked;
+
+ private final CountDownLatch acquireLockFailed;
+
+ private volatile LockRegistryLeaderInitiator initiator;
+
+ CountingPublisher(CountDownLatch granted, CountDownLatch revoked, CountDownLatch acquireLockFailed) {
+ this.granted = granted;
+ this.revoked = revoked;
+ this.acquireLockFailed = acquireLockFailed;
+ }
+
+ CountingPublisher(CountDownLatch granted) {
+ this(granted, new CountDownLatch(1), new CountDownLatch(1));
+ }
+
+ @Override
+ public void publishOnRevoked(Object source, Context context, String role) {
+ this.revoked.countDown();
+ }
+
+ @Override
+ public void publishOnFailedToAcquire(Object source, Context context, String role) {
+ this.acquireLockFailed.countDown();
+ }
+
+ @Override
+ public void publishOnGranted(Object source, Context context, String role) {
+ this.initiator = (LockRegistryLeaderInitiator) source;
+ this.granted.countDown();
+ }
+
+ }
+
+}
diff --git a/src/test/java/org/springframework/integration/aws/lock/DynamoDbLockRegistryTests.java b/src/test/java/org/springframework/integration/aws/lock/DynamoDbLockRegistryTests.java
new file mode 100644
index 0000000..f0ff009
--- /dev/null
+++ b/src/test/java/org/springframework/integration/aws/lock/DynamoDbLockRegistryTests.java
@@ -0,0 +1,319 @@
+/*
+ * Copyright 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.lock;
+
+import static org.assertj.core.api.Java6Assertions.assertThat;
+
+import java.util.Map;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.locks.Lock;
+
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.core.task.AsyncTaskExecutor;
+import org.springframework.core.task.SimpleAsyncTaskExecutor;
+import org.springframework.integration.aws.DynamoDbLocalRunning;
+import org.springframework.integration.test.util.TestUtils;
+import org.springframework.test.annotation.DirtiesContext;
+import org.springframework.test.context.junit4.SpringRunner;
+
+import com.amazonaws.services.dynamodbv2.AmazonDynamoDBAsync;
+import com.amazonaws.services.dynamodbv2.model.DescribeTableRequest;
+import com.amazonaws.waiters.FixedDelayStrategy;
+import com.amazonaws.waiters.MaxAttemptsRetryStrategy;
+import com.amazonaws.waiters.PollingStrategy;
+import com.amazonaws.waiters.Waiter;
+import com.amazonaws.waiters.WaiterParameters;
+
+/**
+ * @author Artem Bilan
+ *
+ * @since 2.0
+ */
+@RunWith(SpringRunner.class)
+@DirtiesContext
+public class DynamoDbLockRegistryTests {
+
+ @ClassRule
+ public static final DynamoDbLocalRunning DYNAMO_DB_RUNNING = DynamoDbLocalRunning.isRunning(4567);
+
+ private final AsyncTaskExecutor taskExecutor = new SimpleAsyncTaskExecutor();
+
+ @Autowired
+ private DynamoDbLockRegistry dynamoDbLockRegistry;
+
+ @BeforeClass
+ public static void setup() {
+ AmazonDynamoDBAsync dynamoDB = DYNAMO_DB_RUNNING.getDynamoDB();
+
+ try {
+ dynamoDB.deleteTableAsync(DynamoDbLockRegistry.DEFAULT_TABLE_NAME);
+
+ Waiter waiter =
+ dynamoDB.waiters()
+ .tableNotExists();
+
+ waiter.run(new WaiterParameters<>(new DescribeTableRequest(DynamoDbLockRegistry.DEFAULT_TABLE_NAME))
+ .withPollingStrategy(new PollingStrategy(new MaxAttemptsRetryStrategy(25),
+ new FixedDelayStrategy(1))));
+ }
+ catch (Exception e) {
+
+ }
+ }
+
+ @Before
+ public void clear() {
+ this.dynamoDbLockRegistry.expireUnusedOlderThan(0);
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testLock() {
+ for (int i = 0; i < 10; i++) {
+ Lock lock = this.dynamoDbLockRegistry.obtain("foo");
+ lock.lock();
+ try {
+ assertThat(TestUtils.getPropertyValue(this.dynamoDbLockRegistry, "locks", Map.class)).hasSize(1);
+ }
+ finally {
+ lock.unlock();
+ }
+ }
+ }
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testLockInterruptibly() throws Exception {
+ for (int i = 0; i < 10; i++) {
+ Lock lock = this.dynamoDbLockRegistry.obtain("foo");
+ lock.lockInterruptibly();
+ try {
+ assertThat(TestUtils.getPropertyValue(this.dynamoDbLockRegistry, "locks", Map.class)).hasSize(1);
+ }
+ finally {
+ lock.unlock();
+ }
+ }
+ }
+
+ @Test
+ public void testReentrantLock() {
+ for (int i = 0; i < 10; i++) {
+ Lock lock1 = this.dynamoDbLockRegistry.obtain("foo");
+ lock1.lock();
+ try {
+ Lock lock2 = this.dynamoDbLockRegistry.obtain("foo");
+ assertThat(lock1).isSameAs(lock2);
+ lock2.lock();
+ lock2.unlock();
+ }
+ finally {
+ lock1.unlock();
+ }
+ }
+ }
+
+ @Test
+ public void testReentrantLockInterruptibly() throws Exception {
+ for (int i = 0; i < 10; i++) {
+ Lock lock1 = this.dynamoDbLockRegistry.obtain("foo");
+ lock1.lockInterruptibly();
+ try {
+ Lock lock2 = this.dynamoDbLockRegistry.obtain("foo");
+ assertThat(lock1).isSameAs(lock2);
+ lock2.lockInterruptibly();
+ lock2.unlock();
+ }
+ finally {
+ lock1.unlock();
+ }
+ }
+ }
+
+ @Test
+ public void testTwoLocks() throws Exception {
+ for (int i = 0; i < 10; i++) {
+ Lock lock1 = this.dynamoDbLockRegistry.obtain("foo");
+ lock1.lockInterruptibly();
+ try {
+ Lock lock2 = this.dynamoDbLockRegistry.obtain("bar");
+ assertThat(lock1).isNotSameAs(lock2);
+ lock2.lockInterruptibly();
+ lock2.unlock();
+ }
+ finally {
+ lock1.unlock();
+ }
+ }
+ }
+
+ @Test
+ public void testTwoThreadsSecondFailsToGetLock() throws Exception {
+ final Lock lock1 = this.dynamoDbLockRegistry.obtain("foo");
+ lock1.lockInterruptibly();
+ final AtomicBoolean locked = new AtomicBoolean();
+ final CountDownLatch latch = new CountDownLatch(1);
+ Future