From 78ac876103ff675c4949a1a15c8b4aecc4786e14 Mon Sep 17 00:00:00 2001
From: Andrii Shvaika <ashvayka@thingsboard.io>
Date: Thu, 29 Oct 2020 13:26:15 +0200
Subject: [PATCH] Bug fixes for Device Profile node

---
 .../server/actors/ruleChain/DefaultTbContext.java |  8 ++++++++
 .../server/dao/rule/RuleNodeStateService.java     |  2 ++
 .../data/device/profile/DeviceProfileAlarm.java   |  3 ++-
 .../server/dao/rule/BaseRuleNodeStateService.java | 14 ++++++++++++++
 .../server/dao/rule/RuleNodeStateDao.java         |  2 ++
 .../server/dao/sql/rule/JpaRuleNodeStateDao.java  |  6 ++++++
 .../dao/sql/rule/RuleNodeStateRepository.java     |  1 +
 .../thingsboard/rule/engine/api/TbContext.java    |  2 ++
 .../rule/engine/profile/TbDeviceProfileNode.java  | 15 +++++++++++----
 .../engine/profile/TbDeviceProfileNodeTest.java   |  3 ++-
 10 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java b/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java
index ca401596f6..72177c02f2 100644
--- a/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java
+++ b/application/src/main/java/org/thingsboard/server/actors/ruleChain/DefaultTbContext.java
@@ -480,6 +480,14 @@ class DefaultTbContext implements TbContext {
         mainCtx.getRuleNodeStateService().removeByRuleNodeId(getTenantId(), getSelfId());
     }
 
