GH-1091: Fix DMLC with Routing Connection Factory

Fixes #1091

The DMLC can open a connection on the calling thread (e.g. `start()`, `setConsumersPerQueue()`).

When using a routing connection factory, the factory key is temporarily bound to the thread for
proper CF selection in connection listeners (e.g. `RabbitAdmin`).

This is disallowed if the calling thread is, itself, a listener container thread.

Add `push/pop` operations to `SimpleResourceHolder`.

**cherry-pick to 2.1.x**

Author: garyrussell

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/connection/SimpleResourceHolder.java

@@ -17,7 +17,9 @@
 package org.springframework.amqp.rabbit.connection;
 
 import java.util.Collections;
+import java.util.Deque;
 import java.util.HashMap;
+import java.util.LinkedList;
 import java.util.Map;
 
 import org.apache.commons.logging.Log;
@@ -28,15 +30,18 @@
 import org.springframework.util.Assert;
 
 /**
- * Central helper that manages resources per thread to be used by resource management code.
- *
- * <p>Supports one resource per key without overwriting, that is, a resource needs
- * to be removed before a new one can be set for the same key.
- *
- * <p>Resource management code should check for thread-bound resources via {@link #has(Object)}.
- *
- * <p>This helper isn't designed for transaction synchronization cases.
- * Use {@code TransactionSynchronizationManager} and {@code ResourceHolder} instead.
+ * Central helper that manages resources per thread to be used by resource management
+ * code.
+ * <p>
+ * {@link #bind(Object, Object)} supports one resource per key without overwriting, that
+ * is, a resource needs to be removed before a new one can be set for the same key. But
+ * see {@link #push(Object, Object)} and {@link #pop(Object)}.
+ * <p>
+ * Resource management code should check for thread-bound resources via
+ * {@link #has(Object)}.
+ * <p>
+ * This helper isn't designed for transaction synchronization cases. Use
+ * {@code TransactionSynchronizationManager} and {@code ResourceHolder} instead.
  *
  * @author Artem Bilan
  * @author Gary Russell
@@ -48,11 +53,14 @@ public final class SimpleResourceHolder {
 
 	private static final String BOUND_TO_THREAD = "] bound to thread [";
 
-	private static final Log logger = LogFactory.getLog(SimpleResourceHolder.class); // NOSONAR lower case
+	private static final Log LOGGER = LogFactory.getLog(SimpleResourceHolder.class);
 
-	private static final ThreadLocal<Map<Object, Object>> resources = // NOSONAR lower case
+	private static final ThreadLocal<Map<Object, Object>> RESOURCES =
 			new NamedThreadLocal<Map<Object, Object>>("Simple resources");
 
+	private static final ThreadLocal<Map<Object, Deque<Object>>> STACK =
+			new NamedThreadLocal<Map<Object, Deque<Object>>>("Simple resources");
+
 	/**
 	 * Return all resources that are bound to the current thread.
 	 * <p>Mainly for debugging purposes. Resource managers should always invoke
@@ -63,7 +71,7 @@ public final class SimpleResourceHolder {
 	 * @see #has
 	 */
 	public static Map<Object, Object> getResources() {
-		Map<Object, Object> map = resources.get();
+		Map<Object, Object> map = RESOURCES.get();
 		return (map != null ? Collections.unmodifiableMap(map) : Collections.emptyMap());
 	}
 
