INT-2502, INT-1117 Improve Group Locking

Add initial support for Global lock registry. Fix race condition described in INT-2502.

INT-2502 addressed PR comments by @garyrussell

INT-2502 polishing

spring-integration-core/src/main/java/org/springframework/integration/aggregator/AbstractCorrelatingMessageHandler.java

@@ -17,8 +17,7 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.locks.Lock;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
@@ -37,6 +36,8 @@ import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.MessageStore;
 import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.integration.store.SimpleMessageStore;
+import org.springframework.integration.util.DefaultLockRegistry;
+import org.springframework.integration.util.LockRegistry;
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 
@@ -83,15 +84,16 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 
 	private boolean sendPartialResultOnExpiry = false;
 	
-	private final Object correlationLocksMonitor = new Object();
-
-	private final ConcurrentMap<Object, Object> locks = new ConcurrentHashMap<Object, Object>();
-
 	private volatile boolean sequenceAware = false;
 
+	private volatile LockRegistry lockRegistry = new DefaultLockRegistry();
+
+	private boolean lockRegistrySet = false;
+
 	public AbstractCorrelatingMessageHandler(MessageGroupProcessor processor, MessageGroupStore store,
 									 CorrelationStrategy correlationStrategy, ReleaseStrategy releaseStrategy) {
 		Assert.notNull(processor);
+
 		Assert.notNull(store);
 		setMessageStore(store);
 		this.outputProcessor = processor;
@@ -110,6 +112,12 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 		this(processor, new SimpleMessageStore(0), null, null);
 	}
 
+	public void setLockRegistry(LockRegistry lockRegistry) {
+		Assert.isTrue(!lockRegistrySet, "'this.lockRegistry' can not be reset once its been set");
+		Assert.notNull("'lockRegistry' must not be null");
+		this.lockRegistry = lockRegistry;
+		this.lockRegistrySet = true;
+	}
 
 	public void setMessageStore(MessageGroupStore store) {
 		this.messageStore = store;
@@ -183,9 +191,10 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 		}
 
 		// TODO: INT-1117 - make the lock global?
-		Object lock = getLock(correlationKey);
+		Lock lock = this.lockRegistry.obtain(correlationKey);
 
-		synchronized (lock) {
+		lock.lockInterruptibly();
+		try {
 			MessageGroup messageGroup = messageStore.getMessageGroup(correlationKey);
 			if (this.sequenceAware){
 				messageGroup = new SequenceAwareMessageGroup(messageGroup);
@@ -206,10 +215,6 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 						// Always clean up even if there was an exception
 						// processing messages						
 						this.afterRelease(messageGroup, completedMessages);
-		
-						synchronized(correlationLocksMonitor){
-							locks.remove(messageGroup.getGroupId());
-						}
 					}
 				} 				
 			} 
@@ -217,9 +222,11 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 				discardChannel.send(message);
 			}
 		}
+		finally {
+			lock.unlock();
+		}
 	}
 
