Add getAssignmentsByClientId() to containers

* Remove redundant `final` from `private` method

Co-authored-by: Artem Bilan <abilan@pivotal.io>

Author: garyrussell

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

@@ -109,6 +109,20 @@ public Collection<TopicPartition> getAssignedPartitions() {
 		}
 	}
 
+	@Override
+	public Map<String, Collection<TopicPartition>> getAssignmentsByClientId() {
+		synchronized (this.lifecycleMonitor) {
+			Map<String, Collection<TopicPartition>> assignments = new HashMap<>();
+			this.containers.forEach(container -> {
+				Map<String, Collection<TopicPartition>> byClientId = container.getAssignmentsByClientId();
+				if (byClientId != null) {
+					assignments.putAll(byClientId);
+				}
+			});
+			return assignments;
+		}
+	}
+
 	@Override
 	public boolean isContainerPaused() {
 		synchronized (this.lifecycleMonitor) {

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

@@ -245,6 +245,18 @@ else if (partitionsListenerConsumer.assignedPartitions != null) {
 		}
 	}
 
+	@Override
+	@Nullable
+	public Map<String, Collection<TopicPartition>> getAssignmentsByClientId() {
+		ListenerConsumer partitionsListenerConsumer = this.listenerConsumer;
+		if (this.listenerConsumer != null) {
+			return Collections.singletonMap(partitionsListenerConsumer.getClientId(), getAssignedPartitions());
+		}
+		else {
+			return null;
+		}
+	}
+
 	@Override
 	public boolean isContainerPaused() {
 		return isPaused() && this.listenerConsumer != null && this.listenerConsumer.isConsumerPaused();
@@ -255,11 +267,7 @@ public boolean isContainerPaused() {
 		ListenerConsumer listenerConsumerForMetrics = this.listenerConsumer;
 		if (listenerConsumerForMetrics != null) {
 			Map<MetricName, ? extends Metric> metrics = listenerConsumerForMetrics.consumer.metrics();
-			Iterator<MetricName> metricIterator = metrics.keySet().iterator();
-			if (metricIterator.hasNext()) {
-				String clientId = metricIterator.next().tags().get("client-id");
-				return Collections.singletonMap(clientId, metrics);
-			}
+			return Collections.singletonMap(listenerConsumerForMetrics.getClientId(), metrics);
 		}
 		return Collections.emptyMap();
 	}
@@ -541,6 +549,8 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 
 		private final Map<TopicPartition, OffsetAndMetadata> commitsDuringRebalance = new HashMap<>();
 
+		private final String clientId;
+
 		private Map<TopicPartition, OffsetMetadata> definedPartitions;
 
 		private int count;
@@ -585,6 +595,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 							KafkaMessageListenerContainer.this.clientIdSuffix,
 							consumerProperties);
 
+			this.clientId = determineClientId();
 			this.transactionTemplate = determineTransactionTemplate();
 			this.genericListener = listener;
 			this.consumerSeekAwareListener = checkConsumerSeekAware(listener);
@@ -658,6 +669,19 @@ else if (listener instanceof MessageListener) {
 			this.subBatchPerPartition = setupSubBatchPerPartition();
 		}
 
+		String getClientId() {
+			return this.clientId;
+		}
+
+		private String determineClientId() {
+			Map<MetricName, ? extends Metric> metrics = this.consumer.metrics();
+			Iterator<MetricName> metricIterator = metrics.keySet().iterator();
+			if (metricIterator.hasNext()) {
+				return metricIterator.next().tags().get("client-id");
+			}
+			return "unknown.client.id";
+		}
+
 		private void checkGroupInstance(Properties properties, ConsumerFactory<K, V> consumerFactory) {
 			String groupInstance = properties.getProperty(ConsumerConfig.GROUP_INSTANCE_ID_CONFIG);
 			if (!StringUtils.hasText(groupInstance)) {

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

@@ -65,10 +65,21 @@ default ContainerProperties getContainerProperties() {
 	 * @return the topics/partitions.
 	 * @since 2.1.3
 	 */
+	@Nullable
 	default Collection<TopicPartition> getAssignedPartitions() {
 		throw new UnsupportedOperationException("This container doesn't support retrieving its assigned partitions");
 	}
 
+	/**
+	 * Return the assigned topics/partitions for this container, by client.id.
+	 * @return the topics/partitions.
+	 * @since 2.5
+	 */
+	@Nullable
+	default Map<String, Collection<TopicPartition>> getAssignmentsByClientId() {
+		throw new UnsupportedOperationException("This container doesn't support retrieving its assigned partitions");
+	}
+
 	/**
 	 * Pause this container before the next poll().
 	 * @since 2.1.3

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

@@ -132,6 +132,7 @@ protected Consumer<Integer, String> createKafkaConsumer(String groupId, String c
 		};
 		ContainerProperties containerProps = new ContainerProperties(topic1);
 		containerProps.setLogContainerConfig(true);
+		containerProps.setClientId("client");
 
 		final CountDownLatch latch = new CountDownLatch(3);
 		final Set<String> listenerThreadNames = new ConcurrentSkipListSet<>();
@@ -164,6 +165,10 @@ protected Consumer<Integer, String> createKafkaConsumer(String groupId, String c
 
 		ContainerTestUtils.waitForAssignment(container, embeddedKafka.getPartitionsPerTopic());
 		assertThat(container.getAssignedPartitions()).hasSize(2);
+		Map<String, Collection<TopicPartition>> assignments = container.getAssignmentsByClientId();
+		assertThat(assignments).hasSize(2);
+		assertThat(assignments.get("client-0")).isNotNull();
+		assertThat(assignments.get("client-1")).isNotNull();
 
 		Map<String, Object> senderProps = KafkaTestUtils.producerProps(embeddedKafka);
 		ProducerFactory<Integer, String> pf = new DefaultKafkaProducerFactory<>(senderProps);

src/reference/asciidoc/kafka.adoc

@@ -2299,8 +2299,8 @@ BeforeTx
 
 |pause
 Requested
-|n/a
-|True if a consumer pause has been requested (read only).
+|(read only)
+|True if a consumer pause has been requested.
 
 |record
 Interceptor
@@ -2323,8 +2323,14 @@ Timeout
 
 |assigned
 Partitions
-|n/a
-|The partitions currently assigned to this container (explicitly or not) (read only).
+|(read only)
+|The partitions currently assigned to this container (explicitly or not).
+
+|assigned
+Partitions
+ByClientId
+|(read only)
+|The partitions currently assigned to this container (explicitly or not).
 
 |clientId
 Suffix
@@ -2345,8 +2351,14 @@ Suffix
 
 |assigned
 Partitions
-|n/a
-|The aggregate of partitions currently assigned to this container's child `KafkaMessageListenerContainer` s (explicitly or not) (read only).
+|(read only)
+|The aggregate of partitions currently assigned to this container's child `KafkaMessageListenerContainer` s (explicitly or not).
+
+|assigned
+Partitions
+ByClientId
+|(read only)
+|The partitions currently assigned to this container's child `KafkaMessageListenerContainer` s (explicitly or not), keyed by the child container's consumer's `client.id` property.
 
 |concurrency
 |1

src/reference/asciidoc/whats-new.adoc

@@ -61,6 +61,9 @@ See <<error-handlers>> for more information.
 You can now control the level at which exceptions intentionally thrown by standard error handlers are logged.
 See <<error-handlers>> for more information.
 
+The `getAssignmentsByClientId()` method has been added, making it easier to determine which consumers in a concurrent container are assigned which partition(s).
+See <<container-props>> for more information.
+
 [[x25-template]]
 ==== KafkaTemplate Changes