From 142a422b7a630d6653195c579f9535388a2c0ac8 Mon Sep 17 00:00:00 2001
From: vparomskiy <paromskiy@gmail.com>
Date: Tue, 28 Aug 2018 16:54:31 +0300
Subject: [PATCH] deduplicate JS scripts in JsScriptEngine

---
 .../AbstractNashornJsSandboxService.java      | 59 ++++++++++++++-----
 .../script/RuleNodeJsScriptEngine.java        |  2 +-
 2 files changed, 45 insertions(+), 16 deletions(-)

diff --git a/application/src/main/java/org/thingsboard/server/service/script/AbstractNashornJsSandboxService.java b/application/src/main/java/org/thingsboard/server/service/script/AbstractNashornJsSandboxService.java
index 7f274ec6d7..724f652609 100644
--- a/application/src/main/java/org/thingsboard/server/service/script/AbstractNashornJsSandboxService.java
+++ b/application/src/main/java/org/thingsboard/server/service/script/AbstractNashornJsSandboxService.java
@@ -22,6 +22,7 @@ import delight.nashornsandbox.NashornSandbox;
 import delight.nashornsandbox.NashornSandboxes;
 import jdk.nashorn.api.scripting.NashornScriptEngineFactory;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.tuple.Pair;
 
 import javax.annotation.PostConstruct;
 import javax.annotation.PreDestroy;
@@ -42,9 +43,10 @@ public abstract class AbstractNashornJsSandboxService implements JsSandboxServic
     private ScriptEngine engine;
     private ExecutorService monitorExecutorService;
 
-    private Map<UUID, String> functionsMap = new ConcurrentHashMap<>();
-
-    private Map<UUID,AtomicInteger> blackListedFunctions = new ConcurrentHashMap<>();
+    private final Map<UUID, String> functionsMap = new ConcurrentHashMap<>();
+    private final Map<UUID,AtomicInteger> blackListedFunctions = new ConcurrentHashMap<>();
+    private final Map<String, Pair<UUID, AtomicInteger>> scriptToId = new ConcurrentHashMap<>();
+    private final Map<UUID, AtomicInteger> scriptIdToCount = new ConcurrentHashMap<>();
 
     @PostConstruct
     public void init() {
@@ -78,19 +80,27 @@ public abstract class AbstractNashornJsSandboxService implements JsSandboxServic
 
     @Override
     public ListenableFuture<UUID> eval(JsScriptType scriptType, String scriptBody, String... argNames) {
-        UUID scriptId = UUID.randomUUID();
-        String functionName = "invokeInternal_" + scriptId.toString().replace('-','_');
-        String jsScript = generateJsScript(scriptType, functionName, scriptBody, argNames);
-        try {
-            if (useJsSandbox()) {
-                sandbox.eval(jsScript);
-            } else {
-                engine.eval(jsScript);
+        Pair<UUID, AtomicInteger> deduplicated = deduplicate(scriptType, scriptBody);
+        UUID scriptId = deduplicated.getLeft();
+        AtomicInteger duplicateCount = deduplicated.getRight();
+
+        if(duplicateCount.compareAndSet(0, 1)) {
+            String functionName = "invokeInternal_" + scriptId.toString().replace('-', '_');
+            String jsScript = generateJsScript(scriptType, functionName, scriptBody, argNames);
+            try {
+                if (useJsSandbox()) {
+                    sandbox.eval(jsScript);
+                } else {
+                    engine.eval(jsScript);
+                }
+                functionsMap.put(scriptId, functionName);
+            } catch (Exception e) {
+                duplicateCount.decrementAndGet();
+                log.warn("Failed to compile JS script: {}", e.getMessage(), e);
+                return Futures.immediateFailedFuture(e);
             }
-            functionsMap.put(scriptId, functionName);
-        } catch (Exception e) {
-            log.warn("Failed to compile JS script: {}", e.getMessage(), e);
-            return Futures.immediateFailedFuture(e);
+        } else {
+            duplicateCount.incrementAndGet();
         }
         return Futures.immediateFuture(scriptId);
     }
@@ -122,6 +132,13 @@ public abstract class AbstractNashornJsSandboxService implements JsSandboxServic
 
     @Override
     public ListenableFuture<Void> release(UUID scriptId) {
+        AtomicInteger count = scriptIdToCount.get(scriptId);
+        if(count != null) {
+            if(count.decrementAndGet() > 0) {
+                return Futures.immediateFuture(null);
+            }
+        }
+
         String functionName = functionsMap.get(scriptId);
         if (functionName != null) {
             try {
@@ -156,4 +173,16 @@ public abstract class AbstractNashornJsSandboxService implements JsSandboxServic
                 throw new RuntimeException("No script factory implemented for scriptType: " + scriptType);
         }
     }
+
+    private Pair<UUID, AtomicInteger> deduplicate(JsScriptType scriptType, String scriptBody) {
+        Pair<UUID, AtomicInteger> precomputed = Pair.of(UUID.randomUUID(), new AtomicInteger());
+
+        Pair<UUID, AtomicInteger> pair = scriptToId.computeIfAbsent(deduplicateKey(scriptType, scriptBody), i -> precomputed);
+        AtomicInteger duplicateCount = scriptIdToCount.computeIfAbsent(pair.getLeft(), i -> pair.getRight());
+        return Pair.of(pair.getLeft(), duplicateCount);
+    }
+
+    private String deduplicateKey(JsScriptType scriptType, String scriptBody) {
+        return scriptType + "_" + scriptBody;
+    }
 }
diff --git a/application/src/main/java/org/thingsboard/server/service/script/RuleNodeJsScriptEngine.java b/application/src/main/java/org/thingsboard/server/service/script/RuleNodeJsScriptEngine.java
index 98b0e77a0a..2ba87ec8da 100644
--- a/application/src/main/java/org/thingsboard/server/service/script/RuleNodeJsScriptEngine.java
+++ b/application/src/main/java/org/thingsboard/server/service/script/RuleNodeJsScriptEngine.java
@@ -45,7 +45,7 @@ public class RuleNodeJsScriptEngine implements org.thingsboard.rule.engine.api.S
         try {
             this.scriptId = this.sandboxService.eval(JsScriptType.RULE_NODE_SCRIPT, script, argNames).get();
         } catch (Exception e) {
-            throw new IllegalArgumentException("Can't compile script: " + e.getMessage());
+            throw new IllegalArgumentException("Can't compile script: " + e.getMessage(), e);
         }
     }