GH-10431: Propagate observation to error handler

Fixes: #10431

The `SourcePollingChannelAdapter` is able to start an observation for a just received message from a source.
This observation is completed when the polling task is done (including transaction).
However, a thrown from downstream error goes to the `errorChannel` already after finishing the mentioned
observation (even if that includes error).
According to end-user experience, the error handling must be a part of the same trace started for the received message.

* Fix `AbstractPollingEndpoint` to perform `donePollingTask()` only for a successful message reception
* Fix `SourcePollingChannelAdapter` to propagate its `ObservationRegistry` to the `ErrorHandlingTaskExecutor`
* Introduce an `Observation` handling in the `ErrorHandlingTaskExecutor`.
Essentially, stop currently on scope `Observation` after error handling for th running task failure
* Prove the feature with a new `SourcePollingChannelAdapterErrorObservationTests`

**Cherry-pick to `6.5.x`**

Author: artembilan

spring-integration-core/src/main/java/org/springframework/integration/endpoint/AbstractPollingEndpoint.java

@@ -156,8 +156,8 @@ public void setAdviceChain(@Nullable List<Advice> adviceChain) {
 	 * A negative number means retrieve unlimited messages until the {@code MessageSource} returns {@code null}.
 	 * Zero means do not poll for any records -
 	 * it can be considered as pausing if 'maxMessagesPerPoll' is later changed to a non-zero value.
-	 * The polling cycle may exit earlier if the source returns null for the current receive call.
-	 * @param maxMessagesPerPoll the number of message to poll per schedule.
+	 * The polling cycle may exit earlier if the source returns null for the current {@code receive} call.
+	 * @param maxMessagesPerPoll the number of messages to poll per schedule.
 	 */
 	@ManagedAttribute
 	public void setMaxMessagesPerPoll(long maxMessagesPerPoll) {
@@ -184,14 +184,14 @@ public void setTransactionSynchronizationFactory(
 	}
 
 	/**
-	 * Return the default error channel if the error handler is explicitly provided and
-	 * it is a {@link MessagePublishingErrorHandler}.
+	 * Return the default error channel if the error handler is explicitly provided,
+	 * and it is a {@link MessagePublishingErrorHandler}.
 	 * @return the channel or null.
 	 * @since 4.3
 	 */
 	public @Nullable MessageChannel getDefaultErrorChannel() {
-		if (!this.errorHandlerIsDefault && this.errorHandler
-				instanceof MessagePublishingErrorHandler messagePublishingErrorHandler) {
+		if (!this.errorHandlerIsDefault &&
+				this.errorHandler instanceof MessagePublishingErrorHandler messagePublishingErrorHandler) {
 
 			return messagePublishingErrorHandler.getDefaultErrorChannel();
 		}
@@ -333,11 +333,9 @@ protected void doStart() {
 		List<Advice> advices = this.adviceChain;
 		if (!CollectionUtils.isEmpty(advices)) {
 			ProxyFactory proxyFactory = new ProxyFactory(task);
-			if (!CollectionUtils.isEmpty(advices)) {
-				advices.stream()
-						.filter(advice -> !isReceiveOnlyAdvice(advice))
-						.forEach(proxyFactory::addAdvice);
-			}
+			advices.stream()
+					.filter(advice -> !isReceiveOnlyAdvice(advice))
+					.forEach(proxyFactory::addAdvice);
 			task = (Callable<@Nullable Message<?>>) proxyFactory.getProxy(this.beanClassLoader);
 		}
 		if (!CollectionUtils.isEmpty(receiveOnlyAdviceChain)) {
@@ -419,12 +417,12 @@ private Flux<Message<?>> createFluxGenerator() {
 	}
 
 	private @Nullable Message<?> pollForMessage() {
-		Exception pollingTaskError = null;
 		try {
-			return this.pollingTask.call();
+			Message<?> message = this.pollingTask.call();
+			donePollingTask(message);
+			return message;
 		}
 		catch (Exception ex) {
-			pollingTaskError = ex;
 			if (ex instanceof MessagingException messagingException) {
 				throw messagingException;
 			}
@@ -451,7 +449,6 @@ private Flux<Message<?>> createFluxGenerator() {
 					TransactionSynchronizationManager.unbindResource(resource);
 				}
 			}
-			donePollingTask(pollingTaskError);
 		}
 	}
 
@@ -501,7 +498,12 @@ protected void messageReceived(@Nullable IntegrationResourceHolder holder, Messa
 		}
 	}
 
-	protected void donePollingTask(@Nullable Exception pollingTaskError) {
+	/**
+	 * The callback of a received message (if any) after the polling task is done.
+	 * If a transaction is enabled, it is committed at this point.
+	 * @param message the message result from the polling task.
+	 */
+	protected void donePollingTask(@Nullable Message<?> message) {
 
 	}
 

spring-integration-core/src/main/java/org/springframework/integration/endpoint/SourcePollingChannelAdapter.java

@@ -44,6 +44,7 @@
 import org.springframework.integration.support.management.observation.MessageReceiverContext;
 import org.springframework.integration.support.management.observation.MessageReceiverObservationConvention;
 import org.springframework.integration.transaction.IntegrationResourceHolder;
+import org.springframework.integration.util.ErrorHandlingTaskExecutor;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessagingException;
@@ -141,6 +142,12 @@ public void setShouldTrack(boolean shouldTrack) {
 	@Override
 	public void registerObservationRegistry(ObservationRegistry observationRegistry) {
 		this.observationRegistry = observationRegistry;
+		if (isObserved()) {
+			ErrorHandlingTaskExecutor taskExecutor = (ErrorHandlingTaskExecutor) getTaskExecutor();
+			if (taskExecutor.getObservationRegistry() == null) {
+				taskExecutor.setObservationRegistry(observationRegistry);
+			}
+		}
 	}
 
 	/**
@@ -263,31 +270,30 @@ protected void handleMessage(Message<?> messageArg) {
 	@Override
 	protected void messageReceived(@Nullable IntegrationResourceHolder holder, Message<?> message) {
 		if (isObserved()) {
-			Observation observation =
-					IntegrationObservation.HANDLER.observation(this.observationConvention,
+			// Cannot use 'Observation.observe()' API
+			// since transaction needs to be aware of the message
+			// and error handling happens in the caller of 'doPoll()' - 'ErrorHandler'
+			IntegrationObservation.HANDLER.observation(this.observationConvention,
 							DefaultMessageReceiverObservationConvention.INSTANCE,
 							() -> new MessageReceiverContext(message, getComponentName(), "message-source"),
-							this.observationRegistry);
-
-			observation.start().openScope();
+							this.observationRegistry)
+					.start()
+					.openScope();
 		}
+
 		super.messageReceived(holder, message);
 	}
 
 	/**
 	 * Stop an observation (and close its scope) previously started
 	 * from the {@link #messageReceived(IntegrationResourceHolder, Message)}.
-	 * @param pollingTaskError an optional error as a result of the polling task.
 	 */
 	@Override
-	protected void donePollingTask(@Nullable Exception pollingTaskError) {
+	protected void donePollingTask(@Nullable Message<?> message) {
 		Observation.Scope currentObservationScope = this.observationRegistry.getCurrentObservationScope();
 		if (currentObservationScope != null) {
 			currentObservationScope.close();
 			Observation currentObservation = currentObservationScope.getCurrentObservation();
-			if (pollingTaskError != null) {
-				currentObservation.error(pollingTaskError);
-			}
 			currentObservation.stop();
 		}
 	}
@@ -302,8 +308,7 @@ protected String getResourceKey() {
 		return IntegrationResourceHolder.MESSAGE_SOURCE;
 	}
 
-	@Nullable
-	private static Object extractProxyTarget(@Nullable Object target) {
+	private static @Nullable Object extractProxyTarget(@Nullable Object target) {
 		if (!(target instanceof Advised advised)) {
 			return target;
 		}

spring-integration-core/src/main/java/org/springframework/integration/util/ErrorHandlingTaskExecutor.java

@@ -18,6 +18,10 @@
 
 import java.util.concurrent.Executor;
 
+import io.micrometer.observation.Observation;
+import io.micrometer.observation.ObservationRegistry;
+import org.jspecify.annotations.Nullable;
+
 import org.springframework.core.task.SyncTaskExecutor;
 import org.springframework.core.task.TaskExecutor;
 import org.springframework.util.Assert;
@@ -27,24 +31,43 @@
  * A {@link TaskExecutor} implementation that wraps an existing Executor
  * instance in order to catch any exceptions. If an exception is thrown, it
  * will be handled by the provided {@link ErrorHandler}.
+ * <p>
+ * If an {@link ObservationRegistry} is provided, the current observation in scope
+ * will be closed after error handling.
  *
  * @author Jonas Partner
  * @author Mark Fisher
  * @author Gary Russell
+ * @author Artem Bilan
  */
 public class ErrorHandlingTaskExecutor implements TaskExecutor {
 
 	private final Executor executor;
 
 	private final ErrorHandler errorHandler;
 
+	private @Nullable ObservationRegistry observationRegistry;
+
 	public ErrorHandlingTaskExecutor(Executor executor, ErrorHandler errorHandler) {
 		Assert.notNull(executor, "executor must not be null");
 		Assert.notNull(errorHandler, "errorHandler must not be null");
 		this.executor = executor;
 		this.errorHandler = errorHandler;
 	}
 
+	/**
+	 * Set an {@link ObservationRegistry} to close current observation in scope after error handling.
+	 * @param observationRegistry the {@link ObservationRegistry} to use.
+	 * @since 6.5.3
+	 */
+	public void setObservationRegistry(ObservationRegistry observationRegistry) {
+		this.observationRegistry = observationRegistry;
+	}
+
+	public @Nullable ObservationRegistry getObservationRegistry() {
+		return this.observationRegistry;
+	}
+
 	public boolean isSyncExecutor() {
 		return this.executor instanceof SyncTaskExecutor;
 	}
@@ -55,8 +78,21 @@ public void execute(final Runnable task) {
 			try {
 				task.run();
 			}
-			catch (Throwable t) { //NOSONAR
-				ErrorHandlingTaskExecutor.this.errorHandler.handleError(t);
+			catch (Throwable throwable) {
+				try {
+					ErrorHandlingTaskExecutor.this.errorHandler.handleError(throwable);
+				}
+				finally {
+					if (this.observationRegistry != null) {
+						var	currentObservationScope = this.observationRegistry.getCurrentObservationScope();
+						if (currentObservationScope != null) {
+							currentObservationScope.close();
+							Observation currentObservation = currentObservationScope.getCurrentObservation();
+							currentObservation.error(throwable);
+							currentObservation.stop();
+						}
+					}
+				}
 			}
 		});
 	}

spring-integration-core/src/test/java/org/springframework/integration/support/management/observation/SourcePollingChannelAdapterErrorObservationTests.java

@@ -0,0 +1,122 @@
+/*
+ * Copyright 2025-present 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.integration.support.management.observation;
+
+import io.micrometer.observation.ObservationRegistry;
+import io.micrometer.tracing.Span;
+import io.micrometer.tracing.test.SampleTestRunner;
+import io.micrometer.tracing.test.simple.SpansAssert;
+
+import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.integration.annotation.BridgeTo;
+import org.springframework.integration.annotation.EndpointId;
+import org.springframework.integration.annotation.InboundChannelAdapter;
+import org.springframework.integration.annotation.Poller;
+import org.springframework.integration.annotation.ServiceActivator;
+import org.springframework.integration.channel.DirectChannel;
+import org.springframework.integration.config.EnableIntegration;
+import org.springframework.integration.config.EnableIntegrationManagement;
+import org.springframework.integration.scheduling.PollerMetadata;
+import org.springframework.integration.test.util.OnlyOnceTrigger;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.awaitility.Awaitility.await;
+
+/**
+ * @author Artem Bilan
+ *
+ * @since 6.5.3
+ */
+public class SourcePollingChannelAdapterErrorObservationTests extends SampleTestRunner {
+
+	@Override
+	public TracingSetup[] getTracingSetup() {
+		return new TracingSetup[] {TracingSetup.IN_MEMORY_BRAVE};
+	}
+
+	@Override
+	public SampleTestRunnerConsumer yourCode() {
+		return (bb, meterRegistry) -> {
+			ObservationRegistry observationRegistry = getObservationRegistry();
+
+			try (AnnotationConfigApplicationContext applicationContext = new AnnotationConfigApplicationContext()) {
+				applicationContext.registerBean(ObservationRegistry.class, () -> observationRegistry);
+				applicationContext.register(ObservationIntegrationTestConfiguration.class);
+				applicationContext.refresh();
+			}
+
+			await().untilAsserted(() -> assertThat(bb.getFinishedSpans()).hasSize(5));
+
+			SpansAssert.assertThat(bb.getFinishedSpans())
+					.haveSameTraceId()
+					.hasASpanWithName("dataMessageSource receive", spanAssert -> spanAssert
+							.hasTag(IntegrationObservation.GatewayTags.COMPONENT_TYPE.asString(), "message-source")
+							.hasKindEqualTo(Span.Kind.CONSUMER))
+					.hasASpanWithName("inputChannel send", spanAssert -> spanAssert
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_NAME.asString(), "inputChannel")
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_TYPE.asString(), "producer")
+							.hasKindEqualTo(Span.Kind.PRODUCER))
+					.hasASpanWithName("dataHandler receive", spanAssert -> spanAssert
+							.hasTag(IntegrationObservation.HandlerTags.COMPONENT_NAME.asString(), "dataHandler")
+							.hasTag(IntegrationObservation.HandlerTags.COMPONENT_TYPE.asString(), "handler")
+							.hasKindEqualTo(Span.Kind.CONSUMER))
+					.hasASpanWithName("errorChannel send", spanAssert -> spanAssert
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_NAME.asString(), "errorChannel")
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_TYPE.asString(), "producer")
+							.hasKindEqualTo(Span.Kind.PRODUCER))
+					.hasASpanWithName("errorChannel.bridgeTo receive", spanAssert -> spanAssert
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_NAME.asString(), "errorChannel.bridgeTo")
+							.hasTag(IntegrationObservation.ProducerTags.COMPONENT_TYPE.asString(), "handler")
+							.hasKindEqualTo(Span.Kind.CONSUMER));
+		};
+	}
+
+	@Configuration
+	@EnableIntegration
+	@EnableIntegrationManagement(observationPatterns = "*")
+	static class ObservationIntegrationTestConfiguration {
+
+		@Bean
+		PollerMetadata pollerMetadata() {
+			PollerMetadata pollerMetadata = new PollerMetadata();
+			pollerMetadata.setTrigger(new OnlyOnceTrigger());
+			return pollerMetadata;
+		}
+
+		@EndpointId("dataMessageSource")
+		@InboundChannelAdapter(channel = "inputChannel", poller = @Poller("pollerMetadata"))
+		String errorData() {
+			return "some data";
+		}
+
+		@EndpointId("dataHandler")
+		@ServiceActivator(inputChannel = "inputChannel")
+		void processData(String data) {
+			throw new RuntimeException("intentional");
+		}
+
+		@Bean
+		@BridgeTo("nullChannel")
+		DirectChannel errorChannel() {
+			return new DirectChannel();
+		}
+
+	}
+
+}