From 504b7619bdb9580bd6346bd256263239283edf07 Mon Sep 17 00:00:00 2001 From: rstoyanchev Date: Wed, 14 Feb 2024 16:26:18 +0000 Subject: [PATCH] WebSocketMessageBrokerConfigurer allows to configure Lifecycle phase Closes gh-32205 --- .../SimpAnnotationMethodMessageHandler.java | 19 ++++++- .../broker/AbstractBrokerMessageHandler.java | 19 ++++++- .../AbstractMessageBrokerConfiguration.java | 57 +++++++++++++++---- .../user/UserDestinationMessageHandler.java | 19 ++++++- ...ngWebSocketMessageBrokerConfiguration.java | 13 ++++- ...cketMessageBrokerConfigurationSupport.java | 21 +++++-- .../WebSocketMessageBrokerConfigurer.java | 22 ++++++- .../SubProtocolWebSocketHandler.java | 17 ++++++ .../support/WebSocketHandlerMapping.java | 19 ++++++- ...essageBrokerConfigurationSupportTests.java | 53 ++++++++++++++++- 10 files changed, 238 insertions(+), 21 deletions(-) diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java index b44383bf61d..9974aa8ea52 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -120,6 +120,9 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan @Nullable private MessageHeaderInitializer headerInitializer; + @Nullable + private Integer phase; + private volatile boolean running; private final Object lifecycleMonitor = new Object(); @@ -271,6 +274,20 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan return this.headerInitializer; } + /** + * Set the phase that this handler should run in. + *

By default, this is {@link SmartLifecycle#DEFAULT_PHASE}. + * @since 6.1.4 + */ + public void setPhase(int phase) { + this.phase = phase; + } + + @Override + public int getPhase() { + return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase()); + } + @Override public final void start() { diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java index 1abc3fa86ee..e71eb4ad29f 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -75,6 +75,9 @@ public abstract class AbstractBrokerMessageHandler private boolean autoStartup = true; + @Nullable + private Integer phase; + private volatile boolean running; private final Object lifecycleMonitor = new Object(); @@ -197,6 +200,20 @@ public abstract class AbstractBrokerMessageHandler return this.autoStartup; } + /** + * Set the phase that this handler should run in. + *

By default, this is {@link SmartLifecycle#DEFAULT_PHASE}. + * @since 6.1.4 + */ + public void setPhase(int phase) { + this.phase = phase; + } + + @Override + public int getPhase() { + return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase()); + } + @Override public void start() { diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java index 50c90329723..b38cfc27d29 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2022 the original author or authors. + * Copyright 2002-2024 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. @@ -27,6 +27,7 @@ import org.springframework.beans.factory.BeanInitializationException; import org.springframework.beans.factory.annotation.Qualifier; import org.springframework.context.ApplicationContext; import org.springframework.context.ApplicationContextAware; +import org.springframework.context.SmartLifecycle; import org.springframework.context.annotation.Bean; import org.springframework.context.event.SmartApplicationListener; import org.springframework.core.task.TaskExecutor; @@ -59,6 +60,7 @@ import org.springframework.messaging.support.AbstractSubscribableChannel; import org.springframework.messaging.support.ExecutorSubscribableChannel; import org.springframework.messaging.support.ImmutableMessageChannelInterceptor; import org.springframework.scheduling.TaskScheduler; +import org.springframework.scheduling.concurrent.ExecutorConfigurationSupport; import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.util.Assert; @@ -132,6 +134,9 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC @Nullable private MessageBrokerRegistry brokerRegistry; + @Nullable + private Integer phase; + /** * Protected constructor. @@ -166,8 +171,12 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC @Bean public TaskExecutor clientInboundChannelExecutor() { - return getTaskExecutor(getClientInboundChannelRegistration(), - "clientInboundChannel-", this::defaultTaskExecutor); + ChannelRegistration registration = getClientInboundChannelRegistration(); + TaskExecutor executor = getTaskExecutor(registration, "clientInboundChannel-", this::defaultTaskExecutor); + if (executor instanceof ExecutorConfigurationSupport executorSupport) { + executorSupport.setPhase(getPhase()); + } + return executor; } protected final ChannelRegistration getClientInboundChannelRegistration() { @@ -180,6 +189,17 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC return this.clientInboundChannelRegistration; } + protected final int getPhase() { + if (this.phase == null) { + this.phase = initPhase(); + } + return this.phase; + } + + protected int initPhase() { + return SmartLifecycle.DEFAULT_PHASE; + } + /** * A hook for subclasses to customize the message channel for inbound messages * from WebSocket clients. @@ -193,17 +213,21 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC ExecutorSubscribableChannel channel = new ExecutorSubscribableChannel(executor); channel.setLogger(SimpLogging.forLog(channel.getLogger())); - ChannelRegistration reg = getClientOutboundChannelRegistration(); - if (reg.hasInterceptors()) { - channel.setInterceptors(reg.getInterceptors()); + ChannelRegistration registration = getClientOutboundChannelRegistration(); + if (registration.hasInterceptors()) { + channel.setInterceptors(registration.getInterceptors()); } return channel; } @Bean public TaskExecutor clientOutboundChannelExecutor() { - return getTaskExecutor(getClientOutboundChannelRegistration(), - "clientOutboundChannel-", this::defaultTaskExecutor); + ChannelRegistration registration = getClientOutboundChannelRegistration(); + TaskExecutor executor = getTaskExecutor(registration, "clientOutboundChannel-", this::defaultTaskExecutor); + if (executor instanceof ExecutorConfigurationSupport executorSupport) { + executorSupport.setPhase(getPhase()); + } + return executor; } protected final ChannelRegistration getClientOutboundChannelRegistration() { @@ -244,7 +268,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC MessageBrokerRegistry registry = getBrokerRegistry(clientInboundChannel, clientOutboundChannel); ChannelRegistration registration = registry.getBrokerChannelRegistration(); - return getTaskExecutor(registration, "brokerChannel-", () -> { + TaskExecutor executor = getTaskExecutor(registration, "brokerChannel-", () -> { // Should never be used ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor(); threadPoolTaskExecutor.setCorePoolSize(0); @@ -252,6 +276,10 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC threadPoolTaskExecutor.setQueueCapacity(0); return threadPoolTaskExecutor; }); + if (executor instanceof ExecutorConfigurationSupport executorSupport) { + executorSupport.setPhase(getPhase()); + } + return executor; } private TaskExecutor defaultTaskExecutor() { @@ -316,6 +344,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC handler.setDestinationPrefixes(brokerRegistry.getApplicationDestinationPrefixes()); handler.setMessageConverter(brokerMessageConverter); handler.setValidator(simpValidator()); + handler.setPhase(getPhase()); List argumentResolvers = new ArrayList<>(); addArgumentResolvers(argumentResolvers); @@ -329,6 +358,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC if (pathMatcher != null) { handler.setPathMatcher(pathMatcher); } + return handler; } @@ -342,8 +372,11 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC AbstractSubscribableChannel clientInboundChannel, AbstractSubscribableChannel clientOutboundChannel, SimpMessagingTemplate brokerMessagingTemplate) { - return new SimpAnnotationMethodMessageHandler( + SimpAnnotationMethodMessageHandler handler = new SimpAnnotationMethodMessageHandler( clientInboundChannel, clientOutboundChannel, brokerMessagingTemplate); + + handler.setPhase(getPhase()); + return handler; } protected void addArgumentResolvers(List argumentResolvers) { @@ -364,6 +397,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC return null; } updateUserDestinationResolver(handler, userDestinationResolver, registry.getUserDestinationPrefix()); + handler.setPhase(getPhase()); return handler; } @@ -403,6 +437,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC } handler.setSystemSubscriptions(subscriptions); updateUserDestinationResolver(handler, userDestinationResolver, registry.getUserDestinationPrefix()); + handler.setPhase(getPhase()); return handler; } @@ -419,6 +454,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC if (destination != null) { handler.setBroadcastDestination(destination); } + handler.setPhase(getPhase()); return handler; } @@ -446,6 +482,7 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC scheduler.setThreadNamePrefix("MessageBroker-"); scheduler.setPoolSize(Runtime.getRuntime().availableProcessors()); scheduler.setRemoveOnCancelPolicy(true); + scheduler.setPhase(getPhase()); return scheduler; } diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java index f60d261ef78..20cd5e5e479 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -77,6 +77,9 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec private volatile boolean running; + @Nullable + private Integer phase; + private final Object lifecycleMonitor = new Object(); @@ -154,6 +157,20 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec return this.headerInitializer; } + /** + * Set the phase that this handler should run in. + *

By default, this is {@link SmartLifecycle#DEFAULT_PHASE}. + * @since 6.1.4 + */ + public void setPhase(int phase) { + this.phase = phase; + } + + @Override + public int getPhase() { + return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase()); + } + @Override public final void start() { diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/DelegatingWebSocketMessageBrokerConfiguration.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/DelegatingWebSocketMessageBrokerConfiguration.java index 1e0d704ce8a..fc91e97367c 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/DelegatingWebSocketMessageBrokerConfiguration.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/DelegatingWebSocketMessageBrokerConfiguration.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2020 the original author or authors. + * Copyright 2002-2024 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. @@ -114,4 +114,15 @@ public class DelegatingWebSocketMessageBrokerConfiguration extends WebSocketMess } } + @Override + protected int initPhase() { + for (WebSocketMessageBrokerConfigurer configurer : this.configurers) { + Integer phase = configurer.getPhase(); + if (phase != null) { + return phase; + } + } + return super.initPhase(); + } + } diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java index aad0ba3db39..a3cd5f1c78b 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -35,12 +35,14 @@ import org.springframework.messaging.simp.user.SimpUserRegistry; import org.springframework.messaging.support.AbstractSubscribableChannel; import org.springframework.scheduling.TaskScheduler; import org.springframework.web.servlet.HandlerMapping; +import org.springframework.web.servlet.handler.AbstractHandlerMapping; import org.springframework.web.socket.WebSocketHandler; import org.springframework.web.socket.config.WebSocketMessageBrokerStats; import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory; import org.springframework.web.socket.messaging.DefaultSimpUserRegistry; import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler; import org.springframework.web.socket.messaging.WebSocketAnnotationMethodMessageHandler; +import org.springframework.web.socket.server.support.WebSocketHandlerMapping; /** * Extends {@link AbstractMessageBrokerConfiguration} and adds configuration for @@ -66,8 +68,11 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac AbstractSubscribableChannel clientInboundChannel,AbstractSubscribableChannel clientOutboundChannel, SimpMessagingTemplate brokerMessagingTemplate) { - return new WebSocketAnnotationMethodMessageHandler( + WebSocketAnnotationMethodMessageHandler handler = new WebSocketAnnotationMethodMessageHandler( clientInboundChannel, clientOutboundChannel, brokerMessagingTemplate); + + handler.setPhase(getPhase()); + return handler; } @Override @@ -93,14 +98,22 @@ public abstract class WebSocketMessageBrokerConfigurationSupport extends Abstrac } registerStompEndpoints(registry); OrderedMessageChannelDecorator.configureInterceptor(clientInboundChannel, registry.isPreserveReceiveOrder()); - return registry.getHandlerMapping(); + AbstractHandlerMapping handlerMapping = registry.getHandlerMapping(); + if (handlerMapping instanceof WebSocketHandlerMapping webSocketMapping) { + webSocketMapping.setPhase(getPhase()); + } + return handlerMapping; } @Bean public WebSocketHandler subProtocolWebSocketHandler( AbstractSubscribableChannel clientInboundChannel, AbstractSubscribableChannel clientOutboundChannel) { - return new SubProtocolWebSocketHandler(clientInboundChannel, clientOutboundChannel); + SubProtocolWebSocketHandler handler = + new SubProtocolWebSocketHandler(clientInboundChannel, clientOutboundChannel); + + handler.setPhase(getPhase()); + return handler; } protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) { diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurer.java b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurer.java index 24036055c09..5ac4cd2618e 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurer.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurer.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2022 the original author or authors. + * Copyright 2002-2024 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. @@ -18,6 +18,8 @@ package org.springframework.web.socket.config.annotation; import java.util.List; +import org.springframework.context.SmartLifecycle; +import org.springframework.lang.Nullable; import org.springframework.messaging.converter.MessageConverter; import org.springframework.messaging.handler.invocation.HandlerMethodArgumentResolver; import org.springframework.messaging.handler.invocation.HandlerMethodReturnValueHandler; @@ -110,4 +112,22 @@ public interface WebSocketMessageBrokerConfigurer { default void configureMessageBroker(MessageBrokerRegistry registry) { } + /** + * Return the {@link SmartLifecycle#getPhase() phase} that WebSocket message + * handling beans of type {@link SmartLifecycle} should run in. + *

The default implementation returns {@link null} which allows other + * configurers to decide. As soon as any configurer returns a value, that + * value is used. If no configurer returns a value, then by default + * {@link SmartLifecycle#DEFAULT_PHASE} is used. + *

It is recommended to use a phase value such as 0 in order to ensure that + * components start before the web server in Spring Boot application. In 6.2.0, + * the default used will change to 0. + * @since 6.1.4 + * @see SmartLifecycle + */ + @Nullable + default Integer getPhase() { + return null; + } + } diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/SubProtocolWebSocketHandler.java b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/SubProtocolWebSocketHandler.java index d93970ea293..867e780abe9 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/messaging/SubProtocolWebSocketHandler.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/messaging/SubProtocolWebSocketHandler.java @@ -103,6 +103,9 @@ public class SubProtocolWebSocketHandler private final DefaultStats stats = new DefaultStats(); + @Nullable + private Integer phase; + private volatile boolean running; private final Object lifecycleMonitor = new Object(); @@ -249,6 +252,20 @@ public class SubProtocolWebSocketHandler return this.timeToFirstMessage; } + /** + * Set the phase that this handler should run in. + *

By default, this is {@link SmartLifecycle#DEFAULT_PHASE}. + * @since 6.1.4 + */ + public void setPhase(int phase) { + this.phase = phase; + } + + @Override + public int getPhase() { + return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase()); + } + /** * Return a String describing internal state and counters. * Effectively {@code toString()} on {@link #getStats() getStats()}. diff --git a/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java b/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java index ebd03280c60..0d825e50d83 100644 --- a/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java +++ b/spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -40,6 +40,9 @@ public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping implements private boolean webSocketUpgradeMatch; + @Nullable + private Integer phase; + private volatile boolean running; @@ -57,6 +60,20 @@ public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping implements this.webSocketUpgradeMatch = match; } + /** + * Set the phase that this handler should run in. + *

By default, this is {@link SmartLifecycle#DEFAULT_PHASE}. + * @since 6.1.4 + */ + public void setPhase(int phase) { + this.phase = phase; + } + + @Override + public int getPhase() { + return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase()); + } + @Override protected void initServletContext(ServletContext servletContext) { diff --git a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java index 2a48863016c..fce2e7f2957 100644 --- a/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java +++ b/spring-websocket/src/test/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupportTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2002-2023 the original author or authors. + * Copyright 2002-2024 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. @@ -21,6 +21,7 @@ import java.util.List; import java.util.Map; import java.util.Set; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.function.Consumer; import org.junit.jupiter.api.Test; @@ -35,6 +36,7 @@ import org.springframework.messaging.handler.annotation.MessageMapping; import org.springframework.messaging.handler.annotation.SendTo; import org.springframework.messaging.simp.SimpMessageType; import org.springframework.messaging.simp.annotation.SubscribeMapping; +import org.springframework.messaging.simp.annotation.support.SimpAnnotationMethodMessageHandler; import org.springframework.messaging.simp.broker.DefaultSubscriptionRegistry; import org.springframework.messaging.simp.broker.SimpleBrokerMessageHandler; import org.springframework.messaging.simp.config.MessageBrokerRegistry; @@ -45,6 +47,7 @@ import org.springframework.messaging.support.AbstractSubscribableChannel; import org.springframework.messaging.support.ChannelInterceptor; import org.springframework.messaging.support.ExecutorSubscribableChannel; import org.springframework.messaging.support.ImmutableMessageChannelInterceptor; +import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler; import org.springframework.stereotype.Controller; import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping; @@ -57,6 +60,7 @@ import org.springframework.web.socket.messaging.StompSubProtocolHandler; import org.springframework.web.socket.messaging.StompTextMessageBuilder; import org.springframework.web.socket.messaging.SubProtocolHandler; import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler; +import org.springframework.web.socket.server.support.WebSocketHandlerMapping; import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler; import static org.assertj.core.api.Assertions.assertThat; @@ -207,6 +211,27 @@ class WebSocketMessageBrokerConfigurationSupportTests { assertThat(actual).matches(expected); } + @Test + void lifecyclePhase() { + ApplicationContext context = createContext(LifecyclePhaseConfig.class); + + int phase = 99; + Consumer executorTester = beanName -> + assertThat(context.getBean(beanName, ThreadPoolTaskExecutor.class).getPhase()).isEqualTo(phase); + + executorTester.accept("clientInboundChannelExecutor"); + executorTester.accept("clientOutboundChannelExecutor"); + executorTester.accept("brokerChannelExecutor"); + + assertThat(context.getBean(SimpAnnotationMethodMessageHandler.class).getPhase()).isEqualTo(phase); + assertThat(context.getBean(UserDestinationMessageHandler.class).getPhase()).isEqualTo(phase); + assertThat(context.getBean(ThreadPoolTaskScheduler.class).getPhase()).isEqualTo(phase); + assertThat(context.getBean(WebSocketHandlerMapping.class).getPhase()).isEqualTo(phase); + + assertThat(context.getBean(SubProtocolWebSocketHandler.class).getPhase()).isEqualTo(phase); + assertThat(context.getBean(SimpleBrokerMessageHandler.class).getPhase()).isEqualTo(phase); + } + @Test void webSocketHandlerDecorator() throws Exception { ApplicationContext context = createContext(WebSocketHandlerDecoratorConfig.class); @@ -347,4 +372,30 @@ class WebSocketMessageBrokerConfigurationSupportTests { } } + + @Configuration + static class LifecyclePhaseConfig extends DelegatingWebSocketMessageBrokerConfiguration { + + @Bean + public WebSocketMessageBrokerConfigurer getConfigurer() { + return new WebSocketMessageBrokerConfigurer() { + + @Override + public void registerStompEndpoints(StompEndpointRegistry registry) { + registry.addEndpoint("/broker"); + } + + @Override + public void configureMessageBroker(MessageBrokerRegistry registry) { + registry.enableSimpleBroker(); + } + + @Override + public Integer getPhase() { + return 99; + } + }; + } + } + }