Improve SubscriptionRegistry subscription removal

This a backport of:
ce20abde51
083d415fc4

Issue: SPR-11931, SPR-11930

spring-messaging/src/main/java/org/springframework/messaging/simp/broker/DefaultSubscriptionRegistry.java

@@ -96,8 +96,8 @@ public class DefaultSubscriptionRegistry extends AbstractSubscriptionRegistry {
 		SessionSubscriptionInfo info = this.subscriptionRegistry.getSubscriptions(sessionId);
 		if (info != null) {
 			String destination = info.removeSubscription(subsId);
-			if (destination != null && info.getSubscriptions(destination) == null) {
+			if (destination != null) {
-				this.destinationCache.updateAfterRemovedSubscription(destination, sessionId, subsId);
+				this.destinationCache.updateAfterRemovedSubscription(sessionId, subsId);
 			}
 		}
 	}
@@ -187,54 +187,51 @@ public class DefaultSubscriptionRegistry extends AbstractSubscriptionRegistry {
 			}
 		}
 
-		public void updateAfterRemovedSubscription(String destination, String sessionId, String subsId) {
+		public void updateAfterRemovedSubscription(String sessionId, String subsId) {
 			synchronized (this.updateCache) {
 				Set<String> destinationsToRemove = new HashSet<String>();
 				for (Map.Entry<String, MultiValueMap<String, String>> entry : this.updateCache.entrySet()) {
-					String cachedDestination = entry.getKey();
+					String destination = entry.getKey();
-					if (getPathMatcher().match(destination, cachedDestination)) {
+					MultiValueMap<String, String> sessionMap = entry.getValue();
-						MultiValueMap<String, String> subs = entry.getValue();
+					List<String> subscriptions = sessionMap.get(sessionId);
-						List<String> subsIds = subs.get(sessionId);
+					if (subscriptions != null) {
-						subsIds.remove(subsId);
+						subscriptions.remove(subsId);
-						if (subsIds.isEmpty()) {
+						if (subscriptions.isEmpty()) {
-							subs.remove(sessionId);
+							sessionMap.remove(sessionId);
 						}
-						if (subs.isEmpty()) {
+						if (sessionMap.isEmpty()) {
-							destinationsToRemove.add(cachedDestination);
+							destinationsToRemove.add(destination);
 						}
 						else {
-							this.accessCache.put(cachedDestination, new LinkedMultiValueMap<String, String>(subs));
+							this.accessCache.put(destination, new LinkedMultiValueMap<String, String>(sessionMap));
 						}
 					}
 				}
-				for (String destinationToRemove : destinationsToRemove) {
+				for (String destination : destinationsToRemove) {
-					this.updateCache.remove(destinationToRemove);
+					this.updateCache.remove(destination);
-					this.accessCache.remove(destinationToRemove);
+					this.accessCache.remove(destination);
 				}
 			}
 		}
 
 		public void updateAfterRemovedSession(SessionSubscriptionInfo info) {
 			synchronized (this.updateCache) {
-				for (String destination : info.getDestinations()) {
+				Set<String> destinationsToRemove = new HashSet<String>();
-					Set<String> destinationsToRemove = new HashSet<String>();
+				for (Map.Entry<String, MultiValueMap<String, String>> entry : this.updateCache.entrySet()) {
-					for (Map.Entry<String, MultiValueMap<String, String>> entry : this.updateCache.entrySet()) {
+					String destination = entry.getKey();
-						String cachedDestination = entry.getKey();
+					MultiValueMap<String, String> sessionMap = entry.getValue();
-						if (getPathMatcher().match(destination, cachedDestination)) {
+					if (sessionMap.remove(info.getSessionId()) != null) {
-							MultiValueMap<String, String> subs = entry.getValue();
+						if (sessionMap.isEmpty()) {
-							subs.remove(info.getSessionId());
+							destinationsToRemove.add(destination);
-							if (subs.isEmpty()) {
+						}
-								destinationsToRemove.add(cachedDestination);
+						else {
-							}
+							this.accessCache.put(destination, new LinkedMultiValueMap<String, String>(sessionMap));
-							else {
-								this.accessCache.put(cachedDestination,new LinkedMultiValueMap<String, String>(subs));
-							}
 						}
 					}
-					for (String d : destinationsToRemove) {
+				}
-						this.updateCache.remove(d);
+				for (String destination : destinationsToRemove) {
-						this.accessCache.remove(d);
+					this.updateCache.remove(destination);
-					}
+					this.accessCache.remove(destination);
 				}
 			}
 		}

spring-messaging/src/test/java/org/springframework/messaging/simp/broker/DefaultSubscriptionRegistryTests.java

@@ -270,6 +270,27 @@ public class DefaultSubscriptionRegistryTests {
 		assertEquals(subscriptionIds, sort(actual.get(sessIds.get(2))));
 	}
 
+	// SPR-11931
+
+	@Test
+	public void registerTwiceAndUnregisterSubscriptions() {
+
+		this.registry.registerSubscription(subscribeMessage("sess01", "subs01", "/foo"));
+		this.registry.registerSubscription(subscribeMessage("sess01", "subs02", "/foo"));
+		MultiValueMap<String, String> actual = this.registry.findSubscriptions(message("/foo"));
+		assertEquals("Expected 1 element", 1, actual.size());
+		assertEquals(Arrays.asList("subs01", "subs02"), actual.get("sess01"));
+
+		this.registry.unregisterSubscription(unsubscribeMessage("sess01", "subs01"));
+		actual = this.registry.findSubscriptions(message("/foo"));
+		assertEquals("Expected 1 element", 1, actual.size());
+		assertEquals(Arrays.asList("subs02"), actual.get("sess01"));
+
+		this.registry.unregisterSubscription(unsubscribeMessage("sess01", "subs02"));
+		actual = this.registry.findSubscriptions(message("/foo"));
+		assertEquals("Expected no element", 0, actual.size());
+	}
+
 	@Test
 	public void unregisterAllSubscriptions() {