diff --git a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/MultiServerUserRegistry.java b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/MultiServerUserRegistry.java index b8d7007f13..144d7f6ca6 100644 --- a/spring-messaging/src/main/java/org/springframework/messaging/simp/user/MultiServerUserRegistry.java +++ b/spring-messaging/src/main/java/org/springframework/messaging/simp/user/MultiServerUserRegistry.java @@ -55,6 +55,9 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati private final boolean delegateApplicationEvents; + /* Cross-server session lookup (e.g. same user connected to multiple servers) */ + private final SessionLookup sessionLookup = new SessionLookup(); + /** * Create an instance wrapping the local user registry. @@ -109,36 +112,34 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati @Override public SimpUser getUser(String userName) { - SimpUser user = this.localRegistry.getUser(userName); - if (user != null) { - return user; - } + // Prefer remote registries due to cross-server SessionLookup for (UserRegistrySnapshot registry : this.remoteRegistries.values()) { - user = registry.getUserMap().get(userName); + SimpUser user = registry.getUserMap().get(userName); if (user != null) { return user; } } - return null; + return this.localRegistry.getUser(userName); } @Override public Set getUsers() { + // Prefer remote registries due to cross-server SessionLookup Set result = new HashSet(); - result.addAll(this.localRegistry.getUsers()); for (UserRegistrySnapshot registry : this.remoteRegistries.values()) { result.addAll(registry.getUserMap().values()); } + result.addAll(this.localRegistry.getUsers()); return result; } @Override public Set findSubscriptions(SimpSubscriptionMatcher matcher) { Set result = new HashSet(); - result.addAll(this.localRegistry.findSubscriptions(matcher)); for (UserRegistrySnapshot registry : this.remoteRegistries.values()) { result.addAll(registry.findSubscriptions(matcher)); } + result.addAll(this.localRegistry.findSubscriptions(matcher)); return result; } @@ -151,7 +152,7 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati void addRemoteRegistryDto(Message message, MessageConverter converter, long expirationPeriod) { UserRegistrySnapshot registry = (UserRegistrySnapshot) converter.fromMessage(message, UserRegistrySnapshot.class); if (registry != null && !registry.getId().equals(this.id)) { - registry.init(expirationPeriod); + registry.init(expirationPeriod, this.sessionLookup); this.remoteRegistries.put(registry.getId(), registry); } } @@ -228,10 +229,10 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati } - public void init(long expirationPeriod) { + public void init(long expirationPeriod, SessionLookup sessionLookup) { this.expirationTime = System.currentTimeMillis() + expirationPeriod; for (TransferSimpUser user : this.users.values()) { - user.afterDeserialization(); + user.afterDeserialization(sessionLookup); } } @@ -267,8 +268,12 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati private String name; + /* User sessions from "this" registry only (i.e. one server) */ private Set sessions; + /* Cross-server session lookup (e.g. user connected to multiple servers) */ + private SessionLookup sessionLookup; + /** * Default constructor for JSON deserialization. @@ -301,11 +306,17 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati @Override public boolean hasSessions() { + if (this.sessionLookup != null) { + return !this.sessionLookup.findSessions(getName()).isEmpty(); + } return !this.sessions.isEmpty(); } @Override public SimpSession getSession(String sessionId) { + if (this.sessionLookup != null) { + return this.sessionLookup.findSessions(getName()).get(sessionId); + } for (TransferSimpSession session : this.sessions) { if (session.getId().equals(sessionId)) { return session; @@ -320,16 +331,28 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati @Override public Set getSessions() { + if (this.sessionLookup != null) { + Map sessions = this.sessionLookup.findSessions(getName()); + return new HashSet(sessions.values()); + } return new HashSet(this.sessions); } - private void afterDeserialization() { + private void afterDeserialization(SessionLookup sessionLookup) { + this.sessionLookup = sessionLookup; for (TransferSimpSession session : this.sessions) { session.setUser(this); session.afterDeserialization(); } } + private void addSessions(Map map) { + for (SimpSession session : this.sessions) { + map.put(session.getId(), session); + } + } + + @Override public boolean equals(Object other) { return (this == other || (other instanceof SimpUser && this.name.equals(((SimpUser) other).getName()))); @@ -506,4 +529,29 @@ public class MultiServerUserRegistry implements SimpUserRegistry, SmartApplicati } } + + /** + * Helper class to find user sessions across all servers. + */ + private class SessionLookup { + + public Map findSessions(String userName) { + Map map = new HashMap(1); + SimpUser user = localRegistry.getUser(userName); + if (user != null) { + for (SimpSession session : user.getSessions()) { + map.put(session.getId(), session); + } + } + for (UserRegistrySnapshot registry : remoteRegistries.values()) { + TransferSimpUser transferUser = registry.getUserMap().get(userName); + if (transferUser != null) { + transferUser.addSessions(map); + } + } + return map; + } + + } + } diff --git a/spring-messaging/src/test/java/org/springframework/messaging/simp/user/MultiServerUserRegistryTests.java b/spring-messaging/src/test/java/org/springframework/messaging/simp/user/MultiServerUserRegistryTests.java index 44e74c6480..5ce5dd708a 100644 --- a/spring-messaging/src/test/java/org/springframework/messaging/simp/user/MultiServerUserRegistryTests.java +++ b/spring-messaging/src/test/java/org/springframework/messaging/simp/user/MultiServerUserRegistryTests.java @@ -16,6 +16,7 @@ package org.springframework.messaging.simp.user; +import static org.hamcrest.Matchers.containsInAnyOrder; import static org.junit.Assert.*; import static org.mockito.Mockito.*; @@ -54,6 +55,7 @@ public class MultiServerUserRegistryTests { this.converter = new MappingJackson2MessageConverter(); } + @Test public void getUserFromLocalRegistry() throws Exception { SimpUser user = Mockito.mock(SimpUser.class); @@ -85,6 +87,7 @@ public class MultiServerUserRegistryTests { assertEquals(1, this.registry.getUsers().size()); SimpUser user = this.registry.getUser("joe"); assertNotNull(user); + assertTrue(user.hasSessions()); assertEquals(1, user.getSessions().size()); SimpSession session = user.getSession("remote-sess"); assertNotNull(session); @@ -98,7 +101,7 @@ public class MultiServerUserRegistryTests { } @Test - public void findUserFromRemoteRegistry() throws Exception { + public void findSubscriptionsFromRemoteRegistry() throws Exception { // Prepare broadcast message from remote server TestSimpUser user1 = new TestSimpUser("joe"); @@ -132,6 +135,43 @@ public class MultiServerUserRegistryTests { assertEquals(new HashSet<>(Arrays.asList("sess1", "sess2")), sessionIds); } + @Test // SPR-13800 + public void getSessionsWhenUserIsConnectedToMultipleServers() throws Exception { + + // Add user to local registry + TestSimpUser localUser = new TestSimpUser("joe"); + TestSimpSession localSession = new TestSimpSession("sess123"); + localUser.addSessions(localSession); + when(this.localRegistry.getUser("joe")).thenReturn(localUser); + + // Prepare broadcast message from remote server + TestSimpUser remoteUser = new TestSimpUser("joe"); + TestSimpSession remoteSession = new TestSimpSession("sess456"); + remoteUser.addSessions(remoteSession); + SimpUserRegistry remoteRegistry = mock(SimpUserRegistry.class); + when(remoteRegistry.getUsers()).thenReturn(Collections.singleton(remoteUser)); + Object remoteRegistryDto = new MultiServerUserRegistry(remoteRegistry).getLocalRegistryDto(); + Message message = this.converter.toMessage(remoteRegistryDto, null); + + // Add remote registry + this.registry.addRemoteRegistryDto(message, this.converter, 20000); + + + assertEquals(1, this.registry.getUsers().size()); + SimpUser user = this.registry.getUsers().iterator().next(); + assertTrue(user.hasSessions()); + assertEquals(2, user.getSessions().size()); + assertThat(user.getSessions(), containsInAnyOrder(localSession, remoteSession)); + assertSame(localSession, user.getSession("sess123")); + assertEquals(remoteSession, user.getSession("sess456")); + + user = this.registry.getUser("joe"); + assertEquals(2, user.getSessions().size()); + assertThat(user.getSessions(), containsInAnyOrder(localSession, remoteSession)); + assertSame(localSession, user.getSession("sess123")); + assertEquals(remoteSession, user.getSession("sess456")); + } + @Test public void purgeExpiredRegistries() throws Exception {