Hide default SockJS TaskScheduler for @EnableWebSocket

Closes gh-27903
This commit is contained in:
rstoyanchev 2022-01-10 16:07:14 +00:00
parent d57bc176f2
commit 368201975a
1 changed files with 53 additions and 19 deletions

View File

@ -1,5 +1,5 @@
/* /*
* Copyright 2002-2020 the original author or authors. * Copyright 2002-2022 the original author or authors.
* *
* Licensed under the Apache License, Version 2.0 (the "License"); * Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License. * you may not use this file except in compliance with the License.
@ -16,6 +16,8 @@
package org.springframework.web.socket.config.annotation; package org.springframework.web.socket.config.annotation;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Bean;
import org.springframework.lang.Nullable; import org.springframework.lang.Nullable;
import org.springframework.scheduling.TaskScheduler; import org.springframework.scheduling.TaskScheduler;
@ -35,16 +37,13 @@ public class WebSocketConfigurationSupport {
@Nullable @Nullable
private ServletWebSocketHandlerRegistry handlerRegistry; private ServletWebSocketHandlerRegistry handlerRegistry;
@Nullable
private TaskScheduler scheduler;
@Bean @Bean
public HandlerMapping webSocketHandlerMapping(@Nullable TaskScheduler defaultSockJsTaskScheduler) { public HandlerMapping webSocketHandlerMapping(DefaultSockJsSchedulerContainer schedulerContainer) {
ServletWebSocketHandlerRegistry registry = initHandlerRegistry(); ServletWebSocketHandlerRegistry registry = initHandlerRegistry();
if (registry.requiresTaskScheduler()) { if (registry.requiresTaskScheduler()) {
TaskScheduler scheduler = defaultSockJsTaskScheduler; TaskScheduler scheduler = schedulerContainer.getScheduler();
Assert.notNull(scheduler, "Expected default TaskScheduler bean"); Assert.notNull(scheduler, "TaskScheduler is required but not initialized");
registry.setTaskScheduler(scheduler); registry.setTaskScheduler(scheduler);
} }
return registry.getHandlerMapping(); return registry.getHandlerMapping();
@ -62,8 +61,9 @@ public class WebSocketConfigurationSupport {
} }
/** /**
* The default TaskScheduler to use if none is registered explicitly via * A container of the default TaskScheduler to use if none was registered
* {@link SockJsServiceRegistration#setTaskScheduler}: * explicitly via {@link SockJsServiceRegistration#setTaskScheduler} as
* follows:
* <pre class="code"> * <pre class="code">
* &#064;Configuration * &#064;Configuration
* &#064;EnableWebSocket * &#064;EnableWebSocket
@ -80,16 +80,50 @@ public class WebSocketConfigurationSupport {
* </pre> * </pre>
*/ */
@Bean @Bean
@Nullable public DefaultSockJsSchedulerContainer defaultSockJsSchedulerContainer() {
public TaskScheduler defaultSockJsTaskScheduler() { return (initHandlerRegistry().requiresTaskScheduler() ?
if (initHandlerRegistry().requiresTaskScheduler()) { new DefaultSockJsSchedulerContainer(initDefaultSockJsScheduler()) :
ThreadPoolTaskScheduler threadPoolScheduler = new ThreadPoolTaskScheduler(); new DefaultSockJsSchedulerContainer(null));
threadPoolScheduler.setThreadNamePrefix("SockJS-");
threadPoolScheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
threadPoolScheduler.setRemoveOnCancelPolicy(true);
this.scheduler = threadPoolScheduler;
} }
private ThreadPoolTaskScheduler initDefaultSockJsScheduler() {
ThreadPoolTaskScheduler scheduler = new ThreadPoolTaskScheduler();
scheduler.setThreadNamePrefix("SockJS-");
scheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
scheduler.setRemoveOnCancelPolicy(true);
return scheduler;
}
private static class DefaultSockJsSchedulerContainer implements InitializingBean, DisposableBean {
@Nullable
private final ThreadPoolTaskScheduler scheduler;
DefaultSockJsSchedulerContainer(@Nullable ThreadPoolTaskScheduler scheduler) {
this.scheduler = scheduler;
}
@Nullable
public ThreadPoolTaskScheduler getScheduler() {
return this.scheduler; return this.scheduler;
} }
@Override
public void afterPropertiesSet() throws Exception {
if (this.scheduler != null) {
this.scheduler.afterPropertiesSet();
}
}
@Override
public void destroy() throws Exception {
if (this.scheduler != null) {
this.scheduler.destroy();
}
}
}
} }