fetch entities only on subs creation

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

@@ -424,8 +424,12 @@ public class DefaultTbEntityDataSubscriptionService implements TbEntityDataSubsc
             long start = System.currentTimeMillis();
             ctx.fetchData();
             long end = System.currentTimeMillis();
-            stats.getAlarmQueryInvocationCnt().incrementAndGet();
+            stats.getRegularQueryInvocationCnt().incrementAndGet();
-            stats.getAlarmQueryTimeSpent().addAndGet(end - start);
+            stats.getRegularQueryTimeSpent().addAndGet(end - start);
+            ctx.cancelTasks();
+            ctx.clearAlarmSubscriptions();
+            ctx.fetchAlarmCount();
+            ctx.createAlarmSubscriptions();
             TbAlarmCountSubCtx finalCtx = ctx;
             ScheduledFuture<?> task = scheduler.scheduleWithFixedDelay(
                     () -> refreshDynamicQuery(finalCtx),
@@ -549,7 +553,7 @@ public class DefaultTbEntityDataSubscriptionService implements TbEntityDataSubsc
     private TbAlarmCountSubCtx createSubCtx(WebSocketSessionRef sessionRef, AlarmCountCmd cmd) {
         Map<Integer, TbAbstractSubCtx> sessionSubs = subscriptionsBySessionId.computeIfAbsent(sessionRef.getSessionId(), k -> new ConcurrentHashMap<>());
         TbAlarmCountSubCtx ctx = new TbAlarmCountSubCtx(serviceId, wsService, entityService, localSubscriptionService,
-                attributesService, stats, alarmService, sessionRef, cmd.getCmdId(), maxEntitiesPerAlarmSubscription);
+                attributesService, stats, alarmService, sessionRef, cmd.getCmdId(), maxEntitiesPerAlarmSubscription, maxAlarmQueriesPerRefreshInterval);
         if (cmd.getQuery() != null) {
             ctx.setAndResolveQuery(cmd.getQuery());
         }

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

@@ -36,7 +36,9 @@ import org.thingsboard.server.service.ws.WebSocketService;
 import org.thingsboard.server.service.ws.WebSocketSessionRef;
 import org.thingsboard.server.service.ws.telemetry.cmd.v2.AlarmCountUpdate;
 
-import java.util.List;
+import java.util.LinkedHashSet;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 @Slf4j
 @ToString(callSuper = true)
@@ -44,56 +46,120 @@ public class TbAlarmCountSubCtx extends TbAbstractEntityQuerySubCtx<AlarmCountQu
 
     private final AlarmService alarmService;
 
+    protected final Map<Integer, EntityId> subToEntityIdMap;
+
+    private LinkedHashSet<EntityId> entitiesIds;
+
     private final int maxEntitiesPerAlarmSubscription;
 
+    private final int maxAlarmQueriesPerRefreshInterval;
+
     @Getter
     @Setter
     private volatile int result;
 
+    @Getter
+    @Setter
+    private boolean tooManyEntities;
+
+    private int alarmCountInvocationAttempts;
+
     public TbAlarmCountSubCtx(String serviceId, WebSocketService wsService,
                               EntityService entityService, TbLocalSubscriptionService localSubscriptionService,
                               AttributesService attributesService, SubscriptionServiceStatistics stats, AlarmService alarmService,
-                              WebSocketSessionRef sessionRef, int cmdId, int maxEntitiesPerAlarmSubscription) {
+                              WebSocketSessionRef sessionRef, int cmdId, int maxEntitiesPerAlarmSubscription, int maxAlarmQueriesPerRefreshInterval) {
         super(serviceId, wsService, entityService, localSubscriptionService, attributesService, stats, sessionRef, cmdId);
         this.alarmService = alarmService;
+        this.subToEntityIdMap = new ConcurrentHashMap<>();
         this.maxEntitiesPerAlarmSubscription = maxEntitiesPerAlarmSubscription;
+        this.maxAlarmQueriesPerRefreshInterval = maxAlarmQueriesPerRefreshInterval;
+        this.entitiesIds = null;
+    }
+
+    @Override
+    public void clearSubscriptions() {
+        clearAlarmSubscriptions();
     }
 
     @Override
     public void fetchData() {
-        result = countAlarms();
+        resetInvocationCounter();
-        sendWsMsg(new AlarmCountUpdate(cmdId, result));
+        if (query.getEntityFilter() != null) {
+            entitiesIds = new LinkedHashSet<>();
+            log.trace("[{}] Fetching data: {}", cmdId, alarmCountInvocationAttempts);
+            PageData<EntityData> data = entityService.findEntityDataByQuery(getTenantId(), getCustomerId(), buildEntityDataQuery());
+            entitiesIds.clear();
+            tooManyEntities = data.hasNext();
+            for (EntityData entityData : data.getData()) {
+                entitiesIds.add(entityData.getEntityId());
+            }
+        }
     }
 
     @Override
     protected void update() {
-        int newCount = countAlarms();
+        resetInvocationCounter();
-        if (newCount != result) {
+        fetchAlarmCount();
-            result = newCount;
+    }
-            sendWsMsg(new AlarmCountUpdate(cmdId, result));
+
+    public void fetchAlarmCount() {
+        alarmCountInvocationAttempts++;
+        log.trace("[{}] Fetching alarms: {}", cmdId, alarmCountInvocationAttempts);
+        if (alarmCountInvocationAttempts <= maxAlarmQueriesPerRefreshInterval) {
+            doFetchAlarmCount();
+        } else {
+            log.trace("[{}] Ignore alarm count fetch due to rate limit: [{}] of maximum [{}]", cmdId, alarmCountInvocationAttempts, maxAlarmQueriesPerRefreshInterval);
         }
     }
 
+    private void doFetchAlarmCount() {
+        result = (int) alarmService.countAlarmsByQuery(getTenantId(), getCustomerId(), query, entitiesIds);
+        sendWsMsg(new AlarmCountUpdate(cmdId, result));
+    }
+
     @Override
     public boolean isDynamic() {
         return true;
     }
 
-    private int countAlarms() {
-        List<EntityId> entityIds = null;
-        if (query.getEntityFilter() != null) {
-            PageData<EntityData> data = entityService.findEntityDataByQuery(getTenantId(), getCustomerId(), buildEntityDataQuery());
-            if (data.getData().isEmpty()) {
-                return 0;
-            }
-            entityIds = data.getData().stream().map(EntityData::getEntityId).toList();
-        }
-        return (int) alarmService.countAlarmsByQuery(getTenantId(), getCustomerId(), query, entityIds);
-    }
-
     private EntityDataQuery buildEntityDataQuery() {
         EntityDataPageLink edpl = new EntityDataPageLink(maxEntitiesPerAlarmSubscription, 0, null,
                 new EntityDataSortOrder(new EntityKey(EntityKeyType.ENTITY_FIELD, ModelConstants.CREATED_TIME_PROPERTY)));
         return new EntityDataQuery(query.getEntityFilter(), edpl, null, null, query.getKeyFilters());
     }
+
+    private void resetInvocationCounter() {
+        alarmCountInvocationAttempts = 0;
+    }
+
+    public void createAlarmSubscriptions() {
+        for (EntityId entityId : entitiesIds) {
+            createAlarmSubscriptionForEntity(entityId);
+        }
+    }
+
+    private void createAlarmSubscriptionForEntity(EntityId entityId) {
+        int subIdx = sessionRef.getSessionSubIdSeq().incrementAndGet();
+        subToEntityIdMap.put(subIdx, entityId);
+        log.trace("[{}][{}][{}] Creating alarms subscription for [{}] ", serviceId, cmdId, subIdx, entityId);
+        TbAlarmsSubscription subscription = TbAlarmsSubscription.builder()
+                .serviceId(serviceId)
+                .sessionId(sessionRef.getSessionId())
+                .subscriptionId(subIdx)
+                .tenantId(sessionRef.getSecurityCtx().getTenantId())
+                .entityId(entityId)
+                .updateProcessor((sub, update) -> fetchAlarmCount())
+                .build();
+        localSubscriptionService.addSubscription(subscription, sessionRef);
+    }
+
+    public void clearAlarmSubscriptions() {
+        if (subToEntityIdMap != null) {
+            for (Integer subId : subToEntityIdMap.keySet()) {
+                localSubscriptionService.cancelSubscription(getTenantId(), getSessionId(), subId);
+            }
+            subToEntityIdMap.clear();
+        }
+    }
+
 }

application/src/test/java/org/thingsboard/server/controller/WebsocketApiTest.java

@@ -84,7 +84,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.
 @DaoSqlTest
 @TestPropertySource(properties = {
         "server.ws.alarms_per_alarm_status_subscription_cache_size=5",
-        "server.ws.dynamic_page_link.refresh_interval=2"
+        "server.ws.dynamic_page_link.refresh_interval=3"
 })
 public class WebsocketApiTest extends AbstractControllerTest {
     @Autowired
@@ -422,7 +422,7 @@ public class WebsocketApiTest extends AbstractControllerTest {
 
         alarm = doPost("/api/alarm", alarm, Alarm.class);
 
-        update = getWsClient().parseAlarmCountReply(getWsClient().waitForUpdate(3000));
+        update = getWsClient().parseAlarmCountReply(getWsClient().waitForUpdate(4000));
         Assert.assertEquals(1, update.getCmdId());
         Assert.assertEquals(1, update.getCount());
 

dao/src/main/java/org/thingsboard/server/dao/sql/query/DefaultAlarmQueryRepository.java

@@ -321,27 +321,35 @@ public class DefaultAlarmQueryRepository implements AlarmQueryRepository {
         if (query.isSearchPropagatedAlarms()) {
             ctx.append("select count(distinct(a.id)) from alarm_info a ");
             ctx.append(JOIN_ENTITY_ALARMS);
-            ctx.append("where a.tenant_id = :tenantId and ea.tenant_id = :tenantId");
-            ctx.addUuidParameter("tenantId", tenantId.getId());
-            if (customerId != null && !customerId.isNullUid()) {
-                ctx.append(" and a.customer_id = :customerId and ea.customer_id = :customerId");
-                ctx.addUuidParameter("customerId", customerId.getId());
-            }
             if (orderedEntityIds != null) {
+                if (orderedEntityIds.isEmpty()) {
+                    return 0;
+                }
                 ctx.addUuidListParameter("entity_filter_entity_ids", orderedEntityIds.stream().map(EntityId::getId).collect(Collectors.toList()));
-                ctx.append(" and ea.entity_id in (:entity_filter_entity_ids)");
+                ctx.append("where ea.entity_id in (:entity_filter_entity_ids)");
+            } else {
+                ctx.append("where a.tenant_id = :tenantId and ea.tenant_id = :tenantId");
+                ctx.addUuidParameter("tenantId", tenantId.getId());
+                if (customerId != null && !customerId.isNullUid()) {
+                    ctx.append(" and a.customer_id = :customerId and ea.customer_id = :customerId");
+                    ctx.addUuidParameter("customerId", customerId.getId());
+                }
             }
         } else {
             ctx.append("select count(id) from alarm_info a ");
-            ctx.append("where a.tenant_id = :tenantId");
-            ctx.addUuidParameter("tenantId", tenantId.getId());
-            if (customerId != null && !customerId.isNullUid()) {
-                ctx.append(" and a.customer_id = :customerId");
-                ctx.addUuidParameter("customerId", customerId.getId());
-            }
             if (orderedEntityIds != null) {
+                if (orderedEntityIds.isEmpty()) {
+                    return 0;
+                }
                 ctx.addUuidListParameter("entity_filter_entity_ids", orderedEntityIds.stream().map(EntityId::getId).collect(Collectors.toList()));
-                ctx.append(" and a.originator_id in (:entity_filter_entity_ids)");
+                ctx.append("where a.originator_id in (:entity_filter_entity_ids)");
+            } else {
+                ctx.append("where a.tenant_id = :tenantId");
+                ctx.addUuidParameter("tenantId", tenantId.getId());
+                if (customerId != null && !customerId.isNullUid()) {
+                    ctx.append(" and a.customer_id = :customerId");
+                    ctx.addUuidParameter("customerId", customerId.getId());
+                }
             }
         }