+    @Override
+    public void removeRuleNodeStateForEntity(EntityId entityId) {
+        if (log.isDebugEnabled()) {
+            log.debug("[{}][{}][{}] Remove Rule Node State for entity.", getTenantId(), getSelfId(), entityId);
+        }
+        mainCtx.getRuleNodeStateService().removeByRuleNodeIdAndEntityId(getTenantId(), getSelfId(), entityId);
+    }
+
     @Override
     public void addDeviceProfileListeners(Consumer<DeviceProfile> profileListener, BiConsumer<DeviceId, DeviceProfile> deviceListener) {
         mainCtx.getDeviceProfileCache().addListener(getTenantId(), getSelfId(), profileListener, deviceListener);
diff --git a/common/dao-api/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateService.java b/common/dao-api/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateService.java
index d5ad9bbbb6..924cad0884 100644
--- a/common/dao-api/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateService.java
+++ b/common/dao-api/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateService.java
@@ -31,4 +31,6 @@ public interface RuleNodeStateService {
     RuleNodeState save(TenantId tenantId, RuleNodeState ruleNodeState);
 
     void removeByRuleNodeId(TenantId tenantId, RuleNodeId selfId);
+
+    void removeByRuleNodeIdAndEntityId(TenantId tenantId, RuleNodeId selfId, EntityId entityId);
 }
diff --git a/common/data/src/main/java/org/thingsboard/server/common/data/device/profile/DeviceProfileAlarm.java b/common/data/src/main/java/org/thingsboard/server/common/data/device/profile/DeviceProfileAlarm.java
index b6437ae7bb..17035f337b 100644
--- a/common/data/src/main/java/org/thingsboard/server/common/data/device/profile/DeviceProfileAlarm.java
+++ b/common/data/src/main/java/org/thingsboard/server/common/data/device/profile/DeviceProfileAlarm.java
@@ -18,6 +18,7 @@ package org.thingsboard.server.common.data.device.profile;
 import lombok.Data;
 import org.thingsboard.server.common.data.alarm.AlarmSeverity;
 
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 
@@ -27,7 +28,7 @@ public class DeviceProfileAlarm {
     private String id;
     private String alarmType;
 
-    private Map<AlarmSeverity, AlarmRule> createRules;
+    private LinkedHashMap<AlarmSeverity, AlarmRule> createRules;
     private AlarmRule clearRule;
 
     // Hidden in advanced settings
diff --git a/dao/src/main/java/org/thingsboard/server/dao/rule/BaseRuleNodeStateService.java b/dao/src/main/java/org/thingsboard/server/dao/rule/BaseRuleNodeStateService.java
index 1d7e8af03f..7bc77e14a3 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/rule/BaseRuleNodeStateService.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/rule/BaseRuleNodeStateService.java
@@ -79,6 +79,20 @@ public class BaseRuleNodeStateService extends AbstractEntityService implements R
         ruleNodeStateDao.removeByRuleNodeId(ruleNodeId.getId());
     }
 
+    @Override
+    public void removeByRuleNodeIdAndEntityId(TenantId tenantId, RuleNodeId ruleNodeId, EntityId entityId) {
+        if (tenantId == null) {
+            throw new DataValidationException("Tenant id should be specified!.");
+        }
+        if (ruleNodeId == null) {
+            throw new DataValidationException("Rule node id should be specified!.");
+        }
+        if (entityId == null) {
+            throw new DataValidationException("Entity id should be specified!.");
+        }
+        ruleNodeStateDao.removeByRuleNodeIdAndEntityId(ruleNodeId.getId(), entityId.getId());
+    }
+
     public RuleNodeState saveOrUpdate(TenantId tenantId, RuleNodeState ruleNodeState, boolean update) {
         try {
             if (update) {
diff --git a/dao/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateDao.java b/dao/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateDao.java
index 89c2baea7f..9f304dd828 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateDao.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/rule/RuleNodeStateDao.java
@@ -35,4 +35,6 @@ public interface RuleNodeStateDao extends Dao<RuleNodeState> {
     RuleNodeState findByRuleNodeIdAndEntityId(UUID ruleNodeId, UUID entityId);
 
     void removeByRuleNodeId(UUID ruleNodeId);
+
+    void removeByRuleNodeIdAndEntityId(UUID ruleNodeId, UUID entityId);
 }
diff --git a/dao/src/main/java/org/thingsboard/server/dao/sql/rule/JpaRuleNodeStateDao.java b/dao/src/main/java/org/thingsboard/server/dao/sql/rule/JpaRuleNodeStateDao.java
index c61e58353c..e32f90a645 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/sql/rule/JpaRuleNodeStateDao.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/sql/rule/JpaRuleNodeStateDao.java
@@ -63,4 +63,10 @@ public class JpaRuleNodeStateDao extends JpaAbstractDao<RuleNodeStateEntity, Rul
     public void removeByRuleNodeId(UUID ruleNodeId) {
         ruleNodeStateRepository.removeByRuleNodeId(ruleNodeId);
     }
+
+    @Transactional
+    @Override
+    public void removeByRuleNodeIdAndEntityId(UUID ruleNodeId, UUID entityId) {
+        ruleNodeStateRepository.removeByRuleNodeIdAndEntityId(ruleNodeId, entityId);
+    }
 }
diff --git a/dao/src/main/java/org/thingsboard/server/dao/sql/rule/RuleNodeStateRepository.java b/dao/src/main/java/org/thingsboard/server/dao/sql/rule/RuleNodeStateRepository.java
index 902bb0ddb9..7bc35adaae 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/sql/rule/RuleNodeStateRepository.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/sql/rule/RuleNodeStateRepository.java
@@ -36,4 +36,5 @@ public interface RuleNodeStateRepository extends PagingAndSortingRepository<Rule
 
     void removeByRuleNodeId(@Param("ruleNodeId") UUID ruleNodeId);
 
+    void removeByRuleNodeIdAndEntityId(@Param("ruleNodeId") UUID ruleNodeId, @Param("entityId") UUID entityId);
 }
diff --git a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java
index 24af4f0d22..a1affda96e 100644
--- a/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java
+++ b/rule-engine/rule-engine-api/src/main/java/org/thingsboard/rule/engine/api/TbContext.java
@@ -223,6 +223,8 @@ public interface TbContext {
 
     RuleNodeState findRuleNodeStateForEntity(EntityId entityId);
 
+    void removeRuleNodeStateForEntity(EntityId entityId);
+
     RuleNodeState saveRuleNodeState(RuleNodeState state);
 
     void clearRuleNodeStates();
diff --git a/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNode.java b/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNode.java
index 0766265368..fa7538bce6 100644
--- a/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNode.java
+++ b/rule-engine/rule-engine-components/src/main/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNode.java
@@ -117,7 +117,7 @@ public class TbDeviceProfileNode implements TbNode {
             if (data.has("profileId")) {
                 invalidateDeviceProfileCache(deviceId, new DeviceProfileId(UUID.fromString(data.get("deviceProfileId").asText())));
             } else {
-                deviceStates.remove(deviceId);
+                removeDeviceState(deviceId);
             }
 
         } else {
@@ -126,7 +126,7 @@ public class TbDeviceProfileNode implements TbNode {
                 if (msg.getType().equals(DataConstants.ENTITY_UPDATED)) {
                     invalidateDeviceProfileCache(deviceId, msg.getData());
                 } else if (msg.getType().equals(DataConstants.ENTITY_DELETED)) {
-                    deviceStates.remove(deviceId);
+                    removeDeviceState(deviceId);
                 } else {
                     DeviceState deviceState = getOrCreateDeviceState(ctx, deviceId, null);
                     if (deviceState != null) {
@@ -209,7 +209,7 @@ public class TbDeviceProfileNode implements TbNode {
             DeviceProfileId currentProfileId = deviceState.getProfileId();
             Device device = JacksonUtil.fromString(deviceJson, Device.class);
             if (!currentProfileId.equals(device.getDeviceProfileId())) {
-                deviceStates.remove(deviceId);
+                removeDeviceState(deviceId);
             }
         }
     }
@@ -218,8 +218,15 @@ public class TbDeviceProfileNode implements TbNode {
         DeviceState deviceState = deviceStates.get(deviceId);
         if (deviceState != null) {
             if (!deviceState.getProfileId().equals(deviceProfileId)) {
-                deviceStates.remove(deviceId);
+                removeDeviceState(deviceId);
             }
         }
     }
+
+    private void removeDeviceState(DeviceId deviceId) {
+        DeviceState state = deviceStates.remove(deviceId);
+        if (config.isPersistAlarmRulesState() && (state != null || !config.isFetchAlarmRulesStateOnStart())) {
+            ctx.removeRuleNodeStateForEntity(deviceId);
+        }
+    }
 }
diff --git a/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNodeTest.java b/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNodeTest.java
index 8562b52bb9..ae6bd7ebbd 100644
--- a/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNodeTest.java
+++ b/rule-engine/rule-engine-components/src/test/java/org/thingsboard/rule/engine/profile/TbDeviceProfileNodeTest.java
@@ -56,6 +56,7 @@ import org.thingsboard.server.dao.alarm.AlarmService;
 import org.thingsboard.server.dao.timeseries.TimeseriesService;
 
 import java.util.Collections;
+import java.util.LinkedHashMap;
 import java.util.UUID;
 
 import static org.mockito.Mockito.verify;
@@ -139,7 +140,7 @@ public class TbDeviceProfileNodeTest {
         DeviceProfileAlarm dpa = new DeviceProfileAlarm();
         dpa.setId("highTemperatureAlarmID");
         dpa.setAlarmType("highTemperatureAlarm");
-        dpa.setCreateRules(Collections.singletonMap(AlarmSeverity.CRITICAL, alarmRule));
+        dpa.setCreateRules(new LinkedHashMap<>(Collections.singletonMap(AlarmSeverity.CRITICAL, alarmRule)));
 
         KeyFilter lowTempFilter = new KeyFilter();
         lowTempFilter.setKey(new EntityKey(EntityKeyType.TIME_SERIES, "temperature"));