Merge pull request #7615 from ViacheslavKlimov/improvement/mvel-invoke-service

[3.4.2] MVEL scripts caching

application/src/main/resources/thingsboard.yml

@@ -625,6 +625,7 @@ mvel:
   max_black_list_duration_sec: "${MVEL_MAX_BLACKLIST_DURATION_SEC:60}"
   # Specify thread pool size for javascript executor service
   thread_pool_size: "${MVEL_THREAD_POOL_SIZE:50}"
+  compiled_scripts_cache_size: "${MVEL_COMPILED_SCRIPTS_CACHE_SIZE:1000}"
   stats:
     enabled: "${TB_MVEL_STATS_ENABLED:false}"
     print_interval_ms: "${TB_MVEL_STATS_PRINT_INTERVAL_MS:10000}"

application/src/test/java/org/thingsboard/server/service/script/MvelInvokeServiceTest.java

@@ -16,6 +16,7 @@
 package org.thingsboard.server.service.script;
 
 import com.fasterxml.jackson.databind.node.ObjectNode;
+import com.github.benmanes.caffeine.cache.Cache;
 import org.junit.Assert;
 import org.junit.jupiter.api.Test;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -24,15 +25,22 @@ import org.springframework.test.context.TestPropertySource;
 import org.thingsboard.common.util.JacksonUtil;
 import org.thingsboard.script.api.ScriptType;
 import org.thingsboard.script.api.mvel.MvelInvokeService;
+import org.thingsboard.script.api.mvel.MvelScript;
 import org.thingsboard.server.common.data.id.TenantId;
 import org.thingsboard.server.controller.AbstractControllerTest;
 import org.thingsboard.server.dao.service.DaoSqlTest;
 
+import java.io.Serializable;
+import java.lang.reflect.Field;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 @DaoSqlTest
@@ -41,6 +49,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy;
         "mvel.max_total_args_size=50",
         "mvel.max_result_size=50",
         "mvel.max_errors=2",
