How to use

SimpUserRegistry

in

org.springframework.messaging.simp.user

@Override
@Nullable
public SimpUser getUser(String userName) {
  // Prefer remote registries due to cross-server SessionLookup
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    SimpUser user = registry.getUserMap().get(userName);
    if (user != null) {
      return user;
    }
  }
  return this.localRegistry.getUser(userName);
}

@Test
public void getUserFromLocalRegistry() throws Exception {
  SimpUser user = Mockito.mock(SimpUser.class);
  Set<SimpUser> users = Collections.singleton(user);
  when(this.localRegistry.getUsers()).thenReturn(users);
  when(this.localRegistry.getUserCount()).thenReturn(1);
  when(this.localRegistry.getUser("joe")).thenReturn(user);
  assertEquals(1, this.registry.getUserCount());
  assertSame(user, this.registry.getUser("joe"));
}

@Override
public Set<SimpSubscription> findSubscriptions(SimpSubscriptionMatcher matcher) {
  Set<SimpSubscription> result = new HashSet<>();
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    result.addAll(registry.findSubscriptions(matcher));
  }
  result.addAll(this.localRegistry.findSubscriptions(matcher));
  return result;
}

@Test
public void handleMessageFromOwnBroadcast() throws Exception {
  TestSimpUser simpUser = new TestSimpUser("joe");
  simpUser.addSessions(new TestSimpSession("123"));
  when(this.localRegistry.getUserCount()).thenReturn(1);
  when(this.localRegistry.getUsers()).thenReturn(Collections.singleton(simpUser));
  assertEquals(1, this.multiServerRegistry.getUserCount());
  Message<?> message = this.converter.toMessage(this.multiServerRegistry.getLocalRegistryDto(), null);
  this.multiServerRegistry.addRemoteRegistryDto(message, this.converter, 20000);
  assertEquals(1, this.multiServerRegistry.getUserCount());
}

@Override
public Set<SimpUser> getUsers() {
  // Prefer remote registries due to cross-server SessionLookup
  Set<SimpUser> result = new HashSet<>();
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    result.addAll(registry.getUserMap().values());
  }
  result.addAll(this.localRegistry.getUsers());
  return result;
}

@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.getUserCount());
  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"));
}

@Override
public int getUserCount() {
  int userCount = 0;
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    userCount += registry.getUserMap().size();
  }
  userCount += this.localRegistry.getUserCount();
  return userCount;
}

/**
 * Constructor to create DTO from a local user registry.
 */
public UserRegistrySnapshot(String id, SimpUserRegistry registry) {
  this.id = id;
  Set<SimpUser> users = registry.getUsers();
  this.users = new HashMap<>(users.size());
  for (SimpUser user : users) {
    this.users.put(user.getName(), new TransferSimpUser(user));
  }
}

@Test
public void handleMessage() throws Exception {
  TestSimpUser simpUser1 = new TestSimpUser("joe");
  TestSimpUser simpUser2 = new TestSimpUser("jane");
  simpUser1.addSessions(new TestSimpSession("123"));
  simpUser2.addSessions(new TestSimpSession("456"));
  HashSet<SimpUser> simpUsers = new HashSet<>(Arrays.asList(simpUser1, simpUser2));
  SimpUserRegistry remoteUserRegistry = mock(SimpUserRegistry.class);
  when(remoteUserRegistry.getUserCount()).thenReturn(2);
  when(remoteUserRegistry.getUsers()).thenReturn(simpUsers);
  MultiServerUserRegistry remoteRegistry = new MultiServerUserRegistry(remoteUserRegistry);
  Message<?> message = this.converter.toMessage(remoteRegistry.getLocalRegistryDto(), null);
  this.handler.handleMessage(message);
  assertEquals(2, remoteRegistry.getUserCount());
  assertNotNull(this.multiServerRegistry.getUser("joe"));
  assertNotNull(this.multiServerRegistry.getUser("jane"));
}

@Override
public int getUserCount() {
  int userCount = 0;
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    userCount += registry.getUserMap().size();
  }
  userCount += this.localRegistry.getUserCount();
  return userCount;
}

private Set<String> getSessionIdsByUser(String userName, @Nullable String sessionId) {
  Set<String> sessionIds;
  SimpUser user = this.userRegistry.getUser(userName);
  if (user != null) {
    if (sessionId != null && user.getSession(sessionId) != null) {
      sessionIds = Collections.singleton(sessionId);
    }
    else {
      Set<SimpSession> sessions = user.getSessions();
      sessionIds = new HashSet<>(sessions.size());
      for (SimpSession session : sessions) {
        sessionIds.add(session.getId());
      }
    }
  }
  else {
    sessionIds = Collections.emptySet();
  }
  return sessionIds;
}

