JettyRequestUpgradeStrategy implements Lifecycle

After this change JettyRequestUpgradeStrategy implements Lifecyle,
which is used to init and cleanup the Jetty WebSocketServerFactory.

Since a RequestUpgradeStrategy is typically created reflectively
within DefaultHandshakeHandler, the Lifecycle events are propagated
from the top, i.e. the Spring MVC HandlerMapping through the
WebSocket/SockJsHttpRequestHandler.

Issue: SPR-13140

spring-websocket/src/main/java/org/springframework/web/socket/config/MessageBrokerBeanDefinitionParser.java

@@ -59,7 +59,6 @@ import org.springframework.util.ClassUtils;
 import org.springframework.util.MimeTypeUtils;
 import org.springframework.util.StringUtils;
 import org.springframework.util.xml.DomUtils;
-import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
 import org.springframework.web.socket.messaging.DefaultSimpUserRegistry;
@@ -67,6 +66,7 @@ import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
 import org.springframework.web.socket.messaging.WebSocketAnnotationMethodMessageHandler;
 import org.springframework.web.socket.server.support.OriginHandshakeInterceptor;
+import org.springframework.web.socket.server.support.WebSocketHandlerMapping;
 import org.springframework.web.socket.server.support.WebSocketHttpRequestHandler;
 import org.springframework.web.socket.sockjs.support.SockJsHttpRequestHandler;
 
@@ -74,7 +74,7 @@ import org.springframework.web.socket.sockjs.support.SockJsHttpRequestHandler;
  * A {@link org.springframework.beans.factory.xml.BeanDefinitionParser} that provides
  * the configuration for the {@code <websocket:message-broker/>} XML namespace element.
  *
- * <p>Registers a Spring MVC {@link org.springframework.web.servlet.handler.SimpleUrlHandlerMapping}
+ * <p>Registers a Spring MVC {@link org.springframework.web.servlet.HandlerMapping}
  * with order 1 to map HTTP WebSocket handshake requests from STOMP/WebSocket clients.
  *
  * <p>Registers the following {@link org.springframework.messaging.MessageChannel}s:
@@ -188,7 +188,7 @@ class MessageBrokerBeanDefinitionParser implements BeanDefinitionParser {
 	private ManagedMap<String, Object> registerHandlerMapping(Element element,
 			ParserContext context, Object source) {
 
-		RootBeanDefinition handlerMappingDef = new RootBeanDefinition(SimpleUrlHandlerMapping.class);
+		RootBeanDefinition handlerMappingDef = new RootBeanDefinition(WebSocketHandlerMapping.class);
 
 		String orderAttribute = element.getAttribute("order");
 		int order = orderAttribute.isEmpty() ? DEFAULT_MAPPING_ORDER : Integer.valueOf(orderAttribute);

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

@@ -27,17 +27,17 @@ import org.springframework.util.Assert;
 import org.springframework.util.MultiValueMap;
 import org.springframework.web.HttpRequestHandler;
 import org.springframework.web.servlet.handler.AbstractHandlerMapping;
-import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
 import org.springframework.web.socket.WebSocketHandler;
 import org.springframework.web.socket.handler.WebSocketHandlerDecorator;
 import org.springframework.web.socket.messaging.StompSubProtocolErrorHandler;
 import org.springframework.web.socket.messaging.StompSubProtocolHandler;
 import org.springframework.web.socket.messaging.SubProtocolWebSocketHandler;
+import org.springframework.web.socket.server.support.WebSocketHandlerMapping;
 import org.springframework.web.util.UrlPathHelper;
 
 /**
  * A registry for STOMP over WebSocket endpoints that maps the endpoints with a
- * {@link SimpleUrlHandlerMapping} for use in Spring MVC.
+ * {@link org.springframework.web.servlet.HandlerMapping} for use in Spring MVC.
  *
  * @author Rossen Stoyanchev
  * @author Artem Bilan
@@ -111,8 +111,9 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 	}
 
 	/**
-	 * Set the order for the resulting {@link SimpleUrlHandlerMapping} relative to
-	 * other handler mappings configured in Spring MVC.
+	 * Set the order for the resulting
+	 * {@link org.springframework.web.servlet.HandlerMapping}
+	 * relative to other handler mappings configured in Spring MVC.
 	 * <p>The default value is 1.
 	 */
 	@Override
@@ -125,7 +126,7 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 	}
 
 	/**
-	 * Set the UrlPathHelper to configure on the {@code SimpleUrlHandlerMapping}
+	 * Set the UrlPathHelper to configure on the {@code HandlerMapping}
 	 * used to map handshake requests.
 	 */
 	@Override
@@ -144,7 +145,8 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 	}
 
 	/**
-	 * Return a handler mapping with the mapped ViewControllers; or {@code null} in case of no registrations.
+	 * Return a handler mapping with the mapped ViewControllers; or {@code null}
+	 * in case of no registrations.
 	 */
 	public AbstractHandlerMapping getHandlerMapping() {
 		Map<String, Object> urlMap = new LinkedHashMap<String, Object>();
@@ -156,7 +158,7 @@ public class WebMvcStompEndpointRegistry implements StompEndpointRegistry {
 				}
 			}
 		}
