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

Author: rstoyanchev

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

@@ -59,22 +59,22 @@
 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;
 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;
 
 /**
  * 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 @@ private RuntimeBeanReference registerUserRegistry(Element element, ParserContext
 	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.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 StompWebSocketEndpointRegistration addEndpoint(String... paths) {
 	}
 
 	/**
-	 * 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 @@ protected int getOrder() {
 	}
 
 	/**
-	 * 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 WebMvcStompEndpointRegistry setErrorHandler(StompSubProtocolErrorHandler
 	}
 
 	/**
-	 * 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 AbstractHandlerMapping getHandlerMapping() {
 				}
 			}
 		}
-		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.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.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 @@
  * @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 Object createWebSocket(ServletUpgradeRequest request, ServletUpgradeRespo
 				// 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 Object createWebSocket(UpgradeRequest request, UpgradeResponse response)
 				return container.getHandler();
 			}
 		});
-		try {
-			this.factory.init();
-		}
-		catch (Exception ex) {
-			throw new IllegalStateException("Unable to initialize Jetty WebSocketServerFactory", ex);
-		}
 	}
 
 
@@ -129,6 +128,33 @@ private List<WebSocketExtension> getWebSocketExtensions() {
 		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 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 @@
  * @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 String[] getSupportedProtocols() {
 		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();
+	}
+
+}