AMQP-621: Fix Consumer.basicCancel Race Condition

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

The `SimpleMessageListenerContainer` has logic to initiate `BlockingQueueConsumer.basicCancel()` and mark `BlockingQueueConsumer` as inactive immediately.
But the `handleCancelOk` answer can come back from the Broker a bit late, in async manner.
During this time window we can still receive messages via `handleDelivery`.
With `prefetchCount = 1` we may end up with the deadlock on `handleDelivery` thread during that time window, when main consumer loop identifies consumer as inactive and exits for `stop` process
and can't do that because the `BlockingQueueConsumer.this.queue.put(new Delivery(consumerTag, envelope, properties, body))` operation blocks the consumer from any other operations.

* Fix the race condition via additional `BlockingQueueConsumer.cancelled()` state to let main loop to spin until `handleCancelOk` honoring any in-flight messages
* The `cancelled` state is based on the `cancelReceived` flag which is changed from the `handleCancelOk` as a Broker response for the `basicCancel`
and in the `handleCancel` event from the Broker as a reason of some unexpected state on Broker, e.g. consuming queue has been removed
* Protect message lost via `basicReject` in case of late arrival into `handleDelivery`, when the consumer has been cancelled already

Address PR comments

* Rename `cancelReceived` -> `cancelled` with JavaDocs
* Initiate `basicCancel()` from `handleCancel(String consumerTag)` to cancel `Consumer` from all other tags.
The Broker canceling initiative is fully similar to `queueChanged` in the listener container
* Do no mark `cancelled` in the `handleCancelOk` until all the tags are cancelled

Polishing

1. Fix case when basicCancel received with only one queue (set cancelled).
2. Add protection to cancelled() in case all cancelOks are not received within shutdownTimeout
   after aborting.
3. If we are aborting, use offer instead of put to insert into the queue to avoid indefinitely
   blocking the client thread.
4. Purge the queue at the end of stop() just in case #2 happens and we still have a hung client thread.

Add Test - Abort with Multiple Queues

Author: artembilan

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

