From 36e16bb877de87a9da918559ea62ceca0879c675 Mon Sep 17 00:00:00 2001
From: Dmytro Skarzhynets <dskarzhynets@thingsboard.io>
Date: Tue, 1 Jul 2025 15:43:13 +0300
Subject: [PATCH] AI rule node: improve cache eviction

---
 .../ai/AiModelSettingsCacheEvictEvent.java    | 28 +++++++++++++++----
 .../dao/ai/AiModelSettingsServiceImpl.java    | 21 ++++++++------
 2 files changed, 34 insertions(+), 15 deletions(-)

diff --git a/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsCacheEvictEvent.java b/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsCacheEvictEvent.java
index 16ce64256b..3bf7dce9ba 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsCacheEvictEvent.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsCacheEvictEvent.java
@@ -15,15 +15,31 @@
  */
 package org.thingsboard.server.dao.ai;
 
-import org.thingsboard.server.common.data.id.AiModelSettingsId;
-import org.thingsboard.server.common.data.id.TenantId;
+import org.thingsboard.server.common.data.ai.AiModelSettings;
 
-import java.util.Set;
+import static java.util.Objects.requireNonNull;
+import static org.thingsboard.server.dao.ai.AiModelSettingsCacheEvictEvent.Deleted;
+import static org.thingsboard.server.dao.ai.AiModelSettingsCacheEvictEvent.Saved;
 
-record AiModelSettingsCacheEvictEvent(Set<AiModelSettingsCacheKey> keys) {
+sealed interface AiModelSettingsCacheEvictEvent permits Saved, Deleted {
+
+    AiModelSettingsCacheKey cacheKey();
+
+    record Saved(AiModelSettingsCacheKey cacheKey, AiModelSettings savedSettings) implements AiModelSettingsCacheEvictEvent {
+
+        public Saved {
+            requireNonNull(cacheKey);
+            requireNonNull(savedSettings);
+        }
+
+    }
+
+    record Deleted(AiModelSettingsCacheKey cacheKey) implements AiModelSettingsCacheEvictEvent {
+
+        public Deleted {
+            requireNonNull(cacheKey);
+        }
 
-    static AiModelSettingsCacheEvictEvent of(TenantId tenantId, AiModelSettingsId settingsId) {
-        return new AiModelSettingsCacheEvictEvent(Set.of(AiModelSettingsCacheKey.of(tenantId, settingsId)));
     }
 
 }
diff --git a/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsServiceImpl.java b/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsServiceImpl.java
index 411a67ca98..2dc3aba623 100644
--- a/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsServiceImpl.java
+++ b/dao/src/main/java/org/thingsboard/server/dao/ai/AiModelSettingsServiceImpl.java
@@ -15,7 +15,6 @@
  */
 package org.thingsboard.server.dao.ai;
 
-import com.google.common.collect.Sets;
 import com.google.common.util.concurrent.FluentFuture;
 import lombok.RequiredArgsConstructor;
 import org.springframework.stereotype.Service;
@@ -38,7 +37,6 @@ import org.thingsboard.server.dao.sql.JpaExecutorService;
 
 import java.util.List;
 import java.util.Optional;
-import java.util.Set;
 
 import static org.thingsboard.server.dao.service.Validator.validatePageLink;
 
@@ -54,7 +52,14 @@ class AiModelSettingsServiceImpl extends CachedVersionedEntityService<AiModelSet
     @Override
     @TransactionalEventListener
     public void handleEvictEvent(AiModelSettingsCacheEvictEvent event) {
-        cache.evict(event.keys());
+        var cacheKey = event.cacheKey();
+        if (event instanceof AiModelSettingsCacheEvictEvent.Saved savedEvent) {
+            cache.put(cacheKey, savedEvent.savedSettings());
+        } else if (event instanceof AiModelSettingsCacheEvictEvent.Deleted) {
+            cache.evict(cacheKey);
+        } else {
+            throw new UnsupportedOperationException("Unsupported event type: " + event.getClass().getSimpleName());
+        }
     }
 
     @Override
@@ -81,7 +86,8 @@ class AiModelSettingsServiceImpl extends CachedVersionedEntityService<AiModelSet
                 .broadcastEvent(true)
                 .build());
 
-        publishEvictEvent(AiModelSettingsCacheEvictEvent.of(savedSettings.getTenantId(), savedSettings.getId()));
+        var cacheKey = AiModelSettingsCacheKey.of(savedSettings.getTenantId(), savedSettings.getId());
+        publishEvictEvent(new AiModelSettingsCacheEvictEvent.Saved(cacheKey, savedSettings));
 
         return savedSettings;
     }
@@ -139,7 +145,7 @@ class AiModelSettingsServiceImpl extends CachedVersionedEntityService<AiModelSet
         boolean deleted = aiModelSettingsDao.deleteByTenantIdAndId(tenantId, settingsId);
         if (deleted) {
             publishDeleteEvent(toDeleteOpt.get());
-            publishEvictEvent(AiModelSettingsCacheEvictEvent.of(tenantId, settingsId));
+            publishEvictEvent(new AiModelSettingsCacheEvictEvent.Deleted(AiModelSettingsCacheKey.of(tenantId, settingsId)));
         }
         return deleted;
     }
@@ -154,13 +160,10 @@ class AiModelSettingsServiceImpl extends CachedVersionedEntityService<AiModelSet
 
         aiModelSettingsDao.deleteByTenantId(tenantId);
 
-        Set<AiModelSettingsCacheKey> cacheKeys = Sets.newHashSetWithExpectedSize(toDelete.size());
         toDelete.forEach(settings -> {
             publishDeleteEvent(settings);
-            cacheKeys.add(AiModelSettingsCacheKey.of(settings.getTenantId(), settings.getId()));
+            publishEvictEvent(new AiModelSettingsCacheEvictEvent.Deleted(AiModelSettingsCacheKey.of(settings.getTenantId(), settings.getId())));
         });
-
-        publishEvictEvent(new AiModelSettingsCacheEvictEvent(cacheKeys));
     }
 
     private void publishDeleteEvent(AiModelSettings settings) {