AMQP-829: RabbitTemplate Req/Resp Fail Fast

JIRA: https://jira.spring.io/browse/AMQP-829

If the channel closes while we are waiting for a reply, throw an
exception instead of waiting to time out.

* Polishing - PR Comments

Author: garyrussell

spring-rabbit/src/main/java/org/springframework/amqp/rabbit/core/RabbitTemplate.java

@@ -25,8 +25,6 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
-import java.util.concurrent.ArrayBlockingQueue;
-import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -1601,6 +1599,8 @@ public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProp
 
 			};
 			ClosingRecoveryListener.addRecoveryListenerIfNecessary(channel);
+			ShutdownListener shutdownListener = c -> pendingReply.completeExceptionally(c);
+			channel.addShutdownListener(shutdownListener);
 			channel.basicConsume(replyTo, true, consumerTag, this.noLocalReplyConsumer, true, null, consumer);
 			Message reply = null;
 			try {
@@ -1612,20 +1612,26 @@ public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProp
 				if (channel.isOpen()) {
 					cancelConsumerQuietly(channel, consumer);
 				}
+				try {
+					channel.removeShutdownListener(shutdownListener);
+				}
+				catch (Exception e) {
+					// NOSONAR - channel might have closed.
+				}
 			}
 			return reply;
 		}, obtainTargetConnectionFactory(this.sendConnectionFactorySelectorExpression, message));
 	}
 
 	private void cancelConsumerQuietly(Channel channel, DefaultConsumer consumer) {
 		try {
-            channel.basicCancel(consumer.getConsumerTag());
-        }
-        catch (Exception e) {
+			channel.basicCancel(consumer.getConsumerTag());
+		}
+		catch (Exception e) {
 			if (this.logger.isDebugEnabled()) {
 				this.logger.debug("Failed to cancel consumer: " + consumer, e);
 			}
-        }
+		}
 	}
 
 	protected Message doSendAndReceiveWithFixed(final String exchange, final String routingKey, final Message message,
@@ -2331,7 +2337,7 @@ private static class PendingReply {
 
 		private volatile String savedCorrelation;
 
-		private final BlockingQueue<Object> queue = new ArrayBlockingQueue<Object>(1);
+		private final CompletableFuture<Message> future = new CompletableFuture<>();
 
 		public String getSavedReplyTo() {
 			return this.savedReplyTo;
@@ -2350,33 +2356,36 @@ public void setSavedCorrelation(String savedCorrelation) {
 		}
 
 		public Message get() throws InterruptedException {
-			Object reply = this.queue.take();
-			return processReply(reply);
+			try {
+				return this.future.get();
+			}
+			catch (ExecutionException e) {
+				throw RabbitExceptionTranslator.convertRabbitAccessException(e.getCause());
+			}
 		}
 
 		public Message get(long timeout, TimeUnit unit) throws InterruptedException {
-			Object reply = this.queue.poll(timeout, unit);
-			return reply == null ? null : processReply(reply);
-		}
-
-		private Message processReply(Object reply) {
-			if (reply instanceof Message) {
-				return (Message) reply;
+			try {
+				return this.future.get(timeout, unit);
 			}
-			else if (reply instanceof AmqpException) {
-				throw (AmqpException) reply;
+			catch (ExecutionException e) {
+				throw RabbitExceptionTranslator.convertRabbitAccessException(e.getCause());
 			}
-			else {
-				throw new AmqpException("Unexpected reply type " + reply.getClass().getName());
+			catch (TimeoutException e) {
+				return null;
 			}
 		}
 
 		public void reply(Message reply) {
-			this.queue.add(reply);
+			this.future.complete(reply);
 		}
 
 		public void returned(AmqpMessageReturnedException e) {
-			this.queue.add(e);
+			completeExceptionally(e);
+		}
+
+		public void completeExceptionally(Throwable t) {
+			this.future.completeExceptionally(t);
 		}
 
 	}

spring-rabbit/src/test/java/org/springframework/amqp/rabbit/core/RabbitTemplateTests.java

@@ -18,10 +18,12 @@
 
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.instanceOf;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertSame;
 import static org.junit.Assert.assertThat;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -36,7 +38,10 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.CountDownLatch;
 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.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
@@ -47,6 +52,7 @@
 import org.mockito.Mockito;
 
 import org.springframework.amqp.AmqpAuthenticationException;
+import org.springframework.amqp.AmqpException;
 import org.springframework.amqp.core.Address;
 import org.springframework.amqp.core.Message;
 import org.springframework.amqp.core.MessageProperties;
@@ -77,6 +83,8 @@
 import com.rabbitmq.client.ConnectionFactory;
 import com.rabbitmq.client.Consumer;
 import com.rabbitmq.client.Envelope;
+import com.rabbitmq.client.ShutdownListener;
+import com.rabbitmq.client.ShutdownSignalException;
 import com.rabbitmq.client.impl.AMQImpl;
 import com.rabbitmq.client.impl.AMQImpl.Queue.DeclareOk;
 
@@ -365,6 +373,48 @@ public void testNestedTxBinding() throws Exception {
 		verify(channel1).txCommit();
 	}
 
+	@Test
+	public void testShutdownWhileWaitingForReply() throws Exception {
+		ConnectionFactory mockConnectionFactory = mock(ConnectionFactory.class);
+		Connection mockConnection = mock(Connection.class);
+		Channel mockChannel = mock(Channel.class);
+
+		given(mockConnectionFactory.newConnection(any(ExecutorService.class), anyString())).willReturn(mockConnection);
+		given(mockConnection.isOpen()).willReturn(true);
+		given(mockConnection.createChannel()).willReturn(mockChannel);
+		given(mockChannel.queueDeclare()).willReturn(new AMQImpl.Queue.DeclareOk("foo", 0, 0));
+		final AtomicReference<ShutdownListener> listener = new AtomicReference<>();
+		final CountDownLatch shutdownLatch = new CountDownLatch(1);
+		willAnswer(invocation -> {
+			listener.set(invocation.getArgument(0));
+			shutdownLatch.countDown();
+			return null;
+		}).given(mockChannel).addShutdownListener(any());
+		SingleConnectionFactory connectionFactory = new SingleConnectionFactory(mockConnectionFactory);
+		connectionFactory.setExecutor(mock(ExecutorService.class));
+		RabbitTemplate template = new RabbitTemplate(connectionFactory);
+		template.setReplyTimeout(60_000);
+		Message input = new Message("Hello, world!".getBytes(), new MessageProperties());
+		ExecutorService exec = Executors.newSingleThreadExecutor();
+		exec.execute(() -> {
+			try {
+				shutdownLatch.await(10, TimeUnit.SECONDS);
+			}
+			catch (InterruptedException e) {
+				Thread.currentThread().interrupt();
+			}
+			listener.get().shutdownCompleted(new ShutdownSignalException(true, false, null, null));
+		});
+		try {
+			template.doSendAndReceiveWithTemporary("foo", "bar", input, null);
+			fail("Expected exception");
+		}
+		catch (AmqpException e) {
+			assertThat(e.getCause(), instanceOf(ShutdownSignalException.class));
+		}
+		exec.shutdownNow();
+	}
+
 	@SuppressWarnings("serial")
 	private class TestTransactionManager extends AbstractPlatformTransactionManager {