-
 	/**
 	 * Allows you to provide additional logic that needs to be performed after the MessageGroup was released. 
 	 * @param group
@@ -230,32 +237,34 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 	private final boolean forceComplete(MessageGroup group) {
 
 		Object correlationKey = group.getGroupId();
-		Object lock = getLock(correlationKey);
-		synchronized (lock) {
-
-			if (group.size() > 0) {
-				try {
-					if (releaseStrategy.canRelease(group)) {
-						this.completeGroup(correlationKey, group);
-					} 
-					else {
-						this.expireGroup(correlationKey, group);
+		Lock lock = this.lockRegistry.obtain(correlationKey);
+		try {
+			lock.lockInterruptibly();
+			try {
+				if (group.size() > 0) {
+					try {
+						if (releaseStrategy.canRelease(group)) {
+							this.completeGroup(correlationKey, group);
+						}
+						else {
+							this.expireGroup(correlationKey, group);
+						}
 					}
+					finally {
+						this.remove(group);
+					}
+					return true;
 				}
-				finally {
-					this.remove(group);
-				}
-				return true;
 			}
-			return false;
+			finally  {
+				lock.unlock();
+			}
 		}
-	}
-
-	private Object getLock(Object correlationKey) {
-		synchronized(correlationLocksMonitor){
-			locks.putIfAbsent(correlationKey, new Object());
-			return locks.get(correlationKey);
+		catch (InterruptedException ie) {
+			Thread.currentThread().interrupt();
+			throw new MessagingException("Thread was interrupted while trying to obtain lock");
 		}
+		return false;
 	}
 
 	void remove(MessageGroup group) {
@@ -411,5 +420,4 @@ public abstract class AbstractCorrelatingMessageHandler extends AbstractMessageH
 			return false;
 		}
 	}
-
 }

spring-integration-core/src/main/java/org/springframework/integration/util/DefaultLockRegistry.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2002-2012 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.util;
+
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
+
+import org.springframework.util.Assert;
+/**
+ * Default implementation of {@link LockRegistry} which uses Masked Hashcode algorithm to obtain locks.
+ * When an instance of this class is created and array of {@link Lock} objects is created. The length of
+ * the array is based on the 'mask' parameter passed in the constructor. The default mask is 0xFF which will create
+ * and array consisting of 256 {@link ReentrantLock} instances.
+ * When the {@link #obtain(Object)} method is called with the lockKey (e.g., Object) the index of the {@link Lock}
+ * is determined by masking the object's hashCode (e.g., object.hashCode & mask) and the {@link Lock} is returned.
+ *
+ * @author Oleg Zhurakousky
+ * @author Gary Russell
+ * @since 2.1.1
+ *
+ */
+public final class DefaultLockRegistry implements LockRegistry {
+
+	private final Lock[] lockTable;
+
+	private final int mask;
+
+	/**
+     * Constructs a DefaultLockRegistry with the default
+     * mask 0xFF with 256 locks.
+     */
+	public DefaultLockRegistry(){
+		this(0xFF);
+	}
+
+	/**
+     * Constructs a DefaultLockRegistry with the supplied
+     * mask - the mask must have a value (2**n) - 1 where n
+     * is 1 to 31, creating a hash of 2**n locks.
+     * <p> Examples:
+     * <li>0x3ff (1023) - 1024 locks</li>
+     * <li>0xfff (4095) - 4096 locks</li>
+     * <p>
+     * @param mask
+     */
+	public DefaultLockRegistry(int mask){
+		String bits = Integer.toBinaryString(mask);
+		Assert.isTrue(bits.lastIndexOf('0') < bits.indexOf('1'), "Mask must be a power of 2 - 1");
+		this.mask = mask;
+		int arraySize = this.mask+1;
+		lockTable = new ReentrantLock[arraySize];
+		for (int i = 0; i < arraySize; i++) {
+			lockTable[i] = new ReentrantLock();
+		}
+	}
+
+	/**
+     * Obtains a lock by masking the lockKey's hashCode() with
+     * the mask and using the result as an index to the lock table.
+     * @param lockKey the object used to derive the lock index.
+     */
+	public Lock obtain(Object lockKey) {
+		Assert.notNull(lockKey, "'lockKey' must not be null");
+		Integer lockIndex = lockKey.hashCode() & this.mask;
+		return this.lockTable[lockIndex];
+	}
+}

spring-integration-core/src/main/java/org/springframework/integration/util/LockRegistry.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2002-2012 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.util;
+
+import java.util.concurrent.locks.Lock;
+
+/**
+ * Strategy for maintaining a registry of shared locks
+ *
+ * @author Oleg Zhurakousky
+ * @author Gary Russell
+ * @since 2.1.1
+ */
+public interface LockRegistry {
+
+	/**
+     * Obtains the lock associated with the parameter object.
+     * @param lockRoot The object with which the lock is associated.
+     * @return The associated lock.
+     */
+	Lock obtain(Object lockKey);
+
+}

spring-integration-core/src/test/java/org/springframework/integration/aggregator/CorrelatingMessageHandlerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2010 the original author or authors.
+ * Copyright 2002-2012 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.
@@ -18,7 +18,6 @@ package org.springframework.integration.aggregator;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 
@@ -28,6 +27,7 @@ import org.junit.runner.RunWith;
 import org.mockito.Mock;
 import org.mockito.internal.stubbing.answers.ThrowsException;
 import org.mockito.runners.MockitoJUnitRunner;
+
 import org.springframework.integration.Message;
 import org.springframework.integration.MessageChannel;
 import org.springframework.integration.MessageHandlingException;
@@ -35,7 +35,6 @@ import org.springframework.integration.store.MessageGroupStore;
 import org.springframework.integration.store.SimpleMessageGroup;
 import org.springframework.integration.store.SimpleMessageStore;
 import org.springframework.integration.support.MessageBuilder;
