INT-4248-4.3.x: Fix LockRegistryLeaderInitiator

JIRA: https://jira.spring.io/browse/INT-4248

The `RedisLockRegistry` doesn't update store for `expire` property on the
subsequent lock interaction.
Only `RedisLockRegistry.obtain()` lets us to know if we expired already or not

To minimize the time window between `lock` and `expire` facts, move `obtain()`
function in the `LockRegistryLeaderInitiator` close to `Lock` usage.

spring-integration-core/src/main/java/org/springframework/integration/support/leader/LockRegistryLeaderInitiator.java

@@ -22,7 +22,6 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.locks.Lock;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -56,6 +55,7 @@ import org.springframework.util.Assert;
  * @author Dave Syer
  * @author Artem Bilan
  * @author Vedran Pavic
+ *
  * @since 4.3.1
  */
 public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBean, ApplicationEventPublisherAware {
@@ -299,8 +299,6 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 
 	protected class LeaderSelector implements Callable<Void> {
 
-		private final Lock lock;
-
 		private final String lockKey;
 
 		private final LockContext context = new LockContext();
@@ -308,7 +306,6 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 		private volatile boolean locked = false;
 
 		LeaderSelector(String lockKey) {
-			this.lock = LockRegistryLeaderInitiator.this.locks.obtain(lockKey);
 			this.lockKey = lockKey;
 		}
 
@@ -318,8 +315,9 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 				while (LockRegistryLeaderInitiator.this.running) {
 					try {
 						// We always try to acquire the lock, in case it expired
-						boolean acquired = this.lock.tryLock(LockRegistryLeaderInitiator.this.heartBeatMillis,
+						// TODO obtain(this.lockKey) because of INT-4248. Should be fixed in 5.0
-								TimeUnit.MILLISECONDS);
+						boolean acquired = LockRegistryLeaderInitiator.this.locks.obtain(this.lockKey)
+								.tryLock(LockRegistryLeaderInitiator.this.heartBeatMillis, TimeUnit.MILLISECONDS);
 						if (!this.locked) {
 							if (acquired) {
 								// Success: we are now leader
@@ -330,7 +328,8 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 						else if (acquired) {
 							// If we were able to acquire it but we were already locked we
 							// should release it
-							this.lock.unlock();
+							LockRegistryLeaderInitiator.this.locks.obtain(this.lockKey)
+									.unlock();
 							// Give it a chance to expire.
 							Thread.sleep(LockRegistryLeaderInitiator.this.heartBeatMillis);
 						}
@@ -344,7 +343,8 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 					}
 					catch (InterruptedException e) {
 						if (this.locked) {
-							this.lock.unlock();
+							LockRegistryLeaderInitiator.this.locks.obtain(this.lockKey)
+									.unlock();
 							this.locked = false;
 							// The lock was broken and we are no longer leader
 							handleRevoked();
@@ -358,7 +358,8 @@ public class LockRegistryLeaderInitiator implements SmartLifecycle, DisposableBe
 			}
 			finally {
 				if (this.locked) {
-					this.lock.unlock();
+					LockRegistryLeaderInitiator.this.locks.obtain(this.lockKey)
+							.unlock();
 					// We are stopping, therefore not leading any more
 					handleRevoked();
 				}

spring-integration-redis/src/test/java/org/springframework/integration/redis/leader/RedisLockRegistryLeaderInitiatorTests.java

@@ -0,0 +1,175 @@
+/*
+ * Copyright 2016-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.redis.leader;
+
+import static org.hamcrest.Matchers.is;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThat;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.Test;
+
+import org.springframework.integration.leader.Context;
+import org.springframework.integration.leader.DefaultCandidate;
+import org.springframework.integration.leader.event.LeaderEventPublisher;
+import org.springframework.integration.redis.rules.RedisAvailable;
+import org.springframework.integration.redis.rules.RedisAvailableTests;
+import org.springframework.integration.redis.util.RedisLockRegistry;
+import org.springframework.integration.support.leader.LockRegistryLeaderInitiator;
+
+/**
+ * @author Artem Bilan
+ * @author Gary Russell
+ *
+ * @since 4.3.9
+ */
+public class RedisLockRegistryLeaderInitiatorTests extends RedisAvailableTests {
+
+	@Test
+	@RedisAvailable
+	public void testDistributedLeaderElection() throws Exception {
+		CountDownLatch granted = new CountDownLatch(1);
+		CountingPublisher countingPublisher = new CountingPublisher(granted);
+		List<LockRegistryLeaderInitiator> initiators = new ArrayList<LockRegistryLeaderInitiator>();
+		for (int i = 0; i < 2; i++) {
+			RedisLockRegistry registry = new RedisLockRegistry(getConnectionFactoryForTest(), "LeaderInitiator");
+			LockRegistryLeaderInitiator initiator =
+					new LockRegistryLeaderInitiator(registry, new DefaultCandidate("foo", "bar"));
+			initiator.setLeaderEventPublisher(countingPublisher);
+			initiators.add(initiator);
+		}
+
+		for (LockRegistryLeaderInitiator initiator : initiators) {
+			initiator.start();
+		}
+
+		assertThat(granted.await(10, TimeUnit.SECONDS), is(true));
+
+		LockRegistryLeaderInitiator initiator1 = countingPublisher.initiator;
+
+		LockRegistryLeaderInitiator initiator2 = null;
+
+		for (LockRegistryLeaderInitiator initiator : initiators) {
+			if (initiator != initiator1) {
+				initiator2 = initiator;
+				break;
+			}
+		}
+
+		assertNotNull(initiator2);
+
+		assertThat(initiator1.getContext().isLeader(), is(true));
+		assertThat(initiator2.getContext().isLeader(), is(false));
+
+		final CountDownLatch granted1 = new CountDownLatch(1);
+		final CountDownLatch granted2 = new CountDownLatch(1);
+		CountDownLatch revoked1 = new CountDownLatch(1);
+		CountDownLatch revoked2 = new CountDownLatch(1);
+		initiator1.setLeaderEventPublisher(new CountingPublisher(granted1, revoked1) {
+
+			@Override
+			public void publishOnRevoked(Object source, Context context, String role) {
+				try {
+					// It's difficult to see round-robin election, so block one initiator until the second is elected.
+					assertThat(granted2.await(10, TimeUnit.SECONDS), is(true));
+				}
+				catch (InterruptedException e) {
+					// No op
+				}
+				super.publishOnRevoked(source, context, role);
+			}
+
+		});
+
+		initiator2.setLeaderEventPublisher(new CountingPublisher(granted2, revoked2) {
+
+			@Override
+			public void publishOnRevoked(Object source, Context context, String role) {
+				try {
+					// It's difficult to see round-robin election, so block one initiator until the second is elected.
+					assertThat(granted1.await(10, TimeUnit.SECONDS), is(true));
+				}
+				catch (InterruptedException e) {
+					// No op
+				}
+				super.publishOnRevoked(source, context, role);
+			}
+
+		});
+
+		initiator1.getContext().yield();
+
+		assertThat(revoked1.await(10, TimeUnit.SECONDS), is(true));
+
+		assertThat(initiator2.getContext().isLeader(), is(true));
+		assertThat(initiator1.getContext().isLeader(), is(false));
+
+		initiator2.getContext().yield();
+
+		assertThat(revoked2.await(10, TimeUnit.SECONDS), is(true));
+
+		assertThat(initiator1.getContext().isLeader(), is(true));
+		assertThat(initiator2.getContext().isLeader(), is(false));
+
+		initiator2.stop();
+
+		CountDownLatch revoked11 = new CountDownLatch(1);
+		initiator1.setLeaderEventPublisher(new CountingPublisher(new CountDownLatch(1), revoked11));
+
+		initiator1.getContext().yield();
+
+		assertThat(revoked11.await(10, TimeUnit.SECONDS), is(true));
+		assertThat(initiator1.getContext().isLeader(), is(false));
+
+		initiator1.stop();
+	}
+
+	private static class CountingPublisher implements LeaderEventPublisher {
+
+		private CountDownLatch granted;
+
+		private CountDownLatch revoked;
+
+		private volatile LockRegistryLeaderInitiator initiator;
+
+		CountingPublisher(CountDownLatch granted, CountDownLatch revoked) {
+			this.granted = granted;
+			this.revoked = revoked;
+		}
+
+		CountingPublisher(CountDownLatch granted) {
+			this(granted, new CountDownLatch(1));
+		}
+
+		@Override
+		public void publishOnRevoked(Object source, Context context, String role) {
+			this.revoked.countDown();
+		}
+
+		@Override
+		public void publishOnGranted(Object source, Context context, String role) {
+			this.initiator = (LockRegistryLeaderInitiator) source;
+			this.granted.countDown();
+		}
+
+	}
+
+}

spring-integration-redis/src/test/java/org/springframework/integration/redis/util/RedisLockRegistryTests.java

@@ -514,6 +514,18 @@ public class RedisLockRegistryTests extends RedisAvailableTests {
 		lock.unlock();
 	}
 
+	@Test
+	@RedisAvailable
+	public void testExpireTwoRegistries() throws Exception {
+		RedisLockRegistry registry1 = new RedisLockRegistry(this.getConnectionFactoryForTest(), this.registryKey, 100);
+		RedisLockRegistry registry2 = new RedisLockRegistry(this.getConnectionFactoryForTest(), this.registryKey, 100);
+		assertTrue(registry1.obtain("foo").tryLock());
+		assertFalse(registry2.obtain("foo").tryLock());
+		waitForExpire("foo");
+		assertTrue(registry1.obtain("foo").tryLock());
+		assertFalse(registry2.obtain("foo").tryLock());
+	}
+
 	private Long getExpire(RedisLockRegistry registry, String lockKey) {
 		RedisTemplate<String, ?> template = this.createTemplate();
 		String registryKey = TestUtils.getPropertyValue(registry, "registryKey", String.class);