Skip to content

Commit 504b761

Browse files
rstoyanchevrstoyanchev
committed
WebSocketMessageBrokerConfigurer allows to configure Lifecycle phase
Closes gh-32205
1 parent f9ae54d commit 504b761

File tree

10 files changed

+238
-21
lines changed

10 files changed

+238
-21
lines changed

‎spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -120,6 +120,9 @@ public class SimpAnnotationMethodMessageHandler extends AbstractMethodMessageHan
120120
@Nullable
121121
private MessageHeaderInitializer headerInitializer;
122122

123+
@Nullable
124+
private Integer phase;
125+
123126
private volatile boolean running;
124127

125128
private final Object lifecycleMonitor = new Object();
@@ -271,6 +274,20 @@ public MessageHeaderInitializer getHeaderInitializer() {
271274
return this.headerInitializer;
272275
}
273276

277+
/**
278+
* Set the phase that this handler should run in.
279+
* <p>By default, this is {@link SmartLifecycle#DEFAULT_PHASE}.
280+
* @since 6.1.4
281+
*/
282+
public void setPhase(int phase) {
283+
this.phase = phase;
284+
}
285+
286+
@Override
287+
public int getPhase() {
288+
return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase());
289+
}
290+
274291

275292
@Override
276293
public final void start() {

‎spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -75,6 +75,9 @@ public abstract class AbstractBrokerMessageHandler
7575

7676
private boolean autoStartup = true;
7777

78+
@Nullable
79+
private Integer phase;
80+
7881
private volatile boolean running;
7982

8083
private final Object lifecycleMonitor = new Object();
@@ -197,6 +200,20 @@ public boolean isAutoStartup() {
197200
return this.autoStartup;
198201
}
199202

203+
/**
204+
* Set the phase that this handler should run in.
205+
* <p>By default, this is {@link SmartLifecycle#DEFAULT_PHASE}.
206+
* @since 6.1.4
207+
*/
208+
public void setPhase(int phase) {
209+
this.phase = phase;
210+
}
211+
212+
@Override
213+
public int getPhase() {
214+
return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase());
215+
}
216+
200217

