diff --git a/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java b/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java index a4d7fe81cf..27a8ca275c 100644 --- a/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java +++ b/application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java @@ -315,7 +315,7 @@ public class DefaultWebSocketService implements WebSocketService { } } - private void processSessionClose(WebSocketSessionRef sessionRef) { + void processSessionClose(WebSocketSessionRef sessionRef) { var tenantProfileConfiguration = getTenantProfileConfiguration(sessionRef); if (tenantProfileConfiguration != null) { String sessionId = "[" + sessionRef.getSessionId() + "]"; @@ -403,7 +403,9 @@ public class DefaultWebSocketService implements WebSocketService { if (tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser() > 0 && UserPrincipal.Type.PUBLIC_ID.equals(sessionRef.getSecurityCtx().getUserPrincipal().getType())) { Set publicUserSessions = publicUserSubscriptionsMap.computeIfAbsent(sessionRef.getSecurityCtx().getTenantId(), id -> ConcurrentHashMap.newKeySet()); synchronized (publicUserSessions) { - if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) { + if (cmd.isUnsubscribe()) { + publicUserSessions.remove(subId); + } else if (publicUserSessions.size() < tenantProfileConfiguration.getMaxWsSubscriptionsPerPublicUser()) { publicUserSessions.add(subId); } else { log.info("[{}][{}][{}] Failed to start subscription. Max public user subscriptions limit reached" diff --git a/application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java b/application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java index 69f918ece7..a533e0369f 100644 --- a/application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java +++ b/application/src/test/java/org/thingsboard/server/service/ws/DefaultWebSocketServiceTest.java @@ -146,6 +146,113 @@ class DefaultWebSocketServiceTest { assertThat(subs.iterator().next()).isEqualTo("[" + sessionId + "]:[" + cmdId + "]"); } + @Test + void processSubscription_unsubscribe_removesEntryFromPublicUserSubscriptionsMap() { + int maxPublicSubscriptions = 5; + TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile = new TenantProfile(); + profile.createDefaultTenantProfileData(); + profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); + willReturn(profile).given(tenantProfileCache).get(tenantId); + + String sessionId = "session-1"; + int cmdId = 1; + WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId); + + service.processSubscription(sessionRef, subscriptionCmd(cmdId)); + + @SuppressWarnings("unchecked") + ConcurrentMap> publicUserSubscriptionsMap = + (ConcurrentMap>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); + assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(1); + + AttributesSubscriptionCmd unsubCmd = subscriptionCmd(cmdId); + unsubCmd.setUnsubscribe(true); + service.processSubscription(sessionRef, unsubCmd); + + assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty(); + } + + @Test + void processSubscription_unsubscribe_freesSlotForNewSubscription() { + int maxPublicSubscriptions = 1; + TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile = new TenantProfile(); + profile.createDefaultTenantProfileData(); + profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); + willReturn(profile).given(tenantProfileCache).get(tenantId); + + WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, "session-1"); + service.processSubscription(sessionRef, subscriptionCmd(1)); + + // slot is full — second subscription on same session should be rejected + assertThat(service.processSubscription(sessionRef, subscriptionCmd(2))).isFalse(); + + // unsubscribe cmd 1 to free the slot + AttributesSubscriptionCmd unsubCmd = subscriptionCmd(1); + unsubCmd.setUnsubscribe(true); + service.processSubscription(sessionRef, unsubCmd); + + // now a new subscription should succeed + assertThat(service.processSubscription(sessionRef, subscriptionCmd(3))) + .as("new subscription should succeed after unsubscribe freed the slot") + .isTrue(); + } + + @Test + void processSessionClose_removesAllSessionSubscriptionsFromPublicUserSubscriptionsMap() { + int maxPublicSubscriptions = 10; + TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile = new TenantProfile(); + profile.createDefaultTenantProfileData(); + profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); + willReturn(profile).given(tenantProfileCache).get(tenantId); + + String sessionId = "closing-session"; + WebSocketSessionRef sessionRef = mockPublicSessionRef(tenantId, sessionId); + + service.processSubscription(sessionRef, subscriptionCmd(1)); + service.processSubscription(sessionRef, subscriptionCmd(2)); + service.processSubscription(sessionRef, subscriptionCmd(3)); + + @SuppressWarnings("unchecked") + ConcurrentMap> publicUserSubscriptionsMap = + (ConcurrentMap>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); + assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3); + + service.processSessionClose(sessionRef); + + assertThat(publicUserSubscriptionsMap.get(tenantId)).isEmpty(); + } + + @Test + void processSessionClose_onlyRemovesClosedSessionSubscriptions() { + int maxPublicSubscriptions = 10; + TenantId tenantId = TenantId.fromUUID(UUID.randomUUID()); + TenantProfile profile = new TenantProfile(); + profile.createDefaultTenantProfileData(); + profile.getDefaultProfileConfiguration().setMaxWsSubscriptionsPerPublicUser(maxPublicSubscriptions); + willReturn(profile).given(tenantProfileCache).get(tenantId); + + WebSocketSessionRef session1 = mockPublicSessionRef(tenantId, "session-1"); + WebSocketSessionRef session2 = mockPublicSessionRef(tenantId, "session-2"); + + service.processSubscription(session1, subscriptionCmd(1)); + service.processSubscription(session1, subscriptionCmd(2)); + service.processSubscription(session2, subscriptionCmd(1)); + + @SuppressWarnings("unchecked") + ConcurrentMap> publicUserSubscriptionsMap = + (ConcurrentMap>) ReflectionTestUtils.getField(service, "publicUserSubscriptionsMap"); + assertThat(publicUserSubscriptionsMap.get(tenantId)).hasSize(3); + + service.processSessionClose(session1); + + Set remaining = publicUserSubscriptionsMap.get(tenantId); + assertThat(remaining).hasSize(1); + assertThat(remaining).allMatch(subId -> subId.startsWith("[session-2]")); + } + private WebSocketSessionRef mockPublicSessionRef(TenantId tenantId, String sessionId) { CustomerId customerId = new CustomerId(UUID.randomUUID()); SecurityUser securityUser = mock(SecurityUser.class);