GH-1672: Option to Immediately Stop the Container

Resolves #1672

Previously (and still, by default), stopping the listener container does
not take effect until the records from the previous poll are all processed.

Add an option to stop after the current record, instead.

**cherry-pick to 2.6.x, 1.5.x**

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/listener/ContainerProperties.java

@@ -255,6 +255,8 @@ public enum EOSMode {
 
 	private boolean stopContainerWhenFenced;
 
+	private boolean stopImmediate;
+
 	/**
 	 * Create properties for a container that will subscribe to the specified topics.
 	 * @param topics the topics.
@@ -744,6 +746,26 @@ public void setStopContainerWhenFenced(boolean stopContainerWhenFenced) {
 		this.stopContainerWhenFenced = stopContainerWhenFenced;
 	}
 
+	/**
+	 * When true, the container will be stopped immediately after processing the current record.
+	 * @return true to stop immediately.
+	 * @since 2.5.11
+	 */
+	public boolean isStopImmediate() {
+		return this.stopImmediate;
+	}
+
+	/**
+	 * Set to true to stop the container after processing the current record (when stop()
+	 * is called). When false (default), the container will stop after all the results of
+	 * the previous poll are processed.
+	 * @param stopImmediate true to stop after the current record.
+	 * @since 2.5.11
+	 */
+	public void setStopImmediate(boolean stopImmediate) {
+		this.stopImmediate = stopImmediate;
+	}
+
 	private void adviseListenerIfNeeded() {
 		if (!CollectionUtils.isEmpty(this.adviceChain)) {
 			if (AopUtils.isAopProxy(this.messageListener)) {

spring-kafka/src/main/java/org/springframework/kafka/listener/KafkaMessageListenerContainer.java

@@ -590,6 +590,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final boolean fixTxOffsets = this.containerProperties.isFixTxOffsets();
 
+		private final boolean stopImmediate = this.containerProperties.isStopImmediate();
+
 		private Map<TopicPartition, OffsetMetadata> definedPartitions;
 
 		private int count;
@@ -1814,6 +1816,9 @@ private void invokeRecordListener(final ConsumerRecords<K, V> records) {
 		private void invokeRecordListenerInTx(final ConsumerRecords<K, V> records) {
 			Iterator<ConsumerRecord<K, V>> iterator = records.iterator();
 			while (iterator.hasNext()) {
+				if (this.stopImmediate && !isRunning()) {
+					break;
+				}
 				final ConsumerRecord<K, V> record = checkEarlyIntercept(iterator.next());
 				if (record == null) {
 					continue;
@@ -1903,6 +1908,9 @@ protected void doInTransactionWithoutResult(TransactionStatus status) {
 		private void doInvokeWithRecords(final ConsumerRecords<K, V> records) {
 			Iterator<ConsumerRecord<K, V>> iterator = records.iterator();
 			while (iterator.hasNext()) {
+				if (this.stopImmediate && !isRunning()) {
+					break;
+				}
 				final ConsumerRecord<K, V> record = checkEarlyIntercept(iterator.next());
 				if (record == null) {
 					continue;

spring-kafka/src/test/java/org/springframework/kafka/listener/KafkaMessageListenerContainerTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2016-2020 the original author or authors.
+ * Copyright 2016-2021 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.
@@ -3139,6 +3139,53 @@ void commitAfterHandleManual() throws InterruptedException {
 		verify(consumer).commitSync(any(), any());
 	}
 
+	@Test
+	@SuppressWarnings({ "unchecked", "rawtypes" })
+	void stopImmediately() throws InterruptedException {
+		ConsumerFactory<Integer, String> cf = mock(ConsumerFactory.class);
+		Consumer<Integer, String> consumer = mock(Consumer.class);
+		given(cf.createConsumer(eq("grp"), eq("clientId"), isNull(), any())).willReturn(consumer);
+		Map<String, Object> cfProps = new HashMap<>();
+		cfProps.put(ConsumerConfig.DEFAULT_API_TIMEOUT_MS_CONFIG, 45000); // wins
+		given(cf.getConfigurationProperties()).willReturn(cfProps);
+		final Map<TopicPartition, List<ConsumerRecord<Integer, String>>> records =
+				Map.of(new TopicPartition("foo", 0), Arrays.asList(new ConsumerRecord<>("foo", 0, 0L, 1, "foo"),
+						new ConsumerRecord<>("foo", 0, 1L, 1, "bar")));
+		ConsumerRecords<Integer, String> consumerRecords = new ConsumerRecords<>(records);
+		ConsumerRecords<Integer, String> emptyRecords = new ConsumerRecords<>(Collections.emptyMap());
+		AtomicBoolean first = new AtomicBoolean(true);
+		given(consumer.poll(any(Duration.class))).willAnswer(i -> {
+			Thread.sleep(50);
+			return first.getAndSet(false) ? consumerRecords : emptyRecords;
+		});
+		TopicPartitionOffset[] topicPartition = new TopicPartitionOffset[] {
+				new TopicPartitionOffset("foo", 0) };
+		ContainerProperties containerProps = new ContainerProperties(topicPartition);
+		containerProps.setGroupId("grp");
+		containerProps.setClientId("clientId");
+		containerProps.setStopImmediate(true);
+		AtomicInteger delivered = new AtomicInteger();
+		AtomicReference<KafkaMessageListenerContainer> containerRef = new AtomicReference<>();
+		containerProps.setMessageListener((MessageListener) r -> {
+			delivered.incrementAndGet();
+			containerRef.get().stop(() -> { });
+		});
+		KafkaMessageListenerContainer<Integer, String> container =
+				new KafkaMessageListenerContainer<>(cf, containerProps);
+		containerRef.set(container);
+		CountDownLatch latch = new CountDownLatch(1);
+		container.setApplicationEventPublisher(event -> {
+			if (event instanceof ConsumerStoppedEvent) {
+				latch.countDown();
+			}
+		});
+		container.start();
+		assertThat(latch.await(10, TimeUnit.SECONDS)).isTrue();
+		container.stop();
+		assertThat(delivered.get()).isEqualTo(1);
+		verify(consumer).commitSync(eq(Map.of(new TopicPartition("foo", 0), new OffsetAndMetadata(1L))), any());
+	}
+
 	private Consumer<?, ?> spyOnConsumer(KafkaMessageListenerContainer<Integer, String> container) {
 		Consumer<?, ?> consumer =
 				KafkaTestUtils.getPropertyValue(container, "listenerConsumer.consumer", Consumer.class);

src/reference/asciidoc/kafka.adoc

@@ -2433,6 +2433,10 @@ Metadata
 |Stop the listener container if a `ProducerFencedException` is thrown.
 See <<after-rollback>> for more information.
 
+|stopImmediate
+|`false`
+|When the container is stopped, stop processing after the current record instead of after processing all the records from the previous poll.
+
 |subBatchPerPartition
 |See desc.
 |When using a batch listener, if this is `true`, the listener is called with the results of the poll split into sub batches, one per partition.

src/reference/asciidoc/whats-new.adoc

@@ -12,4 +12,7 @@ This version requires the 2.7.0 `kafka-clients`.
 ==== Listener Container Changes
 
 The `onlyLogRecordMetadata` container property is now `true` by default.
+
+A new container property `stopImmediate` is now available.
+
 See <<container-props>> for more information.