@@ -86,8 +94,8 @@ public static boolean has(Object key) {
 	@Nullable
 	public static Object get(Object key) {
 		Object value = doGet(key);
-		if (value != null && logger.isTraceEnabled()) {
-			logger.trace("Retrieved value [" + value + FOR_KEY + key + BOUND_TO_THREAD
+		if (value != null && LOGGER.isTraceEnabled()) {
+			LOGGER.trace("Retrieved value [" + value + FOR_KEY + key + BOUND_TO_THREAD
 					+ Thread.currentThread().getName() + "]");
 		}
 		return value;
@@ -100,7 +108,7 @@ public static Object get(Object key) {
 	 */
 	@Nullable
 	private static Object doGet(Object actualKey) {
-		Map<Object, Object> map = resources.get();
+		Map<Object, Object> map = RESOURCES.get();
 		if (map == null) {
 			return null;
 		}
@@ -115,22 +123,71 @@ private static Object doGet(Object actualKey) {
 	 */
 	public static void bind(Object key, Object value) {
 		Assert.notNull(value, "Value must not be null");
-		Map<Object, Object> map = resources.get();
+		Map<Object, Object> map = RESOURCES.get();
 		// set ThreadLocal Map if none found
 		if (map == null) {
 			map = new HashMap<Object, Object>();
-			resources.set(map);
+			RESOURCES.set(map);
 		}
 		Object oldValue = map.put(key, value);
 		Assert.isNull(oldValue, () -> "Already value [" + oldValue + FOR_KEY + key + BOUND_TO_THREAD
 				+ Thread.currentThread().getName() + "]");
 
-		if (logger.isTraceEnabled()) {
-			logger.trace(
+		if (LOGGER.isTraceEnabled()) {
+			LOGGER.trace(
 					"Bound value [" + value + FOR_KEY + key + "] to thread [" + Thread.currentThread().getName() + "]");
 		}
 	}
 
+	/**
+	 * Set the value for this key and push any existing value onto a stack.
+	 * @param key the key.
+	 * @param value the value.
+	 * @since 2.1.11
+	 */
+	public static void push(Object key, Object value) {
+		Object currentValue = get(key);
+		if (currentValue == null) {
+			bind(key, value);
+		}
+		else {
+			Map<Object, Deque<Object>> stack = STACK.get();
+			if (stack == null) {
+				stack = new HashMap<>();
+				STACK.set(stack);
+			}
+			stack.computeIfAbsent(key, k -> new LinkedList<>());
+			stack.get(key).push(currentValue);
+			unbind(key);
+			bind(key, value);
+		}
+	}
+
+	/**
+	 * Unbind the current value and bind the head of the stack if present.
+	 * @param key the key.
+	 * @return the popped value.
+	 * @since 2.1.11
+	 */
+	@Nullable
+	public static Object pop(Object key) {
+		Object popped = unbind(key);
+		Map<Object, Deque<Object>> stack = STACK.get();
+		if (stack != null) {
+			Deque<Object> deque = stack.get(key);
+			if (deque != null && deque.size() > 0) {
+				Object previousValue = deque.pop();
+				if (previousValue != null) {
+					bind(key, previousValue);
+				}
+				if (deque.isEmpty()) {
+					STACK.remove();
+				}
+			}
+		}
+		return popped;
+	}
+
 	/**
 	 * Unbind a resource for the given key from the current thread.
 	 * @param key the key to unbind (usually the resource factory)
@@ -151,18 +208,18 @@ public static Object unbind(Object key) throws IllegalStateException {
 	 */
 	@Nullable
 	public static Object unbindIfPossible(Object key) {
-		Map<Object, Object> map = resources.get();
+		Map<Object, Object> map = RESOURCES.get();
 		if (map == null) {
 			return null;
 		}
 		Object value = map.remove(key);
 		// Remove entire ThreadLocal if empty...
 		if (map.isEmpty()) {
-			resources.remove();
+			RESOURCES.remove();
 		}
 
-		if (value != null && logger.isTraceEnabled()) {
-			logger.trace("Removed value [" + value + FOR_KEY + key + "] from thread ["
+		if (value != null && LOGGER.isTraceEnabled()) {
+			LOGGER.trace("Removed value [" + value + FOR_KEY + key + "] from thread ["
 					+ Thread.currentThread().getName() + "]");
 		}
 		return value;
@@ -172,10 +229,12 @@ public static Object unbindIfPossible(Object key) {
 	 * Clear resources for the current thread.
 	 */
 	public static void clear() {
-		resources.remove();
+		RESOURCES.remove();
+		STACK.remove();
 	}
 
 	private SimpleResourceHolder() {
+		super();
 	}
 
 }

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/listener/DirectMessageListenerContainer.java

@@ -584,7 +584,7 @@ private void startConsumers(final String[] queueNames) {
 	protected void doRedeclareElementsIfNecessary() {
 		String routingLookupKey = getRoutingLookupKey();
 		if (routingLookupKey != null) {
-			SimpleResourceHolder.bind(getRoutingConnectionFactory(), routingLookupKey); // NOSONAR both never null here
+			SimpleResourceHolder.push(getRoutingConnectionFactory(), routingLookupKey); // NOSONAR both never null here
 		}
 		try {
 			redeclareElementsIfNecessary();
@@ -594,7 +594,7 @@ protected void doRedeclareElementsIfNecessary() {
 		}
 		finally {
 			if (routingLookupKey != null) {
-				SimpleResourceHolder.unbind(getRoutingConnectionFactory()); // NOSONAR never null here
+				SimpleResourceHolder.pop(getRoutingConnectionFactory()); // NOSONAR never null here
 			}
 		}
 	}
@@ -660,7 +660,7 @@ private void doConsumeFromQueue(String queue) {
 		}
 		String routingLookupKey = getRoutingLookupKey();
 		if (routingLookupKey != null) {
-			SimpleResourceHolder.bind(getRoutingConnectionFactory(), routingLookupKey); // NOSONAR both never null here
+			SimpleResourceHolder.push(getRoutingConnectionFactory(), routingLookupKey); // NOSONAR both never null here
 		}
 		Connection connection = null; // NOSONAR (close)
 		try {
@@ -675,7 +675,7 @@ private void doConsumeFromQueue(String queue) {
 		}
 		finally {
 			if (routingLookupKey != null) {
-				SimpleResourceHolder.unbind(getRoutingConnectionFactory()); // NOSONAR never null here
+				SimpleResourceHolder.pop(getRoutingConnectionFactory()); // NOSONAR never null here
 			}
 		}
 		SimpleConsumer consumer = consume(queue, connection);

spring-rabbit/src/test/java/org/springframework/amqp/rabbit/connection/RoutingConnectionFactoryTests.java

@@ -41,6 +41,8 @@
 import org.mockito.Mockito;
 
 import org.springframework.amqp.rabbit.listener.DirectMessageListenerContainer;
+import org.springframework.amqp.rabbit.listener.DirectReplyToMessageListenerContainer;
+import org.springframework.amqp.rabbit.listener.DirectReplyToMessageListenerContainer.ChannelHolder;
 import org.springframework.amqp.rabbit.listener.SimpleMessageListenerContainer;
 
 import com.rabbitmq.client.Channel;
@@ -276,4 +278,49 @@ protected synchronized void redeclareElementsIfNecessary() {
 		assertThat(connectionMakerKey2.get()).isEqualTo("xxx[foo]");
 	}
 
+	@Test
+	public void testWithDRTDMLCAndConnectionListenerExistingRFK() throws Exception {
+		ConnectionFactory connectionFactory1 = mock(ConnectionFactory.class);
+		Map<Object, ConnectionFactory> factories = new HashMap<Object, ConnectionFactory>(2);
+		factories.put("xxx[foo]", connectionFactory1);
+		factories.put("xxx[amq.rabbitmq.reply-to]", connectionFactory1);
+
+		final SimpleRoutingConnectionFactory connectionFactory = new SimpleRoutingConnectionFactory();
+		SimpleResourceHolder.bind(connectionFactory, "foo");
+
+		final Connection connection = mock(Connection.class);
+		Channel channel = mock(Channel.class);
+		given(channel.isOpen()).willReturn(true);
+		given(connection.createChannel(anyBoolean())).willReturn(channel);
+		final AtomicReference<Object> connectionMakerKey = new AtomicReference<>();
+		final CountDownLatch latch = new CountDownLatch(2); // early connection in Abstract container
+		willAnswer(i -> {
+			connectionMakerKey.set(connectionFactory.determineCurrentLookupKey());
+			latch.countDown();
+			return connection;
+		}).given(connectionFactory1).createConnection();
+		connectionFactory.setTargetConnectionFactories(factories);
+
+		final AtomicReference<Object> connectionMakerKey2 = new AtomicReference<>();
+		DirectReplyToMessageListenerContainer container = new DirectReplyToMessageListenerContainer(connectionFactory) {
+
+			@Override
+			protected synchronized void redeclareElementsIfNecessary() {
+				connectionMakerKey2.set(connectionFactory.determineCurrentLookupKey());
+			}
+
+		};
+		container.setLookupKeyQualifier("xxx");
+		container.setShutdownTimeout(10);
+		container.afterPropertiesSet();
+		container.start();
+		ChannelHolder channelHolder = container.getChannelHolder();
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		container.releaseConsumerFor(channelHolder, true, "test");
+		container.stop();
+		assertThat(connectionMakerKey.get()).isEqualTo("xxx[amq.rabbitmq.reply-to]");
+		assertThat(connectionMakerKey2.get()).isEqualTo("xxx[amq.rabbitmq.reply-to]");
+		assertThat(SimpleResourceHolder.unbind(connectionFactory)).isEqualTo("foo");
+	}
+
 }