support per-record observations in batch listeners

document recordObservationsInBatch container property
describe new option in the change history
add integration test for per-record observations

Author: Igor Macedo Quintanilha

spring-kafka-docs/src/main/antora/modules/ROOT/pages/appendix/change-history.adoc

@@ -1,6 +1,23 @@
 [[history]]
 = Change History
 
+[[what-s-new-in-3-4-since-3-3]]
+== What's New in 3.4 Since 3.3
+:page-section-summary-toc: 1
+
+This section covers the changes made from version 3.3 to version 3.4.
+For changes in earlier version, see xref:appendix/change-history.adoc[Change History].
+
+[[x34-record-observations-in-batch]]
+=== Individual Record Observations in Batch Processing
+
+A new `ContainerProperties` property `recordObservationsInBatch` has been introduced to enable individual record observations within batch listeners.
+When enabled, and a batch listener is configured with observation enabled, an observation will be started for each record in the batch.
+This allows for proper distributed tracing of individual records within a batch, even when each record originated from different producers with different trace contexts.
+The default value is `false` to maintain backward compatibility.
+
+See xref:kafka/micrometer.adoc#observation[Micrometer Observation] for more details.
+
 [[what-s-new-in-3-3-since-3-2]]
 == What's New in 3.3 Since 3.2
 :page-section-summary-toc: 1

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 3.4
+	 */
+	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 3.4
+	 */
+	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;
@@ -2443,25 +2443,32 @@ private void invokeBatchOnMessageWithRecordsOrList(final ConsumerRecords<K, V> r
 				}
 			}
 			Object sample = startMicrometerSample();
-			try {
-				if (this.wantsFullRecords) {
-					Objects.requireNonNull(this.batchListener).onMessage(records, // NOSONAR
-							this.isAnyManualAck
-									? new ConsumerBatchAcknowledgment(records, recordList)
-									: null,
-							this.consumer);
+			
+			// Handle individual record tracing for batch mode if enabled
+			if (this.containerProperties.isRecordObservationsInBatch() && this.observationEnabled) {
+				invokeBatchWithIndividualRecordTracing(records, recordList, sample);
+			}
+			else {
+				try {
+					if (this.wantsFullRecords) {
+						Objects.requireNonNull(this.batchListener).onMessage(records, // NOSONAR
+								this.isAnyManualAck
+										? new ConsumerBatchAcknowledgment(records, recordList)
+										: null,
+								this.consumer);
+					}
+					else {
+						doInvokeBatchOnMessage(records, recordList); // NOSONAR
+					}
+					batchInterceptAfter(records, null);
+					successTimer(sample, null);
 				}
-				else {
-					doInvokeBatchOnMessage(records, recordList); // NOSONAR
+				catch (RuntimeException e) {
+					this.batchFailed = true;
+					failureTimer(sample, null, e);
+					batchInterceptAfter(records, e);
+					throw e;
 				}
-				batchInterceptAfter(records, null);
-				successTimer(sample, null);
-			}
-			catch (RuntimeException e) {
-				this.batchFailed = true;
-				failureTimer(sample, null, e);
-				batchInterceptAfter(records, e);
-				throw e;
 			}
 		}
 
@@ -4005,6 +4012,69 @@ private Long computeBackwardWhereTo(long offset, boolean toCurrent, TopicPartiti
 
 		}
 
+		private void invokeBatchWithIndividualRecordTracing(final ConsumerRecords<K, V> records,
+				List<ConsumerRecord<K, V>> recordList, @Nullable Object sample) {
+			
+			List<Observation> observations = new ArrayList<>();
+			List<Observation.Scope> scopes = new ArrayList<>();
+			
+			try {
+				// Create individual observations for each record
+				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.start();
+					observations.add(observation);
+					scopes.add(observation.openScope());
+				}
+				
+				// Invoke the batch listener
+				if (this.wantsFullRecords) {
+					Objects.requireNonNull(this.batchListener).onMessage(records,
+							this.isAnyManualAck
+									? new ConsumerBatchAcknowledgment(records, recordList)
+									: null,
+							this.consumer);
+				}
+				else {
+					doInvokeBatchOnMessage(records, recordList);
+				}
+				
+				batchInterceptAfter(records, null);
+				successTimer(sample, null);
+			}
+			catch (RuntimeException e) {
+				this.batchFailed = true;
+				failureTimer(sample, null, e);
+				batchInterceptAfter(records, e);
+				
+				// Mark all observations with error
+				for (Observation observation : observations) {
+					if (!isListenerAdapterObservationAware()) {
+						observation.error(e);
+					}
+				}
+				throw e;
+			}
+			finally {
+				// Close scopes in reverse order
+				for (int i = scopes.size() - 1; i >= 0; i--) {
+					scopes.get(i).close();
+				}
+				
+				// Stop observations
+				for (Observation observation : observations) {
+					if (!isListenerAdapterObservationAware()) {
+						observation.stop();
+					}
+				}
+			}
+		}
+
 	}