Merge pull request #45360 from dmitrysulman

* pr/47287:
  Polish "Register controller advices to RSocket messaging"
  Register controller advices to RSocket messaging

Closes gh-45360

documentation/spring-boot-docs/src/docs/antora/modules/reference/pages/messaging/rsocket.adoc

@@ -69,6 +69,8 @@ Spring Boot will auto-configure the Spring Messaging infrastructure for RSocket.
 
 This means that Spring Boot will create a javadoc:org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler[] bean that will handle RSocket requests to your application.
 
+TIP: You can use {url-spring-framework-docs}/web/webmvc/mvc-controller/ann-advice.html[`@ControllerAdvice`] to handle exceptions.
+
 
 
 [[messaging.rsocket.requester]]

module/spring-boot-rsocket/src/main/java/org/springframework/boot/rsocket/autoconfigure/RSocketMessagingAutoConfiguration.java

@@ -17,22 +17,29 @@
 package org.springframework.boot.rsocket.autoconfigure;
 
 import io.rsocket.transport.netty.server.TcpServerTransport;
+import org.jspecify.annotations.Nullable;
 
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.messaging.handler.MessagingAdviceBean;
 import org.springframework.messaging.rsocket.RSocketRequester;
 import org.springframework.messaging.rsocket.RSocketStrategies;
 import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
+import org.springframework.web.method.ControllerAdviceBean;
 
 /**
  * {@link EnableAutoConfiguration Auto-configuration} for Spring RSocket support in Spring
  * Messaging.
  *
  * @author Brian Clozel
+ * @author Dmitry Sulman
+ * @author Stephane Nicoll
  * @since 4.0.0
  */
 @AutoConfiguration(after = RSocketStrategiesAutoConfiguration.class)
@@ -42,11 +49,71 @@ public final class RSocketMessagingAutoConfiguration {
 	@Bean
 	@ConditionalOnMissingBean
 	RSocketMessageHandler messageHandler(RSocketStrategies rSocketStrategies,
-			ObjectProvider<RSocketMessageHandlerCustomizer> customizers) {
+			ObjectProvider<RSocketMessageHandlerCustomizer> customizers, ApplicationContext context) {
 		RSocketMessageHandler messageHandler = new RSocketMessageHandler();
 		messageHandler.setRSocketStrategies(rSocketStrategies);
 		customizers.orderedStream().forEach((customizer) -> customizer.customize(messageHandler));
+
 		return messageHandler;
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	@ConditionalOnClass(ControllerAdviceBean.class)
+	static class MessagingAdviceConfiguration {
+
+		@Bean
+		MessagingAdviceRSocketMessageHandlerCustomizer messagingAdviceRSocketMessageHandlerCustomizer(
+				ApplicationContext applicationContext) {
+			return new MessagingAdviceRSocketMessageHandlerCustomizer(applicationContext);
+		}
+
+	}
+
+	static final class MessagingAdviceRSocketMessageHandlerCustomizer implements RSocketMessageHandlerCustomizer {
+
+		private final ApplicationContext applicationContext;
+
+		MessagingAdviceRSocketMessageHandlerCustomizer(ApplicationContext applicationContext) {
+			this.applicationContext = applicationContext;
+		}
+
+		@Override
+		public void customize(RSocketMessageHandler messageHandler) {
+			ControllerAdviceBean.findAnnotatedBeans(this.applicationContext)
+				.forEach((controllerAdviceBean) -> messageHandler
+					.registerMessagingAdvice(new ControllerAdviceBeanWrapper(controllerAdviceBean)));
+		}
+
+	}
+
+	private static final class ControllerAdviceBeanWrapper implements MessagingAdviceBean {
+
+		private final ControllerAdviceBean adviceBean;
+
+		private ControllerAdviceBeanWrapper(ControllerAdviceBean adviceBean) {
+			this.adviceBean = adviceBean;
+		}
+
+		@Override
+		public @Nullable Class<?> getBeanType() {
+			return this.adviceBean.getBeanType();
+		}
+
+		@Override
+		public Object resolveBean() {
+			return this.adviceBean.resolveBean();
+		}
+
+		@Override
+		public boolean isApplicableToBeanType(Class<?> beanType) {
+			return this.adviceBean.isApplicableToBeanType(beanType);
+		}
+
+		@Override
+		public int getOrder() {
+			return this.adviceBean.getOrder();
+		}
+
+	}
+
 }

module/spring-boot-rsocket/src/test/java/org/springframework/boot/rsocket/autoconfigure/RSocketMessagingAutoConfigurationTests.java

@@ -16,7 +16,10 @@
 
 package org.springframework.boot.rsocket.autoconfigure;
 
+import io.rsocket.frame.FrameType;
 import org.junit.jupiter.api.Test;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
 
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
@@ -24,9 +27,19 @@ import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.codec.CharSequenceEncoder;
 import org.springframework.core.codec.StringDecoder;
+import org.springframework.messaging.Message;
+import org.springframework.messaging.handler.DestinationPatternsMessageCondition;
+import org.springframework.messaging.handler.annotation.MessageExceptionHandler;
+import org.springframework.messaging.handler.annotation.MessageMapping;
 import org.springframework.messaging.rsocket.RSocketStrategies;
+import org.springframework.messaging.rsocket.annotation.support.RSocketFrameTypeMessageCondition;
 import org.springframework.messaging.rsocket.annotation.support.RSocketMessageHandler;
+import org.springframework.messaging.support.MessageBuilder;
+import org.springframework.messaging.support.MessageHeaderAccessor;
+import org.springframework.stereotype.Controller;
 import org.springframework.util.MimeType;
+import org.springframework.util.RouteMatcher;
+import org.springframework.web.bind.annotation.ControllerAdvice;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -72,6 +85,21 @@ class RSocketMessagingAutoConfigurationTests {
 		});
 	}
 
+	@Test
+	void shouldRegisterControllerAdvice() {
+		this.contextRunner.withBean(TestControllerAdvice.class).withBean(TestController.class).run((context) -> {
+			RSocketMessageHandler handler = context.getBean(RSocketMessageHandler.class);
+			MessageHeaderAccessor headers = new MessageHeaderAccessor();
+			RouteMatcher.Route route = handler.getRouteMatcher().parseRoute("exception");
+			headers.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, route);
+			headers.setHeader(RSocketFrameTypeMessageCondition.FRAME_TYPE_HEADER, FrameType.REQUEST_FNF);
+			Message<?> message = MessageBuilder.createMessage(Mono.empty(), headers.getMessageHeaders());
+
+			StepVerifier.create(handler.handleMessage(message)).expectComplete().verify();
+			assertThat(context.getBean(TestControllerAdvice.class).isExceptionHandled()).isTrue();
+		});
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class BaseConfiguration {
 
@@ -111,4 +139,30 @@ class RSocketMessagingAutoConfigurationTests {
 
 	}
 
+	@Controller
+	static final class TestController {
+
+		@MessageMapping("exception")
+		void handleWithSimulatedException() {
+			throw new IllegalStateException("simulated exception");
+		}
+
+	}
+
+	@ControllerAdvice
+	static final class TestControllerAdvice {
+
+		boolean exceptionHandled;
+
+		boolean isExceptionHandled() {
+			return this.exceptionHandled;
+		}
+
+		@MessageExceptionHandler
+		void handleException(IllegalStateException ex) {
+			this.exceptionHandled = true;
+		}
+
+	}
+
 }