Instrument pollable consumers with Spring Integration metrics

core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/binder/DefaultPollableMessageSource.java

@@ -18,15 +18,19 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiConsumer;
 
+import io.micrometer.observation.ObservationRegistry;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.aop.support.NameMatchMethodPointcutAdvisor;
+import org.springframework.beans.factory.BeanNameAware;
 import org.springframework.cloud.function.context.catalog.FunctionTypeUtils;
 import org.springframework.context.Lifecycle;
 import org.springframework.core.AttributeAccessor;
@@ -42,6 +46,15 @@
 import org.springframework.integration.support.DefaultErrorMessageStrategy;
 import org.springframework.integration.support.ErrorMessageStrategy;
 import org.springframework.integration.support.ErrorMessageUtils;
+import org.springframework.integration.support.management.IntegrationManagement;
+import org.springframework.integration.support.management.metrics.CounterFacade;
+import org.springframework.integration.support.management.metrics.MeterFacade;
+import org.springframework.integration.support.management.metrics.MetricsCaptor;
+import org.springframework.integration.support.management.metrics.SampleFacade;
+import org.springframework.integration.support.management.metrics.TimerFacade;
+import org.springframework.integration.support.management.observation.DefaultMessageReceiverObservationConvention;
+import org.springframework.integration.support.management.observation.IntegrationObservation;
+import org.springframework.integration.support.management.observation.MessageReceiverContext;
 import org.springframework.lang.Nullable;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
@@ -66,7 +79,7 @@
  * @since 2.0
  *
  */
