Mapping of cmdId to unique sequence per subscription to fix race condition with entity data query commands

application/src/main/java/org/thingsboard/server/service/subscription/TbSubscription.java

@@ -36,8 +36,6 @@ public abstract class TbSubscription<T> {
     private final TbSubscriptionType type;
     private final BiConsumer<TbSubscription<T>, T> updateProcessor;
 
-    protected final AtomicInteger sequence = new AtomicInteger();
-
     @Override
     public boolean equals(Object o) {
         if (this == o) return true;

application/src/main/java/org/thingsboard/server/service/ws/DefaultWebSocketService.java

@@ -259,13 +259,13 @@ public class DefaultWebSocketService implements WebSocketService {
     }
 
     @Override
-    public void sendUpdate(String sessionId, TelemetrySubscriptionUpdate update) {
+    public void sendUpdate(String sessionId, int cmdId, TelemetrySubscriptionUpdate update) {
-        sendUpdate(sessionId, update.getSubscriptionId(), update);
+        doSendUpdate(sessionId, cmdId, update);
     }
 
     @Override
     public void sendUpdate(String sessionId, CmdUpdate update) {
-        sendUpdate(sessionId, update.getCmdId(), update);
+        doSendUpdate(sessionId, update.getCmdId(), update);
     }
 
     @Override
@@ -274,7 +274,7 @@ public class DefaultWebSocketService implements WebSocketService {
         sendUpdate(sessionRef, update);
     }
 
-    private <T> void sendUpdate(String sessionId, int cmdId, T update) {
+    private <T> void doSendUpdate(String sessionId, int cmdId, T update) {
         WsSessionMetaData md = wsSessionsMap.get(sessionId);
         if (md != null) {
             sendUpdate(md.getSessionRef(), cmdId, update);
@@ -288,7 +288,7 @@ public class DefaultWebSocketService implements WebSocketService {
             try {
                 msgEndpoint.close(md.getSessionRef(), status);
             } catch (IOException e) {
-                log.warn("[{}] Failed to send session close: {}", sessionId, e);
+                log.warn("[{}] Failed to send session close", sessionId, e);
             }
         }
     }
@@ -439,7 +439,7 @@ public class DefaultWebSocketService implements WebSocketService {
                 TbAttributeSubscription sub = TbAttributeSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
-                        .subscriptionId(cmd.getCmdId())
+                        .subscriptionId(sessionRef.getSessionSubIdSeq().incrementAndGet())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
                         .queryTs(queryTs)
@@ -449,7 +449,7 @@ public class DefaultWebSocketService implements WebSocketService {
                         .updateProcessor((subscription, update) -> {
                             subLock.lock();
                             try {
-                                sendUpdate(subscription.getSessionId(), update);
+                                sendUpdate(subscription.getSessionId(), cmd.getCmdId(), update);
                             } finally {
                                 subLock.unlock();
                             }
@@ -545,7 +545,7 @@ public class DefaultWebSocketService implements WebSocketService {
                 TbAttributeSubscription sub = TbAttributeSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
-                        .subscriptionId(cmd.getCmdId())
+                        .subscriptionId(sessionRef.getSessionSubIdSeq().incrementAndGet())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
                         .queryTs(queryTs)
@@ -554,7 +554,7 @@ public class DefaultWebSocketService implements WebSocketService {
                         .updateProcessor((subscription, update) -> {
                             subLock.lock();
                             try {
-                                sendUpdate(subscription.getSessionId(), update);
+                                sendUpdate(subscription.getSessionId(), cmd.getCmdId(), update);
                             } finally {
                                 subLock.unlock();
                             }
@@ -643,13 +643,13 @@ public class DefaultWebSocketService implements WebSocketService {
                 TbTimeSeriesSubscription sub = TbTimeSeriesSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
-                        .subscriptionId(cmd.getCmdId())
+                        .subscriptionId(sessionRef.getSessionSubIdSeq().incrementAndGet())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
                         .updateProcessor((subscription, update) -> {
                             subLock.lock();
                             try {
-                                sendUpdate(subscription.getSessionId(), update);
+                                sendUpdate(subscription.getSessionId(), cmd.getCmdId(), update);
                             } finally {
                                 subLock.unlock();
                             }
@@ -698,13 +698,13 @@ public class DefaultWebSocketService implements WebSocketService {
                 TbTimeSeriesSubscription sub = TbTimeSeriesSubscription.builder()
                         .serviceId(serviceId)
                         .sessionId(sessionId)
-                        .subscriptionId(cmd.getCmdId())
+                        .subscriptionId(sessionRef.getSessionSubIdSeq().incrementAndGet())
                         .tenantId(sessionRef.getSecurityCtx().getTenantId())
                         .entityId(entityId)
                         .updateProcessor((subscription, update) -> {
                             subLock.lock();
                             try {
-                                sendUpdate(subscription.getSessionId(), update);
+                                sendUpdate(subscription.getSessionId(), cmd.getCmdId(), update);
                             } finally {
                                 subLock.unlock();
                             }
@@ -836,7 +836,7 @@ public class DefaultWebSocketService implements WebSocketService {
                     try {
                         msgEndpoint.sendPing(md.getSessionRef(), currentTime);
                     } catch (IOException e) {
-                        log.warn("[{}] Failed to send ping: {}", md.getSessionRef().getSessionId(), e);
+                        log.warn("[{}] Failed to send ping:", md.getSessionRef().getSessionId(), e);
                     }
                 }));
     }

application/src/main/java/org/thingsboard/server/service/ws/WebSocketService.java

@@ -29,7 +29,7 @@ public interface WebSocketService {
 
     void handleCommands(WebSocketSessionRef sessionRef, WsCommandsWrapper commandsWrapper);
 
-    void sendUpdate(String sessionId, TelemetrySubscriptionUpdate update);
+    void sendUpdate(String sessionId, int cmdId, TelemetrySubscriptionUpdate update);
 
     void sendUpdate(String sessionId, CmdUpdate update);
 

application/src/main/java/org/thingsboard/server/service/ws/notification/DefaultNotificationCommandsHandler.java

@@ -115,7 +115,7 @@ public class DefaultNotificationCommandsHandler implements NotificationCommandsH
     private void fetchUnreadNotificationsCount(NotificationsCountSubscription subscription) {
         log.trace("[{}, subId: {}] Fetching unread notifications count from DB", subscription.getSessionId(), subscription.getSubscriptionId());
         int unreadCount = notificationService.countUnreadNotificationsByRecipientId(subscription.getTenantId(), (UserId) subscription.getEntityId());
-        subscription.getUnreadCounter().set(unreadCount);
+        subscription.getTotalUnreadCounter().set(unreadCount);
     }
 
 
@@ -196,20 +196,20 @@ public class DefaultNotificationCommandsHandler implements NotificationCommandsH
     private void handleNotificationUpdate(NotificationsCountSubscription subscription, NotificationUpdate update) {
         log.trace("[{}, subId: {}] Handling notification update for count sub: {}", subscription.getSessionId(), subscription.getSubscriptionId(), update);
         if (update.isCreated()) {
-            subscription.getUnreadCounter().incrementAndGet();
+            subscription.getTotalUnreadCounter().incrementAndGet();
             sendUpdate(subscription.getSessionId(), subscription.createUpdate());
         } else if (update.isUpdated()) {
             if (update.getNewStatus() == NotificationStatus.READ) {
                 if (update.isAllNotifications()) {
                     fetchUnreadNotificationsCount(subscription);
                 } else {
-                    subscription.getUnreadCounter().decrementAndGet();
+                    subscription.getTotalUnreadCounter().decrementAndGet();
                 }
                 sendUpdate(subscription.getSessionId(), subscription.createUpdate());
             }
         } else if (update.isDeleted()) {
             if (update.getNotification().getStatus() != NotificationStatus.READ) {
-                subscription.getUnreadCounter().decrementAndGet();
+                subscription.getTotalUnreadCounter().decrementAndGet();
                 sendUpdate(subscription.getSessionId(), subscription.createUpdate());
             }
         }

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/AbstractNotificationSubscription.java

@@ -0,0 +1,38 @@
+/**
+ * Copyright © 2016-2024 The Thingsboard Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.thingsboard.server.service.ws.notification.sub;
+
+
+import lombok.Getter;
+import org.thingsboard.server.common.data.id.EntityId;
+import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.service.subscription.TbSubscription;
+import org.thingsboard.server.service.subscription.TbSubscriptionType;
+
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.BiConsumer;
+
+@Getter
+public abstract class AbstractNotificationSubscription<T> extends TbSubscription<T> {
+
+    protected final AtomicInteger sequence = new AtomicInteger();
+    protected final AtomicInteger totalUnreadCounter = new AtomicInteger();
+
+    public AbstractNotificationSubscription(String serviceId, String sessionId, int subscriptionId, TenantId tenantId, EntityId entityId, TbSubscriptionType type, BiConsumer<TbSubscription<T>, T> updateProcessor) {
+        super(serviceId, sessionId, subscriptionId, tenantId, entityId, type, updateProcessor);
+    }
+
+}

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/NotificationsCountSubscription.java

@@ -27,9 +27,7 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BiConsumer;
 
 @Getter
-public class NotificationsCountSubscription extends TbSubscription<NotificationsSubscriptionUpdate> {
+public class NotificationsCountSubscription extends AbstractNotificationSubscription<NotificationsSubscriptionUpdate> {
-
-    private final AtomicInteger unreadCounter = new AtomicInteger();
 
     @Builder
     public NotificationsCountSubscription(String serviceId, String sessionId, int subscriptionId, TenantId tenantId, EntityId entityId,
@@ -40,7 +38,7 @@ public class NotificationsCountSubscription extends TbSubscription<Notifications
     public UnreadNotificationsCountUpdate createUpdate() {
         return UnreadNotificationsCountUpdate.builder()
                 .cmdId(getSubscriptionId())
-                .totalUnreadCount(unreadCounter.get())
+                .totalUnreadCount(totalUnreadCounter.get())
                 .sequenceNumber(sequence.incrementAndGet())
                 .build();
     }

application/src/main/java/org/thingsboard/server/service/ws/notification/sub/NotificationsSubscription.java

@@ -35,11 +35,10 @@ import java.util.function.BiConsumer;
 import java.util.stream.Collectors;
 
 @Getter
-public class NotificationsSubscription extends TbSubscription<NotificationsSubscriptionUpdate> {
+public class NotificationsSubscription extends AbstractNotificationSubscription<NotificationsSubscriptionUpdate> {
 
     private final Map<UUID, Notification> latestUnreadNotifications = new HashMap<>();
     private final int limit;
-    private final AtomicInteger totalUnreadCounter = new AtomicInteger();
 
     @Builder
     public NotificationsSubscription(String serviceId, String sessionId, int subscriptionId, TenantId tenantId, EntityId entityId,