INT-2502-backport Improve Locking, Race Condition

INT-2502-backport Polishing

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

@@ -13,8 +13,12 @@
 
 package org.springframework.integration.aggregator;
 
+import java.util.Collection;
+import java.util.concurrent.locks.Lock;
+
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.springframework.beans.factory.BeanFactory;
 import org.springframework.integration.Message;
 import org.springframework.integration.MessageChannel;
@@ -24,13 +28,15 @@ import org.springframework.integration.channel.NullChannel;
 import org.springframework.integration.core.MessageProducer;
 import org.springframework.integration.core.MessagingTemplate;
 import org.springframework.integration.handler.AbstractMessageHandler;
-import org.springframework.integration.store.*;
+import org.springframework.integration.store.MessageGroup;
+import org.springframework.integration.store.MessageGroupCallback;
+import org.springframework.integration.store.MessageGroupStore;
+import org.springframework.integration.store.MessageStore;
+import org.springframework.integration.store.SimpleMessageStore;
+import org.springframework.integration.util.DefaultLockRegistry;
+import org.springframework.integration.util.LockRegistry;
 import org.springframework.util.Assert;
 
-import java.util.Collection;
-import java.util.concurrent.ConcurrentHashMap;
-import java.util.concurrent.ConcurrentMap;
-
 /**
  * Message handler that holds a buffer of correlated messages in a
  * {@link MessageStore}. This class takes care of correlated groups of messages
@@ -73,9 +79,9 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 
 	private boolean sendPartialResultOnExpiry = false;
 	
-	private final Object correlationLocksMonitor = new Object();
+	private volatile LockRegistry lockRegistry = new DefaultLockRegistry();
 
-	private final ConcurrentMap<Object, Object> locks = new ConcurrentHashMap<Object, Object>();
+	private boolean lockRegistrySet = false;
 
 
 	public CorrelatingMessageHandler(MessageGroupProcessor processor, MessageGroupStore store,
@@ -98,6 +104,12 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		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;
@@ -168,9 +180,9 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 		}
 
 		// TODO: INT-1117 - make the lock global?
-		Object lock = getLock(correlationKey);
-
-		synchronized (lock) {
+		Lock lock = this.lockRegistry.obtain(correlationKey);
+		lock.lockInterruptibly();
+		try {
 			MessageGroup group = messageStore.getMessageGroup(correlationKey);
 			if (group.canAdd(message)) {
 				if (logger.isTraceEnabled()) {
@@ -203,6 +215,9 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 				discardChannel.send(message);
 			}
 		}
+		finally {
+			lock.unlock();
+		}
 	}
 
 	@SuppressWarnings("rawtypes")
@@ -223,33 +238,32 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 	}
 
 	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)) {
-						completeGroup(correlationKey, group);
-					} else {
-						expireGroup(group, correlationKey);
+		Lock lock = this.lockRegistry.obtain(correlationKey);
+		try {
+			lock.lockInterruptibly();
+			try {
+				if (group.size() > 0) {
+					try {
+						if (releaseStrategy.canRelease(group)) {
+							completeGroup(correlationKey, group);
+						} else {
+							expireGroup(group, correlationKey);
+						}
 					}
+					finally {
+						remove(group);
+					}
+					return true;
 				}
-				finally {
-					remove(group);
-				}
-				return true;
+			} finally {
+				lock.unlock();
 			}
-			return false;
-		}
-	}
-
-	private Object getLock(Object correlationKey) {
-		synchronized(correlationLocksMonitor){
-			locks.putIfAbsent(correlationKey, new Object());
-			return locks.get(correlationKey);
+		} catch (InterruptedException e) {
+			Thread.currentThread().interrupt();
+			throw new MessagingException("Thread was interrupted while trying to obtain lock");
 		}
+		return false;
 	}
 
 	private void mark(MessageGroup group) {
@@ -267,9 +281,6 @@ public class CorrelatingMessageHandler extends AbstractMessageHandler implements
 	private void remove(MessageGroup group) {
 		Object correlationKey = group.getGroupId();
 		messageStore.removeMessageGroup(correlationKey);
-		synchronized(correlationLocksMonitor){
-			locks.remove(correlationKey);
-		}
 	}
 
 	private MessageGroup store(Object correlationKey, Message<?> message) {

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.0.6
+ *
+ */
+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 Math.pow(2, n) - 1 where n
+	 * is 1 to 31, creating a hash of Math.pow(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.length() < 32 && (mask == 0 || 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.0.6
+ */
+public interface LockRegistry {
+
+	/**
+	 * Obtains the lock associated with the parameter object.
+	 * @param lockKey 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

@@ -25,7 +25,6 @@ import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 
@@ -35,7 +34,6 @@ 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;
@@ -43,7 +41,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;
 
 /**
  * @author Iwein Fuld
@@ -84,20 +81,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(CorrelatingMessageHandler 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

@@ -0,0 +1,129 @@
+/*
+ * 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.aggregator.scenarios;
+
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+
+import org.junit.Test;
+
+import org.springframework.context.ApplicationContext;
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+import org.springframework.integration.Message;
+import org.springframework.integration.MessageChannel;
+import org.springframework.integration.channel.QueueChannel;
+import org.springframework.integration.support.MessageBuilder;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+/**
+ * @author Oleg Zhurakousky
+ *
+ */
+public class AggregatorWithCustomReleaseStrategyTests {
+
+	@Test
+	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(1800);
+
+		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(10);
+		int counter = 0;
+		while(message != null && ++counter < 7200){
+			message = resultChannel.receive(10);
+		}
+		assertEquals(7200, counter);
+	}
+}

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

@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xmlns:int="http://www.springframework.org/schema/integration"
+	xsi:schemaLocation="http://www.springframework.org/schema/integration http://www.springframework.org/schema/integration/spring-integration-2.0.xsd
+		http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd">
+
+
+	<int:splitter input-channel="in" output-channel="aggregationChannelFromSplitter"/>
+
+	<int:aggregator input-channel="aggregationChannelFromSplitter" output-channel="resultChannel"
+			release-strategy-expression="size() == 2"/>
+
+	<int:aggregator input-channel="aggregationChannelCustomCorrelation" output-channel="resultChannel"
+			release-strategy-expression="size() == 2"
+			correlation-strategy-expression="headers.correlation"/>
+
+	<int:channel id="resultChannel">
+		<int:queue/>
+	</int:channel>
+</beans>

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

@@ -0,0 +1,184 @@
+/*
+ * 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 org.junit.Test;
+
+import static org.junit.Assert.assertNotSame;
+import static org.junit.Assert.assertSame;
+
+/**
+ * @author Gary Russell
+ * @author Oleg Zhurakousky
+ * @since 2.0.6
+ *
+ */
+public class DefaultLockRegistryTests {
+
+	@Test(expected=IllegalArgumentException.class)
+	public void testBadMask() {
+		new DefaultLockRegistry(4);
+	}
+
+	@Test(expected=IllegalArgumentException.class)
+	public void testBadMaskOutOfRange() {// 32bits
+		new DefaultLockRegistry(0xffffffff);
+	}
+
+	@Test
+	public void testSingleLockCreation() {
+		LockRegistry registry = new DefaultLockRegistry(0);
+		Lock a = registry.obtain(23);
+		Lock b = registry.obtain(new Object());
+		Lock c = registry.obtain("hello");
+		assertSame(a, b);
+		assertSame(a, c);
+		assertSame(b, c);
+	}
+
+	@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]);
+	}
+
+}