-import org.springframework.test.util.ReflectionTestUtils;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
@@ -84,20 +83,14 @@ public class CorrelatingMessageHandlerTests {
 		when(correlationStrategy.getCorrelationKey(isA(Message.class))).thenReturn(correlationKey);
 
 		handler.handleMessage(message1);
-		verifyLocks(handler, 1);
 
 		handler.handleMessage(message2);
-		verifyLocks(handler, 0); // lock is removed when group is complete
 
 		verify(correlationStrategy).getCorrelationKey(message1);
 		verify(correlationStrategy).getCorrelationKey(message2);
 		verify(processor).processMessageGroup(isA(SimpleMessageGroup.class));
 	}
 
-	private void verifyLocks(AggregatingMessageHandler handler, int lockCount) {
-		assertEquals(lockCount, ((Map<?, ?>) ReflectionTestUtils.getField(handler, "locks")).size());
-	}
-
 	@Test
 	public void bufferCompletesWithException() throws Exception {
 

spring-integration-core/src/test/java/org/springframework/integration/aggregator/scenarios/AggregatorWithCustomReleaseStrategyTests.java

@@ -34,38 +34,96 @@ import static org.junit.Assert.assertTrue;
 public class AggregatorWithCustomReleaseStrategyTests {
 
 	@Test
-	public void validateSequenceSizeHasNoAffect() throws Exception{
+	public void testAggregatorsUnderStressWithConcurrency() throws Exception{
+		// this is to be sure  after INT-2502
+		for (int i = 0; i < 10; i++) {
+			this.validateSequenceSizeHasNoAffectCustomCorrelator();
+		}
+		for (int i = 0; i < 10; i++) {
+			this.validateSequenceSizeHasNoAffectWithSplitter();
+		}
+	}
+
+	public void validateSequenceSizeHasNoAffectCustomCorrelator() throws Exception{
+		ApplicationContext context =
+				new ClassPathXmlApplicationContext("aggregator-with-custom-release-strategy.xml", this.getClass());
+		final MessageChannel inputChannel = context.getBean("aggregationChannelCustomCorrelation", MessageChannel.class);
+		QueueChannel resultChannel = context.getBean("resultChannel", QueueChannel.class);
+
+		final CountDownLatch latch = new CountDownLatch(1800);
+
+		for (int i = 0; i < 600; i++) {
+			final int counter = i;
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload("foo").
+							setHeader("correlation", "foo"+counter).build());
+					latch.countDown();
+				}
+			}).start();
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload("bar").
+							setHeader("correlation", "foo"+counter).build());
+					latch.countDown();
+				}
+			}).start();
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload("baz").
+							setHeader("correlation", "foo"+counter).build());
+					latch.countDown();
+				}
+			}).start();
+		}
+
+		assertTrue("Sends failed to complete", latch.await(10, TimeUnit.SECONDS));
+
+		Message<?> message = resultChannel.receive(10);
+		int counter = 0;
+		while(message != null){
+			counter++;
+			message = resultChannel.receive(10);
+		}
+		assertEquals(600, counter);
+	}
+
+	public void validateSequenceSizeHasNoAffectWithSplitter() throws Exception{
 		ApplicationContext context =
 				new ClassPathXmlApplicationContext("aggregator-with-custom-release-strategy.xml", this.getClass());
 		final MessageChannel inputChannel = context.getBean("in", MessageChannel.class);
 		QueueChannel resultChannel = context.getBean("resultChannel", QueueChannel.class);
 
-		final CountDownLatch latch = new CountDownLatch(2);
+		final CountDownLatch latch = new CountDownLatch(1800);
 
-		new Thread(new Runnable() {
-			public void run() {
-				inputChannel.send(MessageBuilder.withPayload(new Integer[]{1, 2, 3, 4, 5, 6, 7, 8}).
-						setHeader("correlation", "foo").build());
-				latch.countDown();
-			}
-		}).start();
-
-		new Thread(new Runnable() {
-			public void run() {
-				inputChannel.send(MessageBuilder.withPayload(new Integer[]{10, 20, 30, 40, 50, 60, 70, 80}).
-						setHeader("correlation", "foo").build());
-				latch.countDown();
-			}
-		}).start();
+		for (int i = 0; i < 600; i++) {
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload(new Integer[]{1, 2, 3, 4, 5, 6, 7, 8}).build());
+					latch.countDown();
+				}
+			}).start();
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload(new Integer[]{9, 10, 11, 12, 13, 14, 15, 16}).build());
+					latch.countDown();
+				}
+			}).start();
+			new Thread(new Runnable() {
+				public void run() {
+					inputChannel.send(MessageBuilder.withPayload(new Integer[]{17, 18, 19, 20, 21, 22, 23, 24}).build());
+					latch.countDown();
+				}
+			}).start();
+		}
 
 		assertTrue("Sends failed to complete", latch.await(10, TimeUnit.SECONDS));
 