+        "mvel.compiled_scripts_cache_size=100"
 })
 class MvelInvokeServiceTest extends AbstractControllerTest {
 
@@ -110,6 +119,89 @@ class MvelInvokeServiceTest extends AbstractControllerTest {
         assertThatScriptIsBlocked(scriptId);
     }
 
+    @Test
+    void givenScriptsWithSameBody_thenCompileAndCacheOnlyOnce() throws Exception {
+        String script = "return msg.temperature > 20;";
+        List<UUID> scriptsIds = new ArrayList<>();
+        for (int i = 0; i < 100; i++) {
+            UUID scriptId = evalScript(script);
+            scriptsIds.add(scriptId);
+        }
+
+        Map<UUID, String> scriptIdToHash = getFieldValue(invokeService, "scriptIdToHash");
+        Map<String, MvelScript> scriptMap = getFieldValue(invokeService, "scriptMap");
+        Cache<String, Serializable> compiledScriptsCache = getFieldValue(invokeService, "compiledScriptsCache");
+
+        String scriptHash = scriptIdToHash.get(scriptsIds.get(0));
+
+        assertThat(scriptsIds.stream().map(scriptIdToHash::get)).containsOnly(scriptHash);
+        assertThat(scriptMap).containsKey(scriptHash);
+        assertThat(compiledScriptsCache.getIfPresent(scriptHash)).isNotNull();
+    }
+
+    @Test
+    public void whenReleasingScript_thenCheckForScriptHashUsages() throws Exception {
+        String script = "return msg.temperature > 20;";
+        List<UUID> scriptsIds = new ArrayList<>();
+        for (int i = 0; i < 10; i++) {
+            UUID scriptId = evalScript(script);
+            scriptsIds.add(scriptId);
+        }
+
+        Map<UUID, String> scriptIdToHash = getFieldValue(invokeService, "scriptIdToHash");
+        Map<String, MvelScript> scriptMap = getFieldValue(invokeService, "scriptMap");
+        Cache<String, Serializable> compiledScriptsCache = getFieldValue(invokeService, "compiledScriptsCache");
+
+        String scriptHash = scriptIdToHash.get(scriptsIds.get(0));
+        for (int i = 0; i < 9; i++) {
+            UUID scriptId = scriptsIds.get(i);
+            assertThat(scriptIdToHash).containsKey(scriptId);
+            invokeService.release(scriptId);
+            assertThat(scriptIdToHash).doesNotContainKey(scriptId);
+        }
+        assertThat(scriptMap).containsKey(scriptHash);
+        assertThat(compiledScriptsCache.getIfPresent(scriptHash)).isNotNull();
+
+        invokeService.release(scriptsIds.get(9));
+        assertThat(scriptMap).doesNotContainKey(scriptHash);
+        assertThat(compiledScriptsCache.getIfPresent(scriptHash)).isNull();
+    }
+
+    @Test
+    public void whenCompiledScriptsCacheIsTooBig_thenRemoveRarelyUsedScripts() throws Exception {
+        Map<UUID, String> scriptIdToHash = getFieldValue(invokeService, "scriptIdToHash");
+        Cache<String, Serializable> compiledScriptsCache = getFieldValue(invokeService, "compiledScriptsCache");
+
+        List<UUID> scriptsIds = new ArrayList<>();
+        for (int i = 0; i < 110; i++) { // mvel.compiled_scripts_cache_size = 100
+            String script = "return msg.temperature > " + i;
+            UUID scriptId = evalScript(script);
+            scriptsIds.add(scriptId);
+
+            for (int j = 0; j < i; j++) {
+                invokeScript(scriptId, "{ \"temperature\": 12 }"); // so that scriptsIds is ordered by number of invocations
+            }
+        }
+
+        ConcurrentMap<String, Serializable> cache = compiledScriptsCache.asMap();
+
+        for (int i = 0; i < 10; i++) { // iterating rarely used scripts
+            UUID scriptId = scriptsIds.get(i);
+            String scriptHash = scriptIdToHash.get(scriptId);
+            assertThat(cache).doesNotContainKey(scriptHash);
+        }
+        for (int i = 10; i < 110; i++) {
+            UUID scriptId = scriptsIds.get(i);
+            String scriptHash = scriptIdToHash.get(scriptId);
+            assertThat(cache).containsKey(scriptHash);
+        }
+
+        UUID scriptRemovedFromCache = scriptsIds.get(0);
+        assertThat(compiledScriptsCache.getIfPresent(scriptIdToHash.get(scriptRemovedFromCache))).isNull();
+        invokeScript(scriptRemovedFromCache, "{ \"temperature\": 12 }");
+        assertThat(compiledScriptsCache.getIfPresent(scriptIdToHash.get(scriptRemovedFromCache))).isNotNull();
+    }
+
     private void assertThatScriptIsBlocked(UUID scriptId) {
         assertThatThrownBy(() -> {
             invokeScript(scriptId, "{}");
@@ -125,4 +217,10 @@ class MvelInvokeServiceTest extends AbstractControllerTest {
         return invokeService.invokeScript(TenantId.SYS_TENANT_ID, null, scriptId, msg, "{}", "POST_TELEMETRY_REQUEST").get().toString();
     }
 
+    private <T> T getFieldValue(Object target, String fieldName) throws Exception {
+        Field field = target.getClass().getDeclaredField(fieldName);
+        field.setAccessible(true);
+        return (T) field.get(target);
+    }
+
 }

common/script/script-api/pom.xml

@@ -56,6 +56,10 @@
             <groupId>com.google.code.gson</groupId>
             <artifactId>gson</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.github.ben-manes.caffeine</groupId>
+            <artifactId>caffeine</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.slf4j</groupId>
             <artifactId>slf4j-api</artifactId>

common/script/script-api/src/main/java/org/thingsboard/script/api/mvel/DefaultMvelInvokeService.java

@@ -15,6 +15,10 @@
  */
 package org.thingsboard.script.api.mvel;
 
+import com.github.benmanes.caffeine.cache.Cache;
+import com.github.benmanes.caffeine.cache.Caffeine;
+import com.google.common.hash.Hasher;
+import com.google.common.hash.Hashing;
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.ListeningExecutorService;
 import com.google.common.util.concurrent.MoreExecutors;
@@ -42,12 +46,15 @@ import org.thingsboard.server.common.stats.TbApiUsageStateClient;
 import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
 import java.io.Serializable;
+import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.Executor;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.regex.Pattern;
 
 @Slf4j
@@ -55,7 +62,10 @@ import java.util.regex.Pattern;
 @Service
 public class DefaultMvelInvokeService extends AbstractScriptInvokeService implements MvelInvokeService {
 
-    protected Map<UUID, MvelScript> scriptMap = new ConcurrentHashMap<>();
+    protected final Map<UUID, String> scriptIdToHash = new ConcurrentHashMap<>();
+    protected final Map<String, MvelScript> scriptMap = new ConcurrentHashMap<>();
+    protected Cache<String, Serializable> compiledScriptsCache;
+
     private SandboxedParserConfiguration parserConfig;
 
     private static final Pattern NEW_KEYWORD_PATTERN = Pattern.compile("new\\s");
@@ -92,8 +102,13 @@ public class DefaultMvelInvokeService extends AbstractScriptInvokeService implem
     @Value("${mvel.max_memory_limit_mb:8}")
     private long maxMemoryLimitMb;
 
+    @Value("${mvel.compiled_scripts_cache_size:1000}")
+    private int compiledScriptsCacheSize;
+
     private ListeningExecutorService executor;
 
+    private final Lock lock = new ReentrantLock();
+
     protected DefaultMvelInvokeService(Optional<TbApiUsageStateClient> apiUsageStateClient, Optional<TbApiUsageReportClient> apiUsageReportClient) {
         super(apiUsageStateClient, apiUsageReportClient);
     }
@@ -115,11 +130,14 @@ public class DefaultMvelInvokeService extends AbstractScriptInvokeService implem
         executor = MoreExecutors.listeningDecorator(ThingsBoardExecutors.newWorkStealingPool(threadPoolSize, "mvel-executor"));
         try {
             // Special command to warm up MVEL engine
-            Serializable script = MVEL.compileExpression("var warmUp = {}; warmUp", new SandboxedParserContext(parserConfig));
+            Serializable script = compileScript("var warmUp = {}; warmUp");
             MVEL.executeTbExpression(script, new ExecutionContext(parserConfig), Collections.emptyMap());
         } catch (Exception e) {
             // do nothing
         }
+        compiledScriptsCache = Caffeine.newBuilder()
+                .maximumSize(compiledScriptsCacheSize)
+                .build();
     }
 
     @PreDestroy
@@ -141,16 +159,26 @@ public class DefaultMvelInvokeService extends AbstractScriptInvokeService implem
 
     @Override
     protected boolean isScriptPresent(UUID scriptId) {
-        return scriptMap.containsKey(scriptId);
+        return scriptIdToHash.containsKey(scriptId);
     }
 
     @Override
     protected ListenableFuture<UUID> doEvalScript(TenantId tenantId, ScriptType scriptType, String scriptBody, UUID scriptId, String[] argNames) {
         return executor.submit(() -> {
             try {
-                Serializable compiledScript = MVEL.compileExpression(scriptBody, new SandboxedParserContext(parserConfig));
-                MvelScript script = new MvelScript(compiledScript, scriptBody, argNames);
-                scriptMap.put(scriptId, script);
+                String scriptHash = hash(scriptBody, argNames);
+                compiledScriptsCache.get(scriptHash, k -> {
+                    return compileScript(scriptBody);
+                });
+                lock.lock();
+                try {
+                    scriptIdToHash.put(scriptId, scriptHash);
+                    scriptMap.computeIfAbsent(scriptHash, k -> {
+                        return new MvelScript(scriptBody, argNames);
+                    });
+                } finally {
+                    lock.unlock();
+                }
                 return scriptId;
             } catch (Exception e) {
                 throw new TbScriptException(scriptId, TbScriptException.ErrorCode.COMPILATION, scriptBody, e);
@@ -162,12 +190,16 @@ public class DefaultMvelInvokeService extends AbstractScriptInvokeService implem
     protected MvelScriptExecutionTask doInvokeFunction(UUID scriptId, Object[] args) {
         ExecutionContext executionContext = new ExecutionContext(this.parserConfig, maxMemoryLimitMb * 1024 * 1024);
         return new MvelScriptExecutionTask(executionContext, executor.submit(() -> {
-            MvelScript script = scriptMap.get(scriptId);
-            if (script == null) {
+            String scriptHash = scriptIdToHash.get(scriptId);
+            if (scriptHash == null) {
                 throw new TbScriptException(scriptId, TbScriptException.ErrorCode.OTHER, null, new RuntimeException("Script not found!"));
             }
+            MvelScript script = scriptMap.get(scriptHash);
+            Serializable compiledScript = compiledScriptsCache.get(scriptHash, k -> {
+                return compileScript(script.getScriptBody());
+            });
             try {
-                return MVEL.executeTbExpression(script.getCompiledScript(), executionContext, script.createVars(args));
+                return MVEL.executeTbExpression(compiledScript, executionContext, script.createVars(args));
             } catch (ScriptMemoryOverflowException e) {
                 throw new TbScriptException(scriptId, TbScriptException.ErrorCode.OTHER, script.getScriptBody(), new RuntimeException("Script memory overflow!"));
             } catch (Exception e) {
@@ -178,6 +210,32 @@ public class DefaultMvelInvokeService extends AbstractScriptInvokeService implem
 
     @Override
     protected void doRelease(UUID scriptId) throws Exception {
-        scriptMap.remove(scriptId);
+        String scriptHash = scriptIdToHash.remove(scriptId);
+        if (scriptHash != null) {
+            lock.lock();
+            try {
+                if (!scriptIdToHash.containsValue(scriptHash)) {
+                    scriptMap.remove(scriptHash);
+                    compiledScriptsCache.invalidate(scriptHash);
+                }
+            } finally {
+                lock.unlock();
+            }
+        }
     }
+
+    private Serializable compileScript(String scriptBody) {
+        return MVEL.compileExpression(scriptBody, new SandboxedParserContext(parserConfig));
+    }
+
+    @SuppressWarnings("UnstableApiUsage")
+    protected String hash(String scriptBody, String[] argNames) {
+        Hasher hasher = Hashing.murmur3_128().newHasher();
+        hasher.putUnencodedChars(scriptBody);
+        for (String argName : argNames) {
+            hasher.putString(argName, StandardCharsets.UTF_8);
+        }
+        return hasher.hash().toString();
+    }
+
 }

common/script/script-api/src/main/java/org/thingsboard/script/api/mvel/MvelScript.java

@@ -24,7 +24,6 @@ import java.util.Map;
 @Data
 public class MvelScript {
 
-    private final Serializable compiledScript;
     private final String scriptBody;
     private final String[] argNames;