-		SimpleUrlHandlerMapping hm = new SimpleUrlHandlerMapping();
+		WebSocketHandlerMapping hm = new WebSocketHandlerMapping();
 		hm.setUrlMap(urlMap);
 		hm.setOrder(this.order);
 		if (this.urlPathHelper != null) {

spring-websocket/src/main/java/org/springframework/web/socket/server/jetty/JettyRequestUpgradeStrategy.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 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.
@@ -21,6 +21,7 @@ import java.security.Principal;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
@@ -34,6 +35,7 @@ import org.eclipse.jetty.websocket.servlet.ServletUpgradeRequest;
 import org.eclipse.jetty.websocket.servlet.ServletUpgradeResponse;
 import org.eclipse.jetty.websocket.servlet.WebSocketCreator;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.core.NamedThreadLocal;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
@@ -57,7 +59,7 @@ import org.springframework.web.socket.server.RequestUpgradeStrategy;
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
+public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy, Lifecycle {
 
 	private static final ThreadLocal<WebSocketHandlerContainer> wsContainerHolder =
 			new NamedThreadLocal<WebSocketHandlerContainer>("WebSocket Handler Container");
@@ -67,6 +69,8 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 
 	private volatile List<WebSocketExtension> supportedExtensions;
 
+	private volatile boolean running = false;
+
 
 	/**
 	 * Default constructor that creates {@link WebSocketServerFactory} through its default
@@ -90,6 +94,7 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 				// Cast to avoid infinite recursion
 				return createWebSocket((UpgradeRequest) request, (UpgradeResponse) response);
 			}
+
 			// For Jetty 9.0.x
 			public Object createWebSocket(UpgradeRequest request, UpgradeResponse response) {
 				WebSocketHandlerContainer container = wsContainerHolder.get();
@@ -99,12 +104,6 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 				return container.getHandler();
 			}
 		});
-		try {
-			this.factory.init();
-		}
-		catch (Exception ex) {
-			throw new IllegalStateException("Unable to initialize Jetty WebSocketServerFactory", ex);
-		}
 	}
 
 
@@ -129,6 +128,33 @@ public class JettyRequestUpgradeStrategy implements RequestUpgradeStrategy {
 		return result;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			try {
+				this.factory.init();
+			}
+			catch (Exception ex) {
+				throw new IllegalStateException("Unable to initialize Jetty WebSocketServerFactory", ex);
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			this.factory.cleanup();
+		}
+	}
+
 	@Override
 	public void upgrade(ServerHttpRequest request, ServerHttpResponse response,
 			String selectedProtocol, List<WebSocketExtension> selectedExtensions, Principal user,

spring-websocket/src/main/java/org/springframework/web/socket/server/support/DefaultHandshakeHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 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.
@@ -28,6 +28,7 @@ import java.util.Map;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.server.ServerHttpRequest;
@@ -59,7 +60,7 @@ import org.springframework.web.socket.server.RequestUpgradeStrategy;
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class DefaultHandshakeHandler implements HandshakeHandler {
+public class DefaultHandshakeHandler implements HandshakeHandler, Lifecycle {
 
 	private static final Charset UTF8_CHARSET = Charset.forName("UTF-8");
 
@@ -86,6 +87,8 @@ public class DefaultHandshakeHandler implements HandshakeHandler {
 
 	private final List<String> supportedProtocols = new ArrayList<String>();
 
+	private volatile boolean running = false;
+
 
 	/**
 	 * Default constructor that autodetects and instantiates a
@@ -161,6 +164,31 @@ public class DefaultHandshakeHandler implements HandshakeHandler {
 		return this.supportedProtocols.toArray(new String[this.supportedProtocols.size()]);
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			if (this.requestUpgradeStrategy instanceof Lifecycle) {
+				((Lifecycle) this.requestUpgradeStrategy).start();
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			if (this.requestUpgradeStrategy instanceof Lifecycle) {
+				((Lifecycle) this.requestUpgradeStrategy).stop();
+			}
+		}
+	}
+
 
 	@Override
 	public final boolean doHandshake(ServerHttpRequest request, ServerHttpResponse response,

spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHandlerMapping.java

@@ -0,0 +1,81 @@
+/*
+ * Copyright 2002-2015 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.web.socket.server.support;
+
+import org.springframework.context.Lifecycle;
+import org.springframework.context.SmartLifecycle;
+import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
+
+/**
+ * An extension of {@link SimpleUrlHandlerMapping} that is also a
+ * {@link SmartLifecycle} container and propagates start and stop calls to any
+ * handlers that implement {@link Lifecycle}. The handlers are typically expected
+ * to be {@code WebSocketHttpRequestHandler} or {@code SockJsHttpRequestHandler}.
+ *
+ * @author Rossen Stoyanchev
+ * @since 4.2
+ */
+public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping implements SmartLifecycle {
+
+	private volatile boolean running = false;
+
+
+	@Override
+	public boolean isAutoStartup() {
+		return true;
+	}
+
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public int getPhase() {
+		return Integer.MAX_VALUE;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			for (Object handler : getUrlMap().values()) {
+				if (handler instanceof Lifecycle) {
+					((Lifecycle) handler).start();
+				}
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			for (Object handler : getUrlMap().values()) {
+				if (handler instanceof Lifecycle) {
+					((Lifecycle) handler).stop();
+				}
+			}
+		}
+	}
+
+	@Override
+	public void stop(Runnable callback) {
+		stop();
+		callback.run();
+	}
+
+}

spring-websocket/src/main/java/org/springframework/web/socket/server/support/WebSocketHttpRequestHandler.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2014 the original author or authors.
+ * Copyright 2002-2015 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.
@@ -28,6 +28,7 @@ import javax.servlet.http.HttpServletResponse;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
 import org.springframework.http.server.ServletServerHttpRequest;
@@ -52,7 +53,7 @@ import org.springframework.web.socket.server.HandshakeInterceptor;
  * @author Rossen Stoyanchev
  * @since 4.0
  */
-public class WebSocketHttpRequestHandler implements HttpRequestHandler {
+public class WebSocketHttpRequestHandler implements HttpRequestHandler, Lifecycle {
 
 	private final Log logger = LogFactory.getLog(WebSocketHttpRequestHandler.class);
 
@@ -62,6 +63,8 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler {
 
 	private final List<HandshakeInterceptor> interceptors = new ArrayList<HandshakeInterceptor>();
 
+	private volatile boolean running = false;
+
 
 	public WebSocketHttpRequestHandler(WebSocketHandler wsHandler) {
 		this(wsHandler, new DefaultHandshakeHandler());
@@ -106,6 +109,32 @@ public class WebSocketHttpRequestHandler implements HttpRequestHandler {
 		return this.interceptors;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			if (this.handshakeHandler instanceof Lifecycle) {
+				((Lifecycle) this.handshakeHandler).start();
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			if (this.handshakeHandler instanceof Lifecycle) {
+				((Lifecycle) this.handshakeHandler).stop();
+			}
+		}
+	}
+
+
 	@Override
 	public void handleRequest(HttpServletRequest servletRequest, HttpServletResponse servletResponse)
 			throws ServletException, IOException {

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/support/SockJsHttpRequestHandler.java

@@ -17,10 +17,12 @@
 package org.springframework.web.socket.sockjs.support;
 
 import java.io.IOException;
+
 import javax.servlet.ServletException;
 import javax.servlet.http.HttpServletRequest;
 import javax.servlet.http.HttpServletResponse;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
 import org.springframework.http.server.ServletServerHttpRequest;
@@ -44,7 +46,8 @@ import org.springframework.web.socket.sockjs.SockJsService;
  * @author Sebastien Deleuze
  * @since 4.0
  */
-public class SockJsHttpRequestHandler implements HttpRequestHandler, CorsConfigurationSource {
+public class SockJsHttpRequestHandler
+		implements HttpRequestHandler, CorsConfigurationSource, Lifecycle {
 
 	// No logging: HTTP transports too verbose and we don't know enough to log anything of value
 
@@ -52,6 +55,8 @@ public class SockJsHttpRequestHandler implements HttpRequestHandler, CorsConfigu
 
 	private final WebSocketHandler webSocketHandler;
 
+	private volatile boolean running = false;
+
 
 	/**
 	 * Create a new SockJsHttpRequestHandler.
@@ -81,6 +86,31 @@ public class SockJsHttpRequestHandler implements HttpRequestHandler, CorsConfigu
 		return this.webSocketHandler;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			if (this.sockJsService instanceof Lifecycle) {
+				((Lifecycle) this.sockJsService).start();
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			if (this.sockJsService instanceof Lifecycle) {
+				((Lifecycle) this.sockJsService).stop();
+			}
+		}
+	}
+
 
 	@Override
 	public void handleRequest(HttpServletRequest servletRequest, HttpServletResponse servletResponse)

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/TransportHandlingSockJsService.java

@@ -27,6 +27,7 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledFuture;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.server.ServerHttpRequest;
@@ -59,7 +60,8 @@ import org.springframework.web.socket.sockjs.support.AbstractSockJsService;
  * @author Sebastien Deleuze
  * @since 4.0
  */
-public class TransportHandlingSockJsService extends AbstractSockJsService implements SockJsServiceConfig {
+public class TransportHandlingSockJsService extends AbstractSockJsService
+		implements SockJsServiceConfig, Lifecycle {
 
 	private static final boolean jackson2Present = ClassUtils.isPresent(
 			"com.fasterxml.jackson.databind.ObjectMapper", TransportHandlingSockJsService.class.getClassLoader());
@@ -75,6 +77,8 @@ public class TransportHandlingSockJsService extends AbstractSockJsService implem
 
 	private ScheduledFuture<?> sessionCleanupTask;
 
+	private boolean running;
+
 
 	/**
 	 * Create a TransportHandlingSockJsService with given {@link TransportHandler handler} types.
@@ -150,6 +154,35 @@ public class TransportHandlingSockJsService extends AbstractSockJsService implem
 		return this.interceptors;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			for (TransportHandler handler : this.handlers.values()) {
+				if (handler instanceof Lifecycle) {
+					((Lifecycle) handler).start();
+				}
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			for (TransportHandler handler : this.handlers.values()) {
+				if (handler instanceof Lifecycle) {
+					((Lifecycle) handler).stop();
+				}
+			}
+		}
+	}
+
 
 	@Override
 	protected void handleRawWebSocketRequest(ServerHttpRequest request, ServerHttpResponse response,

spring-websocket/src/main/java/org/springframework/web/socket/sockjs/transport/handler/WebSocketTransportHandler.java

@@ -18,6 +18,7 @@ package org.springframework.web.socket.sockjs.transport.handler;
 
 import java.util.Map;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.server.ServerHttpRequest;
 import org.springframework.http.server.ServerHttpResponse;
 import org.springframework.util.Assert;
@@ -45,10 +46,12 @@ import org.springframework.web.socket.sockjs.transport.session.WebSocketServerSo
  * @since 4.0
  */
 public class WebSocketTransportHandler extends AbstractTransportHandler
-		implements SockJsSessionFactory, HandshakeHandler {
+		implements SockJsSessionFactory, HandshakeHandler, Lifecycle {
 
 	private final HandshakeHandler handshakeHandler;
 
+	private boolean running;
+
 
 	public WebSocketTransportHandler(HandshakeHandler handshakeHandler) {
 		Assert.notNull(handshakeHandler, "handshakeHandler must not be null");
@@ -65,6 +68,31 @@ public class WebSocketTransportHandler extends AbstractTransportHandler
 		return this.handshakeHandler;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			if (this.handshakeHandler instanceof Lifecycle) {
+				((Lifecycle) this.handshakeHandler).start();
+			}
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			if (this.handshakeHandler instanceof Lifecycle) {
+				((Lifecycle) this.handshakeHandler).stop();
+			}
+		}
+	}
+
 	@Override
 	public AbstractSockJsSession createSession(String id, WebSocketHandler handler, Map<String, Object> attrs) {
 		return new WebSocketServerSockJsSession(id, getServiceConfig(), handler, attrs);
@@ -91,5 +119,4 @@ public class WebSocketTransportHandler extends AbstractTransportHandler
 
 		return this.handshakeHandler.doHandshake(request, response, handler, attributes);
 	}
-
 }