201218
@Override
202219
public void start() {

‎spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java

Lines changed: 47 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2022 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -27,6 +27,7 @@
2727
import org.springframework.beans.factory.annotation.Qualifier;
2828
import org.springframework.context.ApplicationContext;
2929
import org.springframework.context.ApplicationContextAware;
30+
import org.springframework.context.SmartLifecycle;
3031
import org.springframework.context.annotation.Bean;
3132
import org.springframework.context.event.SmartApplicationListener;
3233
import org.springframework.core.task.TaskExecutor;
@@ -59,6 +60,7 @@
5960
import org.springframework.messaging.support.ExecutorSubscribableChannel;
6061
import org.springframework.messaging.support.ImmutableMessageChannelInterceptor;
6162
import org.springframework.scheduling.TaskScheduler;
63+
import org.springframework.scheduling.concurrent.ExecutorConfigurationSupport;
6264
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
6365
import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
6466
import org.springframework.util.Assert;
@@ -132,6 +134,9 @@ public abstract class AbstractMessageBrokerConfiguration implements ApplicationC
132134
@Nullable
133135
private MessageBrokerRegistry brokerRegistry;
134136

137+
@Nullable
138+
private Integer phase;
139+
135140

136141
/**
137142
* Protected constructor.
@@ -166,8 +171,12 @@ public AbstractSubscribableChannel clientInboundChannel(
166171

167172
@Bean
168173
public TaskExecutor clientInboundChannelExecutor() {
169-
return getTaskExecutor(getClientInboundChannelRegistration(),
170-
"clientInboundChannel-", this::defaultTaskExecutor);
174+
ChannelRegistration registration = getClientInboundChannelRegistration();
175+
TaskExecutor executor = getTaskExecutor(registration, "clientInboundChannel-", this::defaultTaskExecutor);
176+
if (executor instanceof ExecutorConfigurationSupport executorSupport) {
177+
executorSupport.setPhase(getPhase());
178+
}
179+
return executor;
171180
}
172181

173182
protected final ChannelRegistration getClientInboundChannelRegistration() {
@@ -180,6 +189,17 @@ protected final ChannelRegistration getClientInboundChannelRegistration() {
180189
return this.clientInboundChannelRegistration;
181190
}
182191

192+
protected final int getPhase() {
193+
if (this.phase == null) {
194+
this.phase = initPhase();
195+
}
196+
return this.phase;
197+
}
198+
199+
protected int initPhase() {
200+
return SmartLifecycle.DEFAULT_PHASE;
201+
}
202+
183203
/**
184204
* A hook for subclasses to customize the message channel for inbound messages
185205
* from WebSocket clients.
@@ -193,17 +213,21 @@ public AbstractSubscribableChannel clientOutboundChannel(
193213

194214
ExecutorSubscribableChannel channel = new ExecutorSubscribableChannel(executor);
195215
channel.setLogger(SimpLogging.forLog(channel.getLogger()));
196-
ChannelRegistration reg = getClientOutboundChannelRegistration();
197-
if (reg.hasInterceptors()) {
198-
channel.setInterceptors(reg.getInterceptors());
216+
ChannelRegistration registration = getClientOutboundChannelRegistration();
217+
if (registration.hasInterceptors()) {
218+
channel.setInterceptors(registration.getInterceptors());
199219
}
200220
return channel;
201221
}
202222

203223
@Bean
204224
public TaskExecutor clientOutboundChannelExecutor() {
205-
return getTaskExecutor(getClientOutboundChannelRegistration(),
206-
"clientOutboundChannel-", this::defaultTaskExecutor);
225+
ChannelRegistration registration = getClientOutboundChannelRegistration();
226+
TaskExecutor executor = getTaskExecutor(registration, "clientOutboundChannel-", this::defaultTaskExecutor);
227+
if (executor instanceof ExecutorConfigurationSupport executorSupport) {
228+
executorSupport.setPhase(getPhase());
229+
}
230+
return executor;
207231
}
208232

209233
protected final ChannelRegistration getClientOutboundChannelRegistration() {
@@ -244,14 +268,18 @@ public TaskExecutor brokerChannelExecutor(
244268

245269
MessageBrokerRegistry registry = getBrokerRegistry(clientInboundChannel, clientOutboundChannel);
246270
ChannelRegistration registration = registry.getBrokerChannelRegistration();
247-
return getTaskExecutor(registration, "brokerChannel-", () -> {
271+
TaskExecutor executor = getTaskExecutor(registration, "brokerChannel-", () -> {
248272
// Should never be used
249273
ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
250274
threadPoolTaskExecutor.setCorePoolSize(0);
251275
threadPoolTaskExecutor.setMaxPoolSize(1);
252276
threadPoolTaskExecutor.setQueueCapacity(0);
253277
return threadPoolTaskExecutor;
254278
});
279+
if (executor instanceof ExecutorConfigurationSupport executorSupport) {
280+
executorSupport.setPhase(getPhase());
281+
}
282+
return executor;
255283
}
256284

257285
private TaskExecutor defaultTaskExecutor() {
@@ -316,6 +344,7 @@ public SimpAnnotationMethodMessageHandler simpAnnotationMethodMessageHandler(
316344
handler.setDestinationPrefixes(brokerRegistry.getApplicationDestinationPrefixes());
317345
handler.setMessageConverter(brokerMessageConverter);
318346
handler.setValidator(simpValidator());
347+
handler.setPhase(getPhase());
319348

320349
List<HandlerMethodArgumentResolver> argumentResolvers = new ArrayList<>();
321350
addArgumentResolvers(argumentResolvers);
@@ -329,6 +358,7 @@ public SimpAnnotationMethodMessageHandler simpAnnotationMethodMessageHandler(
329358
if (pathMatcher != null) {
330359
handler.setPathMatcher(pathMatcher);
331360
}
361+
332362
return handler;
333363
}
334364

@@ -342,8 +372,11 @@ protected SimpAnnotationMethodMessageHandler createAnnotationMethodMessageHandle
342372
AbstractSubscribableChannel clientInboundChannel, AbstractSubscribableChannel clientOutboundChannel,
343373
SimpMessagingTemplate brokerMessagingTemplate) {
344374

345-
return new SimpAnnotationMethodMessageHandler(
375+
SimpAnnotationMethodMessageHandler handler = new SimpAnnotationMethodMessageHandler(
346376
clientInboundChannel, clientOutboundChannel, brokerMessagingTemplate);
377+
378+
handler.setPhase(getPhase());
379+
return handler;
347380
}
348381

349382
protected void addArgumentResolvers(List<HandlerMethodArgumentResolver> argumentResolvers) {
@@ -364,6 +397,7 @@ public AbstractBrokerMessageHandler simpleBrokerMessageHandler(
364397
return null;
365398
}
366399
updateUserDestinationResolver(handler, userDestinationResolver, registry.getUserDestinationPrefix());
400+
handler.setPhase(getPhase());
367401
return handler;
368402
}
369403

@@ -403,6 +437,7 @@ public AbstractBrokerMessageHandler stompBrokerRelayMessageHandler(
403437
}
404438
handler.setSystemSubscriptions(subscriptions);
405439
updateUserDestinationResolver(handler, userDestinationResolver, registry.getUserDestinationPrefix());
440+
handler.setPhase(getPhase());
406441
return handler;
407442
}
408443

@@ -419,6 +454,7 @@ public UserDestinationMessageHandler userDestinationMessageHandler(
419454
if (destination != null) {
420455
handler.setBroadcastDestination(destination);
421456
}
457+
handler.setPhase(getPhase());
422458
return handler;
423459
}
424460

@@ -446,6 +482,7 @@ public TaskScheduler messageBrokerTaskScheduler() {
446482
scheduler.setThreadNamePrefix("MessageBroker-");
447483
scheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
448484
scheduler.setRemoveOnCancelPolicy(true);
485+
scheduler.setPhase(getPhase());
449486
return scheduler;
450487
}
451488

‎spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -77,6 +77,9 @@ public class UserDestinationMessageHandler implements MessageHandler, SmartLifec
7777

7878
private volatile boolean running;
7979

80+
@Nullable
81+
private Integer phase;
82+
8083
private final Object lifecycleMonitor = new Object();
8184

8285

@@ -154,6 +157,20 @@ public MessageHeaderInitializer getHeaderInitializer() {
154157
return this.headerInitializer;
155158
}
156159

160+
/**
161+
* Set the phase that this handler should run in.
162+
* <p>By default, this is {@link SmartLifecycle#DEFAULT_PHASE}.
163+
* @since 6.1.4
164+
*/
165+
public void setPhase(int phase) {
166+
this.phase = phase;
167+
}
168+
169+
@Override
170+
public int getPhase() {
171+
return (this.phase != null ? this.phase : SmartLifecycle.super.getPhase());
172+
}
173+
157174

158175
@Override
159176
public final void start() {

‎spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/DelegatingWebSocketMessageBrokerConfiguration.java

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2020 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -114,4 +114,15 @@ protected void configureMessageBroker(MessageBrokerRegistry registry) {
114114
}
115115
}
116116

117+
@Override
118+
protected int initPhase() {
119+
for (WebSocketMessageBrokerConfigurer configurer : this.configurers) {
120+
Integer phase = configurer.getPhase();
121+
if (phase != null) {
122+
return phase;
123+
}
124+
}
125+
return super.initPhase();
126+
}
127+
117128
}

