Update SpringConfiguration to support beans by type

Issue: SPR-10605

Author: rstoyanchev

spring-messaging/src/test/java/org/springframework/messaging/simp/JettyTestServer.java

@@ -22,7 +22,6 @@
 import org.springframework.util.SocketUtils;
 import org.springframework.web.context.WebApplicationContext;
 import org.springframework.web.servlet.DispatcherServlet;
-import org.springframework.web.socket.TestServer;
 
 
 /**

spring-websocket/src/main/java/org/springframework/web/socket/server/endpoint/SpringConfigurator.java

@@ -16,6 +16,9 @@
 
 package org.springframework.web.socket.server.endpoint;
 
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
 import javax.websocket.server.ServerEndpoint;
 import javax.websocket.server.ServerEndpointConfig.Configurator;
 
@@ -24,6 +27,7 @@
 import org.springframework.core.annotation.AnnotationUtils;
 import org.springframework.stereotype.Component;
 import org.springframework.util.ClassUtils;
+import org.springframework.util.ObjectUtils;
 import org.springframework.web.context.ContextLoader;
 import org.springframework.web.context.WebApplicationContext;
 
@@ -48,7 +52,13 @@ public class SpringConfigurator extends Configurator {
 
 	private static Log logger = LogFactory.getLog(SpringConfigurator.class);
 
+	private static final Map<String, Map<Class<?>, String>> cache =
+			new ConcurrentHashMap<String, Map<Class<?>, String>>();
+
+	private static final String NO_VALUE = ObjectUtils.identityToString(new Object());
 
+
+	@SuppressWarnings("unchecked")
 	@Override
 	public <T> T getEndpointInstance(Class<T> endpointClass) throws InstantiationException {
 
@@ -77,10 +87,44 @@ public <T> T getEndpointInstance(Class<T> endpointClass) throws InstantiationExc
 			return endpoint;
 		}
 
+		beanName = getBeanNameByType(wac, endpointClass);
+		if (beanName != null) {
+			return (T) wac.getBean(beanName);
+		}
+
 		if (logger.isTraceEnabled()) {
 			logger.trace("Creating new @ServerEndpoint instance of type " + endpointClass);
 		}
 		return wac.getAutowireCapableBeanFactory().createBean(endpointClass);
 	}
 
+	private String getBeanNameByType(WebApplicationContext wac, Class<?> endpointClass) {
+
+		String wacId = wac.getId();
+
+		Map<Class<?>, String> beanNamesByType = cache.get(wacId);
+		if (beanNamesByType == null) {
+			beanNamesByType = new ConcurrentHashMap<Class<?>, String>();
+			cache.put(wacId, beanNamesByType);
+		}
+
+		if (!beanNamesByType.containsKey(endpointClass)) {
+			String[] names = wac.getBeanNamesForType(endpointClass);
+			if (names.length == 1) {
+				beanNamesByType.put(endpointClass, names[0]);
+			}
+			else {
+				beanNamesByType.put(endpointClass, NO_VALUE);
+				if (names.length > 1) {
+					String message = "Found multiple @ServerEndpoint's of type " + endpointClass + ", names=" + names;
+					logger.error(message);
+					throw new IllegalStateException(message);
+				}
+			}
+		}
+
+		String beanName = beanNamesByType.get(endpointClass);
+		return NO_VALUE.equals(beanName) ? null : beanName;
+	}
+
 }

spring-websocket/src/test/java/org/springframework/web/socket/server/endpoint/SpringConfiguratorTests.java

@@ -16,9 +16,7 @@
 
 package org.springframework.web.socket.server.endpoint;
 
-import javax.websocket.Endpoint;
-import javax.websocket.EndpointConfig;
-import javax.websocket.Session;
+import javax.websocket.server.ServerEndpoint;
 
 import org.junit.After;
 import org.junit.Before;
@@ -65,22 +63,22 @@ public void destroy() {
 
 
 	@Test
-	public void getEndpointInstancePerConnection() throws Exception {
+	public void getEndpointPerConnection() throws Exception {
 		PerConnectionEchoEndpoint endpoint = this.configurator.getEndpointInstance(PerConnectionEchoEndpoint.class);
 		assertNotNull(endpoint);
 	}
 
 	@Test
-	public void getEndpointInstanceSingletonByType() throws Exception {
+	public void getEndpointSingletonByType() throws Exception {
 		EchoEndpoint expected = this.webAppContext.getBean(EchoEndpoint.class);
 		EchoEndpoint actual = this.configurator.getEndpointInstance(EchoEndpoint.class);
 		assertSame(expected, actual);
 	}
 
 	@Test
-	public void getEndpointInstanceSingletonByComponentName() throws Exception {
-		AnotherEchoEndpoint expected = this.webAppContext.getBean(AnotherEchoEndpoint.class);
-		AnotherEchoEndpoint actual = this.configurator.getEndpointInstance(AnotherEchoEndpoint.class);
+	public void getEndpointSingletonByComponentName() throws Exception {
+		ComponentEchoEndpoint expected = this.webAppContext.getBean(ComponentEchoEndpoint.class);
+		ComponentEchoEndpoint actual = this.configurator.getEndpointInstance(ComponentEchoEndpoint.class);
 		assertSame(expected, actual);
 	}
 
@@ -90,7 +88,7 @@ public void getEndpointInstanceSingletonByComponentName() throws Exception {
 	static class Config {
 
 		@Bean
-		public EchoEndpoint echoEndpoint() {
+		public EchoEndpoint javaConfigEndpoint() {
 			return new EchoEndpoint(echoService());
 		}
 
@@ -100,7 +98,8 @@ public EchoService echoService() {
 		}
 	}
 
-	private static class EchoEndpoint extends Endpoint {
+	@ServerEndpoint("/echo")
+	private static class EchoEndpoint {
 
 		@SuppressWarnings("unused")
 		private final EchoService service;
@@ -109,29 +108,23 @@ private static class EchoEndpoint extends Endpoint {
 		public EchoEndpoint(EchoService service) {
 			this.service = service;
 		}
-
-		@Override
-		public void onOpen(Session session, EndpointConfig config) {
-		}
 	}
 
-	@Component("myEchoEndpoint")
-	private static class AnotherEchoEndpoint extends Endpoint {
+	@Component("myComponentEchoEndpoint")
+	@ServerEndpoint("/echo")
+	private static class ComponentEchoEndpoint {
 
 		@SuppressWarnings("unused")
 		private final EchoService service;
 
 		@Autowired
-		public AnotherEchoEndpoint(EchoService service) {
+		public ComponentEchoEndpoint(EchoService service) {
 			this.service = service;
 		}
-
-		@Override
-		public void onOpen(Session session, EndpointConfig config) {
-		}
 	}
 
-	private static class PerConnectionEchoEndpoint extends Endpoint {
+	@ServerEndpoint("/echo")
+	private static class PerConnectionEchoEndpoint {
 
 		@SuppressWarnings("unused")
 		private final EchoService service;
@@ -140,10 +133,6 @@ private static class PerConnectionEchoEndpoint extends Endpoint {
 		public PerConnectionEchoEndpoint(EchoService service) {
 			this.service = service;
 		}
-
-		@Override
-		public void onOpen(Session session, EndpointConfig config) {
-		}
 	}
 
 	private static class EchoService {	}