Caching script instance instead of script factory

Author: eyalkoren

server/src/main/java/org/elasticsearch/ingest/ConditionalProcessor.java

@@ -11,6 +11,7 @@
 
 import org.elasticsearch.common.logging.DeprecationCategory;
 import org.elasticsearch.common.logging.DeprecationLogger;
+import org.elasticsearch.script.CtxMapWrapper;
 import org.elasticsearch.script.DynamicMap;
 import org.elasticsearch.script.IngestConditionalScript;
 import org.elasticsearch.script.Script;
@@ -56,7 +57,8 @@ public class ConditionalProcessor extends AbstractProcessor implements WrappingP
     private final Processor processor;
     private final IngestMetric metric;
     private final LongSupplier relativeTimeProvider;
-    private final IngestConditionalScript.Factory precompiledConditionScriptFactory;
+    private final IngestConditionalScript precompiledConditionScript;
+    private final CtxMapWrapper ctxMapWrapper;
 
     ConditionalProcessor(String tag, String description, Script script, ScriptService scriptService, Processor processor) {
         this(tag, description, script, scriptService, processor, System::nanoTime);
@@ -76,12 +78,15 @@ public class ConditionalProcessor extends AbstractProcessor implements WrappingP
         this.processor = processor;
         this.metric = new IngestMetric();
         this.relativeTimeProvider = relativeTimeProvider;
+        this.ctxMapWrapper = new CtxMapWrapper();
 
         try {
+            final IngestConditionalScript.Factory factory = scriptService.compile(script, IngestConditionalScript.CONTEXT);
             if (ScriptType.INLINE.equals(script.getType())) {
-                precompiledConditionScriptFactory = scriptService.compile(script, IngestConditionalScript.CONTEXT);
+                precompiledConditionScript = factory.newInstance(script.getParams(), ctxMapWrapper);
             } else {
-                precompiledConditionScriptFactory = null;
+                // stored script, so will have to compile at runtime
+                precompiledConditionScript = null;
             }
         } catch (ScriptException e) {
             throw newConfigurationException(TYPE, tag, null, e);
@@ -139,12 +144,17 @@ public void execute(IngestDocument ingestDocument, BiConsumer<IngestDocument, Ex
     }
 
     boolean evaluate(IngestDocument ingestDocument) {
-        IngestConditionalScript.Factory factory = precompiledConditionScriptFactory;
-        if (factory == null) {
-            factory = scriptService.compile(condition, IngestConditionalScript.CONTEXT);
+        IngestConditionalScript script = precompiledConditionScript;
+        if (script == null) {
+            IngestConditionalScript.Factory factory = scriptService.compile(condition, IngestConditionalScript.CONTEXT);
+            script = factory.newInstance(condition.getParams(), ctxMapWrapper);
+        }
+        ctxMapWrapper.setCtxMap(ingestDocument.getCtxMap());
+        try {
+            return script.execute(new UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)));
+        } finally {
+            ctxMapWrapper.clearCtxMap();
         }
-        return factory.newInstance(condition.getParams(), ingestDocument.getCtxMap())
-            .execute(new UnmodifiableIngestData(new DynamicMap(ingestDocument.getSourceAndMetadata(), FUNCTIONS)));
     }
 
     public Processor getInnerProcessor() {

server/src/main/java/org/elasticsearch/script/CtxMapWrapper.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the "Elastic License
+ * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side
+ * Public License v 1"; you may not use this file except in compliance with, at
+ * your election, the "Elastic License 2.0", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+package org.elasticsearch.script;
+
+/**
+ * A wrapper for a {@link CtxMap} that allows for ad-hoc setting for script execution.
+ * This allows for precompilation of scripts that can be executed with different contexts.
+ * The wrapped {@link CtxMap} should be cleared after use to avoid leaks.
+ */
+public class CtxMapWrapper {
+    private volatile CtxMap<?> ctxMap;
+
+    public CtxMap<?> getCtxMap() {
+        if (ctxMap == null) {
+            throw new IllegalStateException("CtxMap is not set");
+        }
+        return ctxMap;
+    }
+
+    public void setCtxMap(CtxMap<?> ctxMap) {
+        this.ctxMap = ctxMap;
+    }
+
+    public void clearCtxMap() {
+        this.ctxMap = null;
+    }
+}

server/src/main/java/org/elasticsearch/script/IngestConditionalScript.java

@@ -15,10 +15,12 @@
 
 /**
  * A script used by {@link org.elasticsearch.ingest.ConditionalProcessor}.
+ * To properly expose the {@link SourceMapFieldScript#field(String)} API, make sure to provide a valid {@link CtxMap} before execution
+ * through the {@link CtxMapWrapper} passed to the constructor and make sure to clear it after use to avoid leaks.
  */
 public abstract class IngestConditionalScript extends SourceMapFieldScript {
 
-    public static final String[] PARAMETERS = { "ctxdsf" /*todo change param name*/ };
+    public static final String[] PARAMETERS = { "runtime-params" };
 
     /** The context used to compile {@link IngestConditionalScript} factories. */
     public static final ScriptContext<Factory> CONTEXT = new ScriptContext<>(
@@ -33,8 +35,8 @@ public abstract class IngestConditionalScript extends SourceMapFieldScript {
     /** The generic runtime parameters for the script. */
     private final Map<String, Object> params;
 
-    public IngestConditionalScript(Map<String, Object> params, CtxMap<?> ctxMap) {
-        super(ctxMap);
+    public IngestConditionalScript(Map<String, Object> params, CtxMapWrapper ctxMapWrapper) {
+        super(ctxMapWrapper);
         this.params = params;
     }
 
@@ -46,6 +48,6 @@ public Map<String, Object> getParams() {
     public abstract boolean execute(Map<String, Object> params);
 
     public interface Factory {
-        IngestConditionalScript newInstance(Map<String, Object> params, CtxMap<?> ctxMap);
+        IngestConditionalScript newInstance(Map<String, Object> params, CtxMapWrapper ctxMapWrapper);
     }
 }

server/src/main/java/org/elasticsearch/script/SourceMapFieldScript.java

@@ -18,23 +18,23 @@
  * These scripts provide {@code ctx} for backwards compatibility and expose {@link Metadata}.
  */
 public abstract class SourceMapFieldScript {
-    protected final CtxMap<?> ctxMap;
+    protected final CtxMapWrapper ctxMapWrapper;
 
-    public SourceMapFieldScript(CtxMap<?> ctxMap) {
-        this.ctxMap = ctxMap;
+    public SourceMapFieldScript(CtxMapWrapper ctxMapWrapper) {
+        this.ctxMapWrapper = ctxMapWrapper;
     }
 
     /** Provides backwards compatibility access to ctx */
     public Map<String, Object> getCtx() {
-        return ctxMap;
+        return ctxMapWrapper.getCtxMap();
     }
 
     /** Return the metadata for this script */
     public Metadata metadata() {
-        return ctxMap.getMetadata();
+        return ctxMapWrapper.getCtxMap().getMetadata();
     }
 
     public SourceMapField field(String path) {
-        return new SourceMapField(path, ctxMap::getSource);
+        return new SourceMapField(path, () -> ctxMapWrapper.getCtxMap().getSource());
     }
 }

server/src/test/java/org/elasticsearch/ingest/ConditionalProcessorTests.java

@@ -226,7 +226,7 @@ public boolean execute(Map<String, Object> params) {
     public void testRuntimeError() {
         ScriptService scriptService = MockScriptService.singleContext(
             IngestConditionalScript.CONTEXT,
-            code -> (params, ctxMap) -> new IngestConditionalScript(params, ctxMap) {
+            code -> (params, ctxMapWrapper) -> new IngestConditionalScript(params, ctxMapWrapper) {
                 @Override
                 public boolean execute(Map<String, Object> params) {
                     throw new IllegalArgumentException("runtime problem");

test/framework/src/main/java/org/elasticsearch/script/MockScriptEngine.java

@@ -169,7 +169,10 @@ public void execute() {
         } else if (context.instanceClazz.equals(AggregationScript.class)) {
             return context.factoryClazz.cast(new MockAggregationScript(script));
         } else if (context.instanceClazz.equals(IngestConditionalScript.class)) {
-            IngestConditionalScript.Factory factory = (parameters, ctxMap) -> new IngestConditionalScript(parameters, ctxMap) {
+            IngestConditionalScript.Factory factory = (parameters, ctxMapWrapper) -> new IngestConditionalScript(
+                parameters,
+                ctxMapWrapper
+            ) {
                 @Override
                 public boolean execute(Map<String, Object> params) {
                     return (boolean) script.apply(params);