‎spring-websocket/src/main/java/org/springframework/web/socket/config/annotation/WebSocketMessageBrokerConfigurationSupport.java

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2023 the original author or authors.
2+
* Copyright 2002-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -35,12 +35,14 @@
3535
import org.springframework.messaging.support.AbstractSubscribableChannel;
3636
import org.springframework.scheduling.TaskScheduler;
3737
import org.springframework.web.servlet.HandlerMapping;
38+
import org.springframework.web.servlet.handler.AbstractHandlerMapping;
3839
import org.springframework.web.socket.WebSocketHandler;
3940
import org.springframework.web.socket.config.WebSocketMessageBrokerStats;
4041
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
4142
import org.springframework.web.socket.messaging.DefaultSimpUserRegistry;
4243
import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
4344
import org.springframework.web.socket.messaging.WebSocketAnnotationMethodMessageHandler;
45+
import org.springframework.web.socket.server.support.WebSocketHandlerMapping;
4446

4547
/**
4648
* Extends {@link AbstractMessageBrokerConfiguration} and adds configuration for
@@ -66,8 +68,11 @@ protected SimpAnnotationMethodMessageHandler createAnnotationMethodMessageHandle
6668
AbstractSubscribableChannel clientInboundChannel,AbstractSubscribableChannel clientOutboundChannel,
6769
SimpMessagingTemplate brokerMessagingTemplate) {
6870

69-
return new WebSocketAnnotationMethodMessageHandler(
71+
WebSocketAnnotationMethodMessageHandler handler = new WebSocketAnnotationMethodMessageHandler(
7072
clientInboundChannel, clientOutboundChannel, brokerMessagingTemplate);
73+
74+
handler.setPhase(getPhase());
75+
return handler;
7176
}
7277

7378
@Override
@@ -93,14 +98,22 @@ public HandlerMapping stompWebSocketHandlerMapping(
9398
}
9499
registerStompEndpoints(registry);
95100
OrderedMessageChannelDecorator.configureInterceptor(clientInboundChannel, registry.isPreserveReceiveOrder());
96-
return registry.getHandlerMapping();
101+
AbstractHandlerMapping handlerMapping = registry.getHandlerMapping();
102+
if (handlerMapping instanceof WebSocketHandlerMapping webSocketMapping) {
103+
webSocketMapping.setPhase(getPhase());
104+
}
105+
return handlerMapping;
97106
}
98107

99108
@Bean
100109
public WebSocketHandler subProtocolWebSocketHandler(
101110
AbstractSubscribableChannel clientInboundChannel, AbstractSubscribableChannel clientOutboundChannel) {
102111

103-
return new SubProtocolWebSocketHandler(clientInboundChannel, clientOutboundChannel);
112+
SubProtocolWebSocketHandler handler =
113+
new SubProtocolWebSocketHandler(clientInboundChannel, clientOutboundChannel);
114+
115+
handler.setPhase(getPhase());
116+
return handler;
104117
}
105118

106119
protected WebSocketHandler decorateWebSocketHandler(WebSocketHandler handler) {

0 commit comments

Comments
 (0)