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 * @Configuration 068 * @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}