Polish SessionSubscriptionRegistry types

spring-websocket/src/main/java/org/springframework/web/messaging/SessionSubscriptionRegistry.java

@@ -33,4 +33,6 @@ public interface SessionSubscriptionRegistry {
 
 	Set<String> getSessionSubscriptions(String sessionId, String destination);
 
+	Set<SessionSubscriptionRegistration> getRegistrationsByDestination(String destination);
+
 }

spring-websocket/src/main/java/org/springframework/web/messaging/service/SimpleBrokerWebMessageHandler.java

@@ -40,7 +40,7 @@ public class SimpleBrokerWebMessageHandler extends AbstractWebMessageHandler {
 
 	private final MessageChannel clientChannel;
 
-	private CachingSessionSubscriptionRegistry subscriptionRegistry=
+	private SessionSubscriptionRegistry subscriptionRegistry=
 			new CachingSessionSubscriptionRegistry(new DefaultSessionSubscriptionRegistry());
 
 
@@ -56,7 +56,7 @@ public class SimpleBrokerWebMessageHandler extends AbstractWebMessageHandler {
 
 	public void setSubscriptionRegistry(SessionSubscriptionRegistry subscriptionRegistry) {
 		Assert.notNull(subscriptionRegistry, "subscriptionRegistry is required");
-		this.subscriptionRegistry = new CachingSessionSubscriptionRegistry(subscriptionRegistry);
+		this.subscriptionRegistry = subscriptionRegistry;
 	}
 
 	@Override

spring-websocket/src/main/java/org/springframework/web/messaging/support/CachingSessionSubscriptionRegistry.java

@@ -28,8 +28,7 @@ import org.springframework.web.messaging.SessionSubscriptionRegistry;
 
 /**
  * A decorator for a {@link SessionSubscriptionRegistry} that intercepts subscriptions
- * being added and removed, and maintains a cache that tracks registrations for a
- * given destination.
+ * being added and removed and maintains a lookup cache of registrations by destination.
  *
  * @author Rossen Stoyanchev
  * @since 4.0
@@ -49,7 +48,8 @@ public class CachingSessionSubscriptionRegistry implements SessionSubscriptionRe
 
 	@Override
 	public SessionSubscriptionRegistration getRegistration(String sessionId) {
-		return new CachingSessionSubscriptionRegistration(this.delegate.getRegistration(sessionId));
+		SessionSubscriptionRegistration reg = this.delegate.getRegistration(sessionId);
+		return (reg != null) ? new CachingSessionSubscriptionRegistration(reg) : null;
 	}
 
 	@Override
@@ -71,6 +71,7 @@ public class CachingSessionSubscriptionRegistry implements SessionSubscriptionRe
 		return this.delegate.getSessionSubscriptions(sessionId, destination);
 	}
 
+	@Override
 	public Set<SessionSubscriptionRegistration> getRegistrationsByDestination(String destination) {
 		return this.destinationCache.getRegistrations(destination);
 	}

spring-websocket/src/main/java/org/springframework/web/messaging/support/DefaultSessionSubscriptionRegistry.java

@@ -16,6 +16,7 @@
 
 package org.springframework.web.messaging.support;
 
+import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -63,4 +64,20 @@ public class DefaultSessionSubscriptionRegistry implements SessionSubscriptionRe
 		return (registration != null) ? registration.getSubscriptionsByDestination(destination) : null;
 	}
 
+	/**
+	 * The default implementation performs a lookup by destination on each registration.
+	 * For a more efficient algorithm consider decorating an instance of this class with
+	 * {@link CachingSessionSubscriptionRegistry}.
+	 */
+	@Override
+	public Set<SessionSubscriptionRegistration> getRegistrationsByDestination(String destination) {
+		Set<SessionSubscriptionRegistration> result = new HashSet<SessionSubscriptionRegistration>();
+		for (SessionSubscriptionRegistration r : this.registrations.values()) {
+			if (r.getSubscriptionsByDestination(destination) != null) {
+				result.add(r);
+			}
+		}
+		return result.isEmpty() ? null : result;
+	}
+
 }