-		Message<?> message = resultChannel.receive(1000);
+		Message<?> message = resultChannel.receive(10);
 		int counter = 0;
-		while(message != null){
-			counter++;
-			message = resultChannel.receive(1000);
+		while(message != null && ++counter < 7200){
+			message = resultChannel.receive(10);
 		}
-		assertEquals(8, counter);
+		assertEquals(7200, counter);
 	}
 }

spring-integration-core/src/test/java/org/springframework/integration/aggregator/scenarios/aggregator-with-custom-release-strategy.xml

@@ -6,9 +6,13 @@
 		http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd">
 
 
-	<int:splitter input-channel="in" output-channel="aggregationChannel"/>
+	<int:splitter input-channel="in" output-channel="aggregationChannelFromSplitter"/>
 
-	<int:aggregator input-channel="aggregationChannel" output-channel="resultChannel"
+	<int:aggregator input-channel="aggregationChannelFromSplitter" output-channel="resultChannel"
+			release-strategy-expression="size() == 2"
+			expire-groups-upon-completion="true"/>
+
+	<int:aggregator input-channel="aggregationChannelCustomCorrelation" output-channel="resultChannel"
 			release-strategy-expression="size() == 2"
 			correlation-strategy-expression="headers.correlation"
 			expire-groups-upon-completion="true"/>
@@ -16,4 +20,4 @@
 	<int:channel id="resultChannel">
 		<int:queue/>
 	</int:channel>
-</beans>
+</beans>

spring-integration-core/src/test/java/org/springframework/integration/util/DefaultLockRegistryTests.java

@@ -0,0 +1,167 @@
+/*
+ * Copyright 2002-2012 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.util;
+
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+
+import java.util.concurrent.locks.Lock;
+
+import org.junit.Test;
+
+/**
+ * @author Gary Russell
+ * @since 2.1.1
+ *
+ */
+public class DefaultLockRegistryTests {
+
+    @Test(expected=IllegalArgumentException.class)
+    public void testBadMask() {
+        new DefaultLockRegistry(4);
+    }
+
+    @Test
+    public void testSame() {
+        LockRegistry registry = new DefaultLockRegistry();
+        Lock lock1 = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 0;
+            }});
+        Lock lock2 = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 256;
+            }});
+        assertSame(lock1, lock2);
+    }
+
+    @Test
+    public void testDifferent() {
+        LockRegistry registry = new DefaultLockRegistry();
+        Lock lock1 = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 0;
+            }});
+        Lock lock2 = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 255;
+            }});
+        assertNotSame(lock1, lock2);
+    }
+
+    @Test
+    public void testAllDifferentAndSame() {
+        LockRegistry registry = new DefaultLockRegistry(3);
+        Lock[] locks = new Lock[4];
+        locks[0] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 0;
+            }});
+        locks[1] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 1;
+            }});
+        locks[2] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 2;
+            }});
+        locks[3] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 3;
+            }});
+        for (int i = 0; i < 4; i++) {
+            for (int j = 1; j < 4; j++) {
+                if (i != j) {
+                    assertNotSame(locks[i], locks[j]);
+                }
+            }
+        }
+        Lock[] moreLocks = new Lock[4];
+        moreLocks[0] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 0;
+            }});
+        moreLocks[1] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 1;
+            }});
+        moreLocks[2] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 2;
+            }});
+        moreLocks[3] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 3;
+            }});
+        assertSame(locks[0], moreLocks[0]);
+        assertSame(locks[1], moreLocks[1]);
+        assertSame(locks[2], moreLocks[2]);
+        assertSame(locks[3], moreLocks[3]);
+        moreLocks[0] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 4;
+            }});
+        moreLocks[1] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 5;
+            }});
+        moreLocks[2] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 6;
+            }});
+        moreLocks[3] = registry.obtain(new Object() {
+
+            @Override
+            public int hashCode() {
+                return 7;
+            }});
+        assertSame(locks[0], moreLocks[0]);
+        assertSame(locks[1], moreLocks[1]);
+        assertSame(locks[2], moreLocks[2]);
+        assertSame(locks[3], moreLocks[3]);
+    }
+
+}