GH-66: Add DynamoDbLockRegistry implementation (#93)

* GH-66: Add DynamoDbLockRegistry implementation

Fixes https://github.com/spring-projects/spring-integration-aws/issues/66

* * Remove `Lifecycle` from `DynamoDbLockRegistry` in favor of a thread
execution in the `afterPropertiesSet()`
* Fix `lock()` interruptibility logic

* * Remove `mavenLocal()` since the upstream PR is merged
* decrease an amount of expectations in the Kinesis test

* * Upgrade to SC-AWS-2.0.0.RC2

* * Add Docs for the `DynamoDbLockRegistry`

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 

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)
 

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

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}.
+ * <p>
+ * 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<String, DynamoDbLock> 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<Map.Entry<String, DynamoDbLock>> iterator = this.locks.entrySet().iterator();
+		long now = System.currentTimeMillis();
+		while (iterator.hasNext()) {
+			Map.Entry<String, DynamoDbLock> 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
+					+ "]";
+		}
+
+	}
+
+}

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;

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();
 	}

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<DescribeTableRequest> 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<DynamoDbLockRegistry> registries = new ArrayList<>();
+		List<LockRegistryLeaderInitiator> 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();
+		}
+
+	}
+
+}

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<DescribeTableRequest> 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<Object> result = this.taskExecutor.submit(() -> {
+			Lock lock2 = this.dynamoDbLockRegistry.obtain("foo");
+			locked.set(lock2.tryLock(200, TimeUnit.MILLISECONDS));
+			latch.countDown();
+			try {
+				lock2.unlock();
+			}
+			catch (Exception e) {
+				return e;
+			}
+			return null;
+		});
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isFalse();
+		lock1.unlock();
+		Object ise = result.get(10, TimeUnit.SECONDS);
+		assertThat(ise).isInstanceOf(IllegalMonitorStateException.class);
+		assertThat(((Exception) ise).getMessage()).contains("You do not own");
+	}
+
+	@Test
+	public void testTwoThreads() throws Exception {
+		final Lock lock1 = this.dynamoDbLockRegistry.obtain("foo");
+		final AtomicBoolean locked = new AtomicBoolean();
+		final CountDownLatch latch1 = new CountDownLatch(1);
+		final CountDownLatch latch2 = new CountDownLatch(1);
+		final CountDownLatch latch3 = new CountDownLatch(1);
+		lock1.lockInterruptibly();
+		this.taskExecutor.execute(() -> {
+			Lock lock2 = this.dynamoDbLockRegistry.obtain("foo");
+			try {
+				latch1.countDown();
+				lock2.lockInterruptibly();
+				latch2.await(10, TimeUnit.SECONDS);
+				locked.set(true);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+			}
+			finally {
+				lock2.unlock();
+				latch3.countDown();
+			}
+		});
+
+		assertThat(latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isFalse();
+
+		lock1.unlock();
+		latch2.countDown();
+
+		assertThat(latch3.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isTrue();
+	}
+
+	@Test
+	public void testTwoThreadsDifferentRegistries() throws Exception {
+		final DynamoDbLockRegistry registry1 = new DynamoDbLockRegistry(DYNAMO_DB_RUNNING.getDynamoDB());
+		registry1.afterPropertiesSet();
+		final DynamoDbLockRegistry registry2 = new DynamoDbLockRegistry(DYNAMO_DB_RUNNING.getDynamoDB());
+		registry2.afterPropertiesSet();
+
+		final Lock lock1 = registry1.obtain("foo");
+		final AtomicBoolean locked = new AtomicBoolean();
+		final CountDownLatch latch1 = new CountDownLatch(1);
+		final CountDownLatch latch2 = new CountDownLatch(1);
+		final CountDownLatch latch3 = new CountDownLatch(1);
+		lock1.lockInterruptibly();
+		this.taskExecutor.execute(() -> {
+			Lock lock2 = registry2.obtain("foo");
+			try {
+				latch1.countDown();
+				lock2.lockInterruptibly();
+				latch2.await(10, TimeUnit.SECONDS);
+				locked.set(true);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+			}
+			finally {
+				lock2.unlock();
+				latch3.countDown();
+			}
+		});
+		assertThat(latch1.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isFalse();
+
+		lock1.unlock();
+		latch2.countDown();
+
+		assertThat(latch3.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isTrue();
+
+		registry1.destroy();
+		registry2.destroy();
+	}
+
+	@Test
+	public void testTwoThreadsWrongOneUnlocks() throws Exception {
+		final Lock lock = this.dynamoDbLockRegistry.obtain("foo");
+		lock.lockInterruptibly();
+		final AtomicBoolean locked = new AtomicBoolean();
+		final CountDownLatch latch = new CountDownLatch(1);
+		Future<Object> result =
+				this.taskExecutor.submit(() -> {
+					try {
+						lock.unlock();
+					}
+					catch (Exception e) {
+						latch.countDown();
+						return e;
+					}
+					return null;
+				});
+
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		assertThat(locked.get()).isFalse();
+
+		lock.unlock();
+		Object imse = result.get(10, TimeUnit.SECONDS);
+		assertThat(imse).isInstanceOf(IllegalMonitorStateException.class);
+		assertThat(((Exception) imse).getMessage()).contains("You do not own");
+	}
+
+
+	@Configuration
+	public static class ContextConfiguration {
+
+		@Bean
+		public DynamoDbLockRegistry dynamoDbLockRegistry() {
+			DynamoDbLockRegistry dynamoDbLockRegistry = new DynamoDbLockRegistry(DYNAMO_DB_RUNNING.getDynamoDB());
+			dynamoDbLockRegistry.setHeartbeatPeriod(1);
+			dynamoDbLockRegistry.setRefreshPeriod(10);
+			return dynamoDbLockRegistry;
+		}
+
+	}
+
+}