-public class DefaultPollableMessageSource implements PollableMessageSource, Lifecycle {
+public class DefaultPollableMessageSource implements PollableMessageSource, Lifecycle, IntegrationManagement, BeanNameAware {
 
 	private static final Log log = LogFactory.getLog(DefaultPollableMessageSource.class);
 
@@ -98,6 +111,18 @@ public class DefaultPollableMessageSource implements PollableMessageSource, Life
 
 	private boolean running;
 
+	private final Set<MeterFacade> meters = ConcurrentHashMap.newKeySet();
+
+	private String beanName;
+
+	private MetricsCaptor metricsCaptor;
+
+	private ObservationRegistry observationRegistry = ObservationRegistry.NOOP;
+
+	private CounterFacade receiveCounter;
+
+	private TimerFacade successTimer;
+
 	/**
 	 * @param messageConverter instance of {@link SmartMessageConverter}. Can be null.
 	 */
@@ -106,6 +131,45 @@ public DefaultPollableMessageSource(
 		this.messageConverter = messageConverter;
 	}
 
+	@Override
+	public void setBeanName(String name) {
+		this.beanName = name;
+	}
+
+	@Override
+	public String getComponentName() {
+		return this.beanName;
+	}
+
+	@Override
+	public String getComponentType() {
+		return "pollable-message-source";
+	}
+
+	@Override
+	public void registerMetricsCaptor(MetricsCaptor captor) {
+		this.metricsCaptor = captor;
+	}
+
+	@Override
+	public void registerObservationRegistry(ObservationRegistry registry) {
+		Assert.notNull(registry, "'observationRegistry' must not be null");
+		this.observationRegistry = registry;
+	}
+
+	@Override
+	public boolean isObserved() {
+		return !ObservationRegistry.NOOP.equals(this.observationRegistry);
+	}
+
+	@Override
+	public void destroy() {
+		this.meters.forEach(MeterFacade::remove);
+		this.meters.clear();
+		this.receiveCounter = null;
+		this.successTimer = null;
+	}
+
 	public void setSource(MessageSource<?> source) {
 		ProxyFactory pf = new ProxyFactory(source);
 
@@ -235,10 +299,25 @@ public boolean poll(MessageHandler handler) {
 
 	@Override
 	public boolean poll(MessageHandler handler, ParameterizedTypeReference<?> type) {
-		Message<?> message = this.receive(type);
+		Message<?> message;
+		try {
+			message = this.receive(type);
+		}
+		catch (RuntimeException ex) {
+			if (this.metricsCaptor != null) {
+				receiveCounter(false, ex.getClass().getSimpleName()).increment();
+			}
+			throw ex;
+		}
 		if (message == null) {
 			return false;
 		}
+		if (this.metricsCaptor != null) {
+			if (this.receiveCounter == null) {
+				this.receiveCounter = receiveCounter(true, "none");
+			}
+			this.receiveCounter.increment();
+		}
 
 		AcknowledgmentCallback ackCallback = StaticMessageHeaderAccessor.getAcknowledgmentCallback(message);
 
@@ -361,6 +440,33 @@ private Message<?> receive(ParameterizedTypeReference<?> type) {
 	}
 
 	private void doHandleMessage(MessageHandler handler, Message<?> message) {
+		if (isObserved()) {
+			IntegrationObservation.HANDLER.observation(null,
+							DefaultMessageReceiverObservationConvention.INSTANCE,
+							() -> new MessageReceiverContext(message, getComponentName()),
+							this.observationRegistry)
+					.observe(() -> invokeHandler(handler, message));
+		}
+		else if (this.metricsCaptor != null) {
+			SampleFacade sample = this.metricsCaptor.start();
+			try {
+				invokeHandler(handler, message);
+				if (this.successTimer == null) {
+					this.successTimer = sendTimer(true, "none");
+				}
+				sample.stop(this.successTimer);
+			}
+			catch (RuntimeException ex) {
+				sample.stop(sendTimer(false, ex.getClass().getSimpleName()));
+				throw ex;
+			}
+		}
+		else {
+			invokeHandler(handler, message);
+		}
+	}
+
+	private void invokeHandler(MessageHandler handler, Message<?> message) {
 		try {
 			handler.handleMessage(message);
 		}
@@ -369,6 +475,30 @@ private void doHandleMessage(MessageHandler handler, Message<?> message) {
 		}
 	}
 
+	private CounterFacade receiveCounter(boolean success, String exception) {
+		CounterFacade counter = this.metricsCaptor.counterBuilder(RECEIVE_COUNTER_NAME)
+				.tag("name", getComponentName())
+				.tag("type", "source")
+				.tag("result", success ? "success" : "failure")
+				.tag("exception", exception)
+				.description("Messages received")
+				.build();
+		this.meters.add(counter);
+		return counter;
+	}
+
+	private TimerFacade sendTimer(boolean success, String exception) {
+		TimerFacade timer = this.metricsCaptor.timerBuilder(SEND_TIMER_NAME)
+				.tag("type", "handler")
+				.tag("name", getComponentName())
+				.tag("result", success ? "success" : "failure")
+				.tag("exception", exception)
+				.description("Send processing time")
+				.build();
+		this.meters.add(timer);
+		return timer;
+	}
+
 	/**
 	 * If there's a retry template, it will set the attribute holder via the listener. If
 	 * there's no retry template, but there's an error channel, we create a new attribute

core/spring-cloud-stream/src/main/java/org/springframework/cloud/stream/function/BindableFunctionProxyFactory.java

@@ -186,7 +186,8 @@ private void createInput(String name) {
 		if (this.supportedBindableFeatures.isPollable()) {
 			PollableMessageSource pollableSource = (PollableMessageSource) getBindingTargetFactory(PollableMessageSource.class).createInput(name);
 			if (context != null && !context.containsBean(name)) {
-				context.registerBean(name, PollableMessageSource.class, () -> pollableSource);
+				// register with the concrete type so that bean type matching (e.g., IntegrationManagement) works
+				context.registerBean(name, (Class<PollableMessageSource>) pollableSource.getClass(), () -> pollableSource);
 			}
 			this.inputHolders.put(name, new BoundTargetHolder(pollableSource, true));
 		}

core/spring-cloud-stream/src/test/java/org/springframework/cloud/stream/binder/DefaultPollableMessageSourceTests.java

@@ -20,16 +20,19 @@
 import java.util.List;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.integration.channel.DirectChannel;
+import org.springframework.integration.support.management.micrometer.MicrometerMetricsCaptor;
 import org.springframework.messaging.Message;
 import org.springframework.messaging.MessageChannel;
 import org.springframework.messaging.MessageHandler;
 import org.springframework.messaging.support.ChannelInterceptor;
 import org.springframework.messaging.support.MessageBuilder;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
  * Verifies that {@link DefaultPollableMessageSource} invokes the full
@@ -42,6 +45,30 @@ class DefaultPollableMessageSourceTests {
 
 	private final AtomicInteger receiveCount = new AtomicInteger();
 
+	@Test
+	void pollRecordsSpringIntegrationMetrics() {
+		SimpleMeterRegistry registry = new SimpleMeterRegistry();
+		DefaultPollableMessageSource source = newSource();
+		source.setBeanName("pollableInput-in-0");
+		source.registerMetricsCaptor(new MicrometerMetricsCaptor(registry));
+		source.setSource(this::message);
+
+		assertThat(source.poll(noopHandler())).isTrue();
+		assertThatThrownBy(() -> source.poll(m -> {
+			throw new IllegalStateException("boom");
+		})).isInstanceOf(RuntimeException.class);
+
+		assertThat(registry.find("spring.integration.receive")
+				.tag("name", "pollableInput-in-0").tag("type", "source").counter().count()).isEqualTo(2.0);
+		assertThat(registry.find("spring.integration.send").tag("type", "handler")
+				.tag("name", "pollableInput-in-0").tag("result", "success").timer().count()).isEqualTo(1);
+		assertThat(registry.find("spring.integration.send").tag("type", "handler")
+				.tag("name", "pollableInput-in-0").tag("result", "failure").timer().count()).isEqualTo(1);
+
+		source.destroy();
+		assertThat(registry.find("spring.integration.send").timers()).isEmpty();
+	}
+
 	@Test
 	void successfulPollInvokesFullInterceptorLifecycleInOrder() {
 		DefaultPollableMessageSource source = newSource();