From 1b0e95d7d82435d7c532b597eab7cf86265fdbe1 Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Wed, 7 Dec 2016 17:03:56 -0500
Subject: [PATCH] Support for RequestUpgradeStrategy + Lifecycle

Issue: SPR-15527
---
 .../support/HandshakeWebSocketService.java    | 38 ++++++++++++++++++-
 ...tractWebSocketHandlerIntegrationTests.java | 15 +++++---
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java b/spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java
index b96d813963..b32143ea50 100644
--- a/spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java
+++ b/spring-web-reactive/src/main/java/org/springframework/web/reactive/socket/server/support/HandshakeWebSocketService.java
@@ -22,6 +22,7 @@ import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import reactor.core.publisher.Mono;
 
+import org.springframework.context.Lifecycle;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
 import org.springframework.http.server.reactive.ServerHttpRequest;
@@ -43,7 +44,7 @@ import org.springframework.web.server.ServerWebExchange;
  * @author Rossen Stoyanchev
  * @since 5.0
  */
-public class HandshakeWebSocketService implements WebSocketService {
+public class HandshakeWebSocketService implements WebSocketService, Lifecycle {
 
 	private static final String SEC_WEBSOCKET_KEY = "Sec-WebSocket-Key";
 
@@ -58,6 +59,8 @@ public class HandshakeWebSocketService implements WebSocketService {
 
 	private final RequestUpgradeStrategy upgradeStrategy;
 
+	private volatile boolean running = false;
+
 
 	/**
 	 * Default constructor automatic, classpath detection based discovery of the
@@ -104,6 +107,39 @@ public class HandshakeWebSocketService implements WebSocketService {
 		return this.upgradeStrategy;
 	}
 
+	@Override
+	public boolean isRunning() {
+		return this.running;
+	}
+
+	@Override
+	public void start() {
+		if (!isRunning()) {
+			this.running = true;
+			doStart();
+		}
+	}
+
+	protected void doStart() {
+		if (getUpgradeStrategy() instanceof Lifecycle) {
+			((Lifecycle) getUpgradeStrategy()).start();
+		}
+	}
+
+	@Override
+	public void stop() {
+		if (isRunning()) {
+			this.running = false;
+			doStop();
+		}
+	}
+
+	protected void doStop() {
+		if (getUpgradeStrategy() instanceof Lifecycle) {
+			((Lifecycle) getUpgradeStrategy()).stop();
+		}
+	}
+
 
 	@Override
 	public Mono<Void> handleRequest(ServerWebExchange exchange, WebSocketHandler webSocketHandler) {
diff --git a/spring-web-reactive/src/test/java/org/springframework/web/reactive/socket/server/AbstractWebSocketHandlerIntegrationTests.java b/spring-web-reactive/src/test/java/org/springframework/web/reactive/socket/server/AbstractWebSocketHandlerIntegrationTests.java
index 681864cfd6..2489211ba0 100644
--- a/spring-web-reactive/src/test/java/org/springframework/web/reactive/socket/server/AbstractWebSocketHandlerIntegrationTests.java
+++ b/spring-web-reactive/src/test/java/org/springframework/web/reactive/socket/server/AbstractWebSocketHandlerIntegrationTests.java
@@ -104,12 +104,15 @@ public abstract class AbstractWebSocketHandlerIntegrationTests {
 
 		@Bean
 		public WebSocketHandlerAdapter handlerAdapter() {
-			RequestUpgradeStrategy strategy = createUpgradeStrategy();
-			WebSocketService service = new HandshakeWebSocketService(strategy);
-			return new WebSocketHandlerAdapter(service);
+			return new WebSocketHandlerAdapter(webSocketService());
 		}
 
-		protected abstract RequestUpgradeStrategy createUpgradeStrategy();
+		@Bean
+		public WebSocketService webSocketService() {
+			return new HandshakeWebSocketService(getUpgradeStrategy());
+		}
+
+		protected abstract RequestUpgradeStrategy getUpgradeStrategy();
 
 	}
 
@@ -117,7 +120,7 @@ public abstract class AbstractWebSocketHandlerIntegrationTests {
 	static class ReactorNettyConfig extends AbstractHandlerAdapterConfig {
 
 		@Override
-		protected RequestUpgradeStrategy createUpgradeStrategy() {
+		protected RequestUpgradeStrategy getUpgradeStrategy() {
 			return new ReactorNettyRequestUpgradeStrategy();
 		}
 	}
@@ -126,7 +129,7 @@ public abstract class AbstractWebSocketHandlerIntegrationTests {
 	static class RxNettyConfig extends AbstractHandlerAdapterConfig {
 
 		@Override
-		protected RequestUpgradeStrategy createUpgradeStrategy() {
+		protected RequestUpgradeStrategy getUpgradeStrategy() {
 			return new RxNettyRequestUpgradeStrategy();
 		}
 	}