@@ -97,10 +97,12 @@ public class BlockingQueueConsumer {
 
 	private InternalConsumer consumer;
 
+	/**
+	 * The flag indicating that consumer has been cancelled from all queues
+	 * via {@code handleCancelOk} callback replies.
+	 */
 	private final AtomicBoolean cancelled = new AtomicBoolean(false);
 
-	private final AtomicBoolean cancelReceived = new AtomicBoolean(false);
-
 	private final AcknowledgeMode acknowledgeMode;
 
 	private final ConnectionFactory connectionFactory;
@@ -133,6 +135,10 @@ public class BlockingQueueConsumer {
 
 	private BackOffExecution backOffExecution;
 
+	private long shutdownTimeout;
+
+	private volatile long abortStarted;
+
 	/**
 	 * Create a consumer. The consumer must not attempt to use
 	 * the connection factory or communicate with the broker
@@ -249,6 +255,10 @@ public String getConsumerTag() {
 	public final void setQuiesce(long shutdownTimeout) {
 	}
 
+	public void setShutdownTimeout(long shutdownTimeout) {
+		this.shutdownTimeout = shutdownTimeout;
+	}
+
 	/**
 	 * Set the number of retries after passive queue declaration fails.
 	 * @param declarationRetries The number of retries, default 3.
@@ -318,14 +328,18 @@ protected void basicCancel() {
 				break;
 			}
 		}
-		this.consumerTags.clear();
-		this.cancelled.set(true);
+		this.abortStarted = System.currentTimeMillis();
 	}
 
 	protected boolean hasDelivery() {
 		return !this.queue.isEmpty();
 	}
 
+	protected boolean cancelled() {
+		return this.cancelled.get() || (this.abortStarted > 0 &&
+				this.abortStarted + this.shutdownTimeout > System.currentTimeMillis());
+	}
+
 	/**
 	 * Check if we are in shutdown mode and if so throw an exception.
 	 */
@@ -339,7 +353,6 @@ private void checkShutdown() {
 	 * If this is a non-POISON non-null delivery simply return it.
 	 * If this is POISON we are in shutdown mode, throw
 	 * shutdown. If delivery is null, we may be in shutdown mode. Check and see.
-	 * @throws InterruptedException
 	 */
 	private Message handle(Delivery delivery) throws InterruptedException {
 		if ((delivery == null && this.shutdown != null)) {
@@ -391,7 +404,7 @@ public Message nextMessage(long timeout) throws InterruptedException, ShutdownSi
 			checkMissingQueues();
 		}
 		Message message = handle(this.queue.poll(timeout, TimeUnit.MILLISECONDS));
-		if (message == null && this.cancelReceived.get()) {
+		if (message == null && this.cancelled.get()) {
 			throw new ConsumerCancelledException();
 		}
 		return message;
@@ -582,9 +595,8 @@ private void attemptPassiveDeclarations() {
 	}
 
 	public void stop() {
-		this.cancelled.set(true);
 		if (this.consumer != null && this.consumer.getChannel() != null && this.consumerTags.size() > 0
-				&& !this.cancelReceived.get()) {
+				&& !this.cancelled.get()) {
 			try {
 				RabbitUtils.closeMessageConsumer(this.consumer.getChannel(), this.consumerTags.keySet(),
 						this.transactional);
@@ -602,6 +614,7 @@ public void stop() {
 		ConnectionFactoryUtils.releaseResources(this.resourceHolder);
 		this.deliveryTags.clear();
 		this.consumer = null;
+		this.queue.clear(); // in case we still have a client thread blocked
 	}
 
 	/**
@@ -738,19 +751,29 @@ public void handleShutdownSignal(String consumerTag, ShutdownSignalException sig
 		@Override
 		public void handleCancel(String consumerTag) throws IOException {
 			if (logger.isWarnEnabled()) {
-				logger.warn("Cancel received for " + consumerTag + "; " + BlockingQueueConsumer.this);
+				logger.warn("Cancel received for " + consumerTag + " ("
+						+ BlockingQueueConsumer.this.consumerTags.get(consumerTag)
+						+ "); " + BlockingQueueConsumer.this);
 			}
 			BlockingQueueConsumer.this.consumerTags.remove(consumerTag);
-			BlockingQueueConsumer.this.cancelReceived.set(true);
+			if (BlockingQueueConsumer.this.consumerTags.isEmpty()) {
+				BlockingQueueConsumer.this.cancelled.set(true);
+			}
+			else {
+				basicCancel();
+			}
 		}
 
 		@Override
 		public void handleCancelOk(String consumerTag) {
 			if (logger.isDebugEnabled()) {
-				logger.debug("Received cancellation notice for tag " + consumerTag + "; " + BlockingQueueConsumer.this);
+				logger.debug("Received cancelOk for tag " + consumerTag + " ("
+						+ BlockingQueueConsumer.this.consumerTags.get(consumerTag)
+						+ "); " + BlockingQueueConsumer.this);
 			}
-			synchronized (BlockingQueueConsumer.this.consumerTags) {
-				BlockingQueueConsumer.this.consumerTags.remove(consumerTag);
+			BlockingQueueConsumer.this.consumerTags.remove(consumerTag);
+			if (BlockingQueueConsumer.this.consumerTags.isEmpty()) {
+				BlockingQueueConsumer.this.cancelled.set(true);
 			}
 		}
 
@@ -761,7 +784,13 @@ public void handleDelivery(String consumerTag, Envelope envelope, AMQP.BasicProp
 				logger.debug("Storing delivery for " + BlockingQueueConsumer.this);
 			}
 			try {
-				BlockingQueueConsumer.this.queue.put(new Delivery(consumerTag, envelope, properties, body));
+				if (BlockingQueueConsumer.this.abortStarted > 0) {
+					BlockingQueueConsumer.this.queue.offer(new Delivery(consumerTag, envelope, properties, body),
+							BlockingQueueConsumer.this.shutdownTimeout, TimeUnit.MILLISECONDS);
+				}
+				else {
+					BlockingQueueConsumer.this.queue.put(new Delivery(consumerTag, envelope, properties, body));
+				}
 			}
 			catch (InterruptedException e) {
 				Thread.currentThread().interrupt();

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

@@ -1040,6 +1040,7 @@ protected BlockingQueueConsumer createBlockingQueueConsumer() {
 			consumer.setTagStrategy(this.consumerTagStrategy);
 		}
 		consumer.setBackOffExecution(this.recoveryBackOff.start());
+		consumer.setShutdownTimeout(this.shutdownTimeout);
 		return consumer;
 	}
 
@@ -1306,7 +1307,7 @@ public void run() {
 					ConsumerChannelRegistry.registerConsumerChannel(this.consumer.getChannel(), getConnectionFactory());
 				}
 
-				while (isActive(this.consumer) || this.consumer.hasDelivery()) {
+				while (isActive(this.consumer) || this.consumer.hasDelivery() || !this.consumer.cancelled()) {
 					try {
 						boolean receivedOk = receiveAndExecute(this.consumer); // At least one message received
 						if (SimpleMessageListenerContainer.this.maxConcurrentConsumers != null) {

spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/ContainerInitializationTests.java

@@ -46,8 +46,12 @@
  */
 public class ContainerInitializationTests {
 
+	private static final String TEST_MISMATCH = "test.mismatch";
+
+	private static final String TEST_MISMATCH2 = "test.mismatch2";
+
 	@Rule
-	public BrokerRunning brokerRunning = BrokerRunning.isRunningWithEmptyQueues("test.mismatch");
+	public BrokerRunning brokerRunning = BrokerRunning.isRunningWithEmptyQueues(TEST_MISMATCH, TEST_MISMATCH2);
 
 	@After
 	public void tearDown() {
@@ -84,8 +88,23 @@ public void testMismatchedQueue() {
 	public void testMismatchedQueueDuringRestart() throws Exception {
 		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(Config2.class);
 		RabbitAdmin admin = context.getBean(RabbitAdmin.class);
-		admin.deleteQueue("test.mismatch");
-		admin.declareQueue(new Queue("test.mismatch", false, false, true));
+		admin.deleteQueue(TEST_MISMATCH);
+		admin.declareQueue(new Queue(TEST_MISMATCH, false, false, true));
+		SimpleMessageListenerContainer container = context.getBean(SimpleMessageListenerContainer.class);
+		int n = 0;
+		while (n++ < 100 && container.isRunning()) {
+			Thread.sleep(100);
+		}
+		assertFalse(container.isRunning());
+		context.close();
+	}
+
+	@Test
+	public void testMismatchedQueueDuringRestartMultiQueue() throws Exception {
+		AnnotationConfigApplicationContext context = new AnnotationConfigApplicationContext(Config3.class);
+		RabbitAdmin admin = context.getBean(RabbitAdmin.class);
+		admin.deleteQueue(TEST_MISMATCH);
+		admin.declareQueue(new Queue(TEST_MISMATCH, false, false, true));
 		SimpleMessageListenerContainer container = context.getBean(SimpleMessageListenerContainer.class);
 		int n = 0;
 		while (n++ < 100 && container.isRunning()) {
@@ -106,7 +125,7 @@ public ConnectionFactory connectionFactory() {
 		@Bean
 		public SimpleMessageListenerContainer container() {
 			SimpleMessageListenerContainer container = new SimpleMessageListenerContainer(connectionFactory());
-			container.setQueueNames("test.mismatch");
+			container.setQueueNames(TEST_MISMATCH);
 			container.setMessageListener(new MessageListenerAdapter(new Object() {
 
 				@SuppressWarnings("unused")
@@ -120,7 +139,7 @@ public void handleMessage(Message m) {
 
 		@Bean
 		public Queue queue() {
-			return new Queue("test.mismatch", false, false, true); // mismatched
+			return new Queue(TEST_MISMATCH, false, false, true); // mismatched
 		}
 
 	}
@@ -141,7 +160,25 @@ static class Config2 extends Config1 {
 		@Override
 		@Bean
 		public Queue queue() {
-			return new Queue("test.mismatch", true, false, false);
+			return new Queue(TEST_MISMATCH, true, false, false);
+		}
+
+	}
+
+	@Configuration
+	static class Config3 extends Config2 {
+
+
+		@Override
+		public SimpleMessageListenerContainer container() {
+			SimpleMessageListenerContainer container = super.container();
+			container.setQueueNames(TEST_MISMATCH, TEST_MISMATCH2);
+			return container;
+		}
+
+		@Bean
+		public Queue queue2() {
+			return new Queue(TEST_MISMATCH2, true, false, false);
 		}
 
 	}

spring-rabbit/src/test/java/org/springframework/amqp/rabbit/listener/SimpleMessageListenerContainerIntegration2Tests.java

@@ -154,6 +154,7 @@ public void testChangeQueues() throws Exception {
 	public void testDeleteOneQueue() throws Exception {
 		CountDownLatch latch = new CountDownLatch(20);
 		container = createContainer(new MessageListenerAdapter(new PojoListener(latch)), queue.getName(), queue1.getName());
+		container.setFailedDeclarationRetryInterval(100);
 		ApplicationEventPublisher publisher = mock(ApplicationEventPublisher.class);
 		container.setApplicationEventPublisher(publisher);
 		for (int i = 0; i < 10; i++) {