001/*
002 * Copyright 2002-2018 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.socket.config.annotation;
018
019import org.springframework.context.annotation.Bean;
020import org.springframework.lang.Nullable;
021import org.springframework.scheduling.TaskScheduler;
022import org.springframework.scheduling.concurrent.ThreadPoolTaskScheduler;
023import org.springframework.util.Assert;
024import org.springframework.web.servlet.HandlerMapping;
025
026/**
027 * Configuration support for WebSocket request handling.
028 *
029 * @author Rossen Stoyanchev
030 * @since 4.0
031 */
032public class WebSocketConfigurationSupport {
033
034        @Nullable
035        private ServletWebSocketHandlerRegistry handlerRegistry;
036
037        @Nullable
038        private TaskScheduler scheduler;
039
040
041        @Bean
042        public HandlerMapping webSocketHandlerMapping() {
043                ServletWebSocketHandlerRegistry registry = initHandlerRegistry();
044                if (registry.requiresTaskScheduler()) {
045                        TaskScheduler scheduler = defaultSockJsTaskScheduler();
046                        Assert.notNull(scheduler, "Expected default TaskScheduler bean");
047                        registry.setTaskScheduler(scheduler);
048                }
049                return registry.getHandlerMapping();
050        }
051
052        private ServletWebSocketHandlerRegistry initHandlerRegistry() {
053                if (this.handlerRegistry == null) {
054                        this.handlerRegistry = new ServletWebSocketHandlerRegistry();
055                        registerWebSocketHandlers(this.handlerRegistry);
056                }
057                return this.handlerRegistry;
058        }
059
060        protected void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
061        }
062
063        /**
064         * The default TaskScheduler to use if none is registered explicitly via
065         * {@link SockJsServiceRegistration#setTaskScheduler}:
066         * <pre class="code">
067         * &#064;Configuration
068         * &#064;EnableWebSocket
069         * public class WebSocketConfig implements WebSocketConfigurer {
070         *
071         *   public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
072         *     registry.addHandler(myHandler(), "/echo")
073         *             .withSockJS()
074         *             .setTaskScheduler(myScheduler());
075         *   }
076         *
077         *   // ...
078         * }
079         * </pre>
080         */
081        @Bean
082        @Nullable
083        public TaskScheduler defaultSockJsTaskScheduler() {
084                if (initHandlerRegistry().requiresTaskScheduler()) {
085                        ThreadPoolTaskScheduler threadPoolScheduler = new ThreadPoolTaskScheduler();
086                        threadPoolScheduler.setThreadNamePrefix("SockJS-");
087                        threadPoolScheduler.setPoolSize(Runtime.getRuntime().availableProcessors());
088                        threadPoolScheduler.setRemoveOnCancelPolicy(true);
089                        this.scheduler = threadPoolScheduler;
090                }
091                return this.scheduler;
092        }
093
094}