Spring/spring-framework
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

WebSocketMessageBrokerConfigurer allows to configure Lifecycle phase

Browse Source 
Closes gh-32205
This commit is contained in:
rstoyanchev
2024-02-14 16:26:18 +00:00
parent f9ae54d91e
commit 504b7619bd
10 changed files with 238 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/SimpAnnotationMethodMessageHandler.java
View File

@@ -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.
* <p>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() {

19
spring-messaging/src/main/java/org/springframework/messaging/simp/broker/AbstractBrokerMessageHandler.java
View File

@@ -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.
* <p>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() {

57
spring-messaging/src/main/java/org/springframework/messaging/simp/config/AbstractMessageBrokerConfiguration.java
View File

@@ -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<HandlerMethodArgumentResolver> 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<HandlerMethodArgumentResolver> 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;
}

19
spring-messaging/src/main/java/org/springframework/messaging/simp/user/UserDestinationMessageHandler.java
View File

@@ -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.
* <p>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() {