support per-record observations in batch listeners

Signed-off-by: Igor Macedo Quintanilha <[email protected]>

Author: Igor Macedo Quintanilha

spring-kafka-docs/src/main/antora/modules/ROOT/pages/micrometer.adoc

@@ -0,0 +1,60 @@
+= Micrometer Support
+
+Spring for Apache Kafka provides support for Micrometer.
+
+[[listener-obs]]
+== Listener Observations
+
+When you have `micrometer-observation` on the class path, and you configure a
+`ObservationRegistry` in the application context, the container will automatically create
+`Observation`s for consumed records.
+You can disable this feature by setting the `observationEnabled` container property to
+`false`.
+
+The default `KafkaListenerObservationConvention` is `DefaultKafkaListenerObservationConvention`
+which provides the following observation names and key/values.
+
+Listener observation names:
+
+* `spring.kafka.listener`
+
+Key Values:
+
+* `spring.kafka.listener.id` - the listener id.
+* `spring.kafka.group.id` - the consumer group id.
+* `spring.kafka.client.id` - the consumer client id.
+* `spring.kafka.cluster.id` - the cluster id.
+
+When using a record listener (`@KafkaListener` method), a single observation is
+created for each record.
+
+[[batch-listener-obs]]
+=== Batch Listener Observations
+
+When using a batch listener, by default, no observations are created, even if a
+`ObservationRegistry` is present.
+This is because the scope of an observation is tied to the thread, and with a batch
+listener, there is no one-to-one mapping between an observation and a record.
+
+To enable per-record observations in a batch listener, set the container factory
+property `recordObservationsInBatch` to `true`.
+
+[source,java]
+----
+@Bean
+ConcurrentKafkaListenerContainerFactory<?, ?> kafkaListenerContainerFactory(
+        ConcurrentKafkaListenerContainerFactoryConfigurer configurer,
+        ConsumerFactory<Object, Object> kafkaConsumerFactory) {
+
+    ConcurrentKafkaListenerContainerFactory<Object, Object> factory = new ConcurrentKafkaListenerContainerFactory<>();
+    configurer.configure(factory, kafkaConsumerFactory);
+    factory.getContainerProperties().setRecordObservationsInBatch(true);
+    return factory;
+}
+----
+
+When this property is `true`, an observation will be created for each record in the
+batch, similar to a record listener.
+This allows you to have visibility into the processing of each record, even within a
+batch context.
+No spans will be created for the batch itself, only for each record. 

spring-kafka-docs/src/main/antora/modules/ROOT/pages/whats-new.adoc

@@ -76,3 +76,10 @@ For details, see xref:kafka/receiving-messages/rebalance-listeners.adoc#new-reba
 
 The `DefaultKafkaHeaderMapper` and `SimpleKafkaHeaderMapper` support multi-value header mapping for Kafka records.
 More details are available in xref:kafka/headers.adoc#multi-value-header[Support multi-value header mapping].
+
+[[x40-batch-observability]]
+=== Per-Record Observation in Batch Listeners
+
+It is now possible to get an observation for each record when using a batch listener.
+Previously, only a single observation was created for the entire batch.
+See xref:micrometer.adoc#batch-listener-obs[Observability for Batch Listeners] for more information.

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

@@ -310,6 +310,8 @@ public enum EOSMode {
 
 	private boolean restartAfterAuthExceptions;
 
+	private boolean recordObservationsInBatch;
+
 	/**
 	 * Create properties for a container that will subscribe to the specified topics.
 	 * @param topics the topics.
@@ -1091,6 +1093,27 @@ public void setRestartAfterAuthExceptions(boolean restartAfterAuthExceptions) {
 		this.restartAfterAuthExceptions = restartAfterAuthExceptions;
 	}
 
+	/**
+	 * When true, and a batch listener is configured with observation enabled, an observation
+	 * will be started for each record in the batch.
+	 * @return recordObservationsInBatch.
+	 * @since 4.0
+	 */
+	public boolean isRecordObservationsInBatch() {
+		return this.recordObservationsInBatch;
+	}
+
+	/**
+	 * Set whether to enable individual record observations in a batch.
+	 * When true, and a batch listener is configured with observation enabled, an observation
+	 * will be started for each record in the batch. Default false.
+	 * @param recordObservationsInBatch true to enable individual record observations.
+	 * @since 4.0
+	 */
+	public void setRecordObservationsInBatch(boolean recordObservationsInBatch) {
+		this.recordObservationsInBatch = recordObservationsInBatch;
+	}
+
 	@Override
 	public String toString() {
 		return "ContainerProperties ["
@@ -1141,6 +1164,7 @@ public String toString() {
 						? "\n observationRegistry=" + this.observationRegistry
 						: "")
 				+ "\n restartAfterAuthExceptions=" + this.restartAfterAuthExceptions
+				+ "\n recordObservationsInBatch=" + this.recordObservationsInBatch
 				+ "\n]";
 	}
 

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

@@ -898,7 +898,7 @@ private final class ListenerConsumer implements SchedulingAwareRunnable, Consume
 				this.isBatchListener = true;
 				this.wantsFullRecords = this.batchListener.wantsPollResult();
 				this.pollThreadStateProcessor = setUpPollProcessor(true);
-				this.observationEnabled = false;
+				this.observationEnabled =  this.containerProperties.isObservationEnabled() && this.containerProperties.isRecordObservationsInBatch();
 			}
 			else if (listener instanceof MessageListener) {
 				this.listener = (MessageListener<K, V>) listener;
@@ -2423,6 +2423,21 @@ private void ackBatch(final ConsumerRecords<K, V> records) throws InterruptedExc
 			}
 		}
 
+		private void invokeBatchWithIndividualRecordObservation(List<ConsumerRecord<K, V>> recordList) {
+			// Create individual observations for each record without scopes
+			for (ConsumerRecord<K, V> record : recordList) {
+				Observation observation = KafkaListenerObservation.LISTENER_OBSERVATION.observation(
+						this.containerProperties.getObservationConvention(),
+							DefaultKafkaListenerObservationConvention.INSTANCE,
+							() -> new KafkaRecordReceiverContext(record, getListenerId(), getClientId(), this.consumerGroupId,
+									this::clusterId),
+							this.observationRegistry);
+				observation.observe(() -> {
+					this.logger.debug(() -> "Observing record in batch: " + KafkaUtils.format(record));
+				});
+			}
+		}
+
 		private void invokeBatchOnMessageWithRecordsOrList(final ConsumerRecords<K, V> recordsArg,
 				List<ConsumerRecord<K, V>> recordListArg) {
 
@@ -2443,7 +2458,13 @@ private void invokeBatchOnMessageWithRecordsOrList(final ConsumerRecords<K, V> r
 				}
 			}
 			Object sample = startMicrometerSample();
+
+
 			try {
+				if (this.observationEnabled) {
+					invokeBatchWithIndividualRecordObservation(recordList);
+				}
+
 				if (this.wantsFullRecords) {
 					Objects.requireNonNull(this.batchListener).onMessage(records, // NOSONAR
 							this.isAnyManualAck