@Test
public void broadcastRegistry() throws Exception {
  TestSimpUser simpUser1 = new TestSimpUser("joe");
  TestSimpUser simpUser2 = new TestSimpUser("jane");
  simpUser1.addSessions(new TestSimpSession("123"));
  simpUser1.addSessions(new TestSimpSession("456"));
  HashSet<SimpUser> simpUsers = new HashSet<>(Arrays.asList(simpUser1, simpUser2));
  when(this.localRegistry.getUsers()).thenReturn(simpUsers);
  getUserRegistryTask().run();
  ArgumentCaptor<Message> captor = ArgumentCaptor.forClass(Message.class);
  verify(this.brokerChannel).send(captor.capture());
  Message<?> message = captor.getValue();
  assertNotNull(message);
  MessageHeaders headers = message.getHeaders();
  assertEquals("/topic/simp-user-registry", SimpMessageHeaderAccessor.getDestination(headers));
  MultiServerUserRegistry remoteRegistry = new MultiServerUserRegistry(mock(SimpUserRegistry.class));
  remoteRegistry.addRemoteRegistryDto(message, this.converter, 20000);
  assertEquals(2, remoteRegistry.getUserCount());
  assertNotNull(remoteRegistry.getUser("joe"));
  assertNotNull(remoteRegistry.getUser("jane"));
}

@Override
public int getUserCount() {
  int userCount = 0;
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    userCount += registry.getUserMap().size();
  }
  userCount += this.localRegistry.getUserCount();
  return userCount;
}

@Override
public Set<SimpSubscription> findSubscriptions(SimpSubscriptionMatcher matcher) {
  Set<SimpSubscription> result = new HashSet<>();
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    result.addAll(registry.findSubscriptions(matcher));
  }
  result.addAll(this.localRegistry.findSubscriptions(matcher));
  return result;
}

public Map<String, SimpSession> findSessions(String userName) {
  Map<String, SimpSession> map = new HashMap<>(4);
  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;
}

@Test
public void purgeExpiredRegistries() throws Exception {
  // Prepare broadcast message from remote server
  TestSimpUser testUser = new TestSimpUser("joe");
  testUser.addSessions(new TestSimpSession("remote-sub"));
  SimpUserRegistry testRegistry = mock(SimpUserRegistry.class);
  when(testRegistry.getUsers()).thenReturn(Collections.singleton(testUser));
  Object registryDto = new MultiServerUserRegistry(testRegistry).getLocalRegistryDto();
  Message<?> message = this.converter.toMessage(registryDto, null);
  // Add remote registry
  this.registry.addRemoteRegistryDto(message, this.converter, -1);
  assertEquals(1, this.registry.getUserCount());
  this.registry.purgeExpiredRegistries();
  assertEquals(0, this.registry.getUserCount());
}

@Override
public int getUserCount() {
  int userCount = 0;
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    userCount += registry.getUserMap().size();
  }
  userCount += this.localRegistry.getUserCount();
  return userCount;
}

@Override
public Set<SimpSubscription> findSubscriptions(SimpSubscriptionMatcher matcher) {
  Set<SimpSubscription> result = new HashSet<>();
  for (UserRegistrySnapshot registry : this.remoteRegistries.values()) {
    result.addAll(registry.findSubscriptions(matcher));
  }
  result.addAll(this.localRegistry.findSubscriptions(matcher));
  return result;
}

@Before
public void setup() {
  TestSimpUser simpUser = new TestSimpUser("joe");
  simpUser.addSessions(new TestSimpSession("123"));
  this.registry = mock(SimpUserRegistry.class);
  when(this.registry.getUser("joe")).thenReturn(simpUser);
  this.resolver = new DefaultUserDestinationResolver(this.registry);
}

@Test
public void findSubscriptionsFromRemoteRegistry() throws Exception {
  // Prepare broadcast message from remote server
  TestSimpUser user1 = new TestSimpUser("joe");
  TestSimpUser user2 = new TestSimpUser("jane");
  TestSimpUser user3 = new TestSimpUser("jack");
  TestSimpSession session1 = new TestSimpSession("sess1");
  TestSimpSession session2 = new TestSimpSession("sess2");
  TestSimpSession session3 = new TestSimpSession("sess3");
  session1.addSubscriptions(new TestSimpSubscription("sub1", "/match"));
  session2.addSubscriptions(new TestSimpSubscription("sub1", "/match"));
  session3.addSubscriptions(new TestSimpSubscription("sub1", "/not-a-match"));
  user1.addSessions(session1);
  user2.addSessions(session2);
  user3.addSessions(session3);
  SimpUserRegistry userRegistry = mock(SimpUserRegistry.class);
  when(userRegistry.getUsers()).thenReturn(new HashSet<>(Arrays.asList(user1, user2, user3)));
  Object registryDto = new MultiServerUserRegistry(userRegistry).getLocalRegistryDto();
  Message<?> message = this.converter.toMessage(registryDto, null);
  // Add remote registry
  this.registry.addRemoteRegistryDto(message, this.converter, 20000);
  assertEquals(3, this.registry.getUserCount());
  Set<SimpSubscription> matches = this.registry.findSubscriptions(s -> s.getDestination().equals("/match"));
  assertEquals(2, matches.size());
  Iterator<SimpSubscription> iterator = matches.iterator();
  Set<String> sessionIds = new HashSet<>(2);
  sessionIds.add(iterator.next().getSession().getId());
  sessionIds.add(iterator.next().getSession().getId());
  assertEquals(new HashSet<>(Arrays.asList("sess1", "sess2")), sessionIds);
}