Scripting: Cache script results if deterministic (#50106) (#50329)

Cache results from queries that use scripts if they use only
deterministic API calls.  Nondeterministic API calls are marked in the
whitelist with the `@nondeterministic` annotation.  Examples are
`Math.random()` and `new Date()`.

Refs: #49466

Author: stu-elastic

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/NonDeterministicAnnotation.java

@@ -0,0 +1,29 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.spi.annotation;
+
+public class NonDeterministicAnnotation {
+
+    public static final String NAME = "nondeterministic";
+
+    public static final NonDeterministicAnnotation INSTANCE = new NonDeterministicAnnotation();
+
+    private NonDeterministicAnnotation() {}
+}

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/NonDeterministicAnnotationParser.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to Elasticsearch under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.elasticsearch.painless.spi.annotation;
+
+import java.util.Map;
+
+public class NonDeterministicAnnotationParser implements WhitelistAnnotationParser {
+
+    public static final NonDeterministicAnnotationParser INSTANCE = new NonDeterministicAnnotationParser();
+
+    private NonDeterministicAnnotationParser() {}
+
+    @Override
+    public Object parse(Map<String, String> arguments) {
+        if (arguments.isEmpty() == false) {
+            throw new IllegalArgumentException(
+                "unexpected parameters for [@" + NonDeterministicAnnotation.NAME + "] annotation, found " + arguments
+            );
+        }
+
+        return NonDeterministicAnnotation.INSTANCE;
+    }
+}

modules/lang-painless/spi/src/main/java/org/elasticsearch/painless/spi/annotation/WhitelistAnnotationParser.java

@@ -34,7 +34,8 @@ public interface WhitelistAnnotationParser {
     Map<String, WhitelistAnnotationParser> BASE_ANNOTATION_PARSERS = Collections.unmodifiableMap(
             Stream.of(
                     new AbstractMap.SimpleEntry<>(NoImportAnnotation.NAME, NoImportAnnotationParser.INSTANCE),
-                    new AbstractMap.SimpleEntry<>(DeprecatedAnnotation.NAME, DeprecatedAnnotationParser.INSTANCE)
+                    new AbstractMap.SimpleEntry<>(DeprecatedAnnotation.NAME, DeprecatedAnnotationParser.INSTANCE),
+                    new AbstractMap.SimpleEntry<>(NonDeterministicAnnotation.NAME, NonDeterministicAnnotationParser.INSTANCE)
             ).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue))
     );
 

modules/lang-painless/src/main/java/org/elasticsearch/painless/Compiler.java

@@ -26,7 +26,6 @@
 import org.elasticsearch.painless.spi.Whitelist;
 import org.objectweb.asm.util.Printer;
 
-import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.net.MalformedURLException;
 import java.net.URL;
@@ -205,13 +204,14 @@ private static void addFactoryMethod(Map<String, Class<?>> additionalClasses, Cl
      * @param name The name of the script.
      * @param source The source code for the script.
      * @param settings The CompilerSettings to be used during the compilation.
-     * @return An executable script that implements both a specified interface and is a subclass of {@link PainlessScript}
+     * @return The ScriptRoot used to compile
      */
-    Constructor<?> compile(Loader loader, Set<String> extractedVariables, String name, String source, CompilerSettings settings) {
+    ScriptRoot compile(Loader loader, Set<String> extractedVariables, String name, String source,
+            CompilerSettings settings) {
         ScriptClassInfo scriptClassInfo = new ScriptClassInfo(painlessLookup, scriptClass);
         SClass root = Walker.buildPainlessTree(scriptClassInfo, name, source, settings, painlessLookup, null);
         root.extractVariables(extractedVariables);
-        root.analyze(painlessLookup, settings);
+        ScriptRoot scriptRoot = root.analyze(painlessLookup, settings);
         Map<String, Object> statics = root.write();
 
         try {
@@ -225,8 +225,9 @@ Constructor<?> compile(Loader loader, Set<String> extractedVariables, String nam
                 clazz.getField(statik.getKey()).set(null, statik.getValue());
             }
 
-            return clazz.getConstructors()[0];
-        } catch (Exception exception) { // Catch everything to let the user know this is something caused internally.
+            return scriptRoot;
+        } catch (Exception exception) {
+            // Catch everything to let the user know this is something caused internally.
             throw new IllegalStateException("An internal error occurred attempting to define the script [" + name + "].", exception);
         }
     }

modules/lang-painless/src/main/java/org/elasticsearch/painless/PainlessScriptEngine.java

@@ -143,12 +143,13 @@ public Loader run() {
         });
 
         Set<String> extractedVariables = new HashSet<>();
-        compile(contextsToCompilers.get(context), loader, extractedVariables, scriptName, scriptSource, params);
+        ScriptRoot scriptRoot = compile(contextsToCompilers.get(context), loader, extractedVariables, scriptName, scriptSource, params);
 
         if (context.statefulFactoryClazz != null) {
-            return generateFactory(loader, context, extractedVariables, generateStatefulFactory(loader, context, extractedVariables));
+            return generateFactory(loader, context, extractedVariables, generateStatefulFactory(loader, context, extractedVariables),
+                scriptRoot);
         } else {
-            return generateFactory(loader, context, extractedVariables, WriterConstants.CLASS_TYPE);
+            return generateFactory(loader, context, extractedVariables, WriterConstants.CLASS_TYPE, scriptRoot);
         }
     }
 
@@ -270,14 +271,16 @@ private <T extends ScriptFactory> Type generateStatefulFactory(
      * @param context The {@link ScriptContext}'s semantics are used to define the factory class.
      * @param classType The type to be instaniated in the newFactory or newInstance method.  Depends
      *                  on whether a {@link ScriptContext#statefulFactoryClazz} is specified.
+     * @param scriptRoot the {@link ScriptRoot} used to do the compilation
      * @param <T> The factory class.
      * @return A factory class that will return script instances.
      */
     private <T extends ScriptFactory> T generateFactory(
         Loader loader,
         ScriptContext<T> context,
         Set<String> extractedVariables,
-        Type classType
+        Type classType,
+        ScriptRoot scriptRoot
     ) {
         int classFrames = ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS;
         int classAccess = Opcodes.ACC_PUBLIC | Opcodes.ACC_SUPER| Opcodes.ACC_FINAL;
@@ -330,8 +333,19 @@ private <T extends ScriptFactory> T generateFactory(
         adapter.endMethod();
 
         writeNeedsMethods(context.factoryClazz, writer, extractedVariables);
-        writer.visitEnd();
 
+        String methodName = "isResultDeterministic";
+        org.objectweb.asm.commons.Method isResultDeterministic = new org.objectweb.asm.commons.Method(methodName,
+            MethodType.methodType(boolean.class).toMethodDescriptorString());
+
+        GeneratorAdapter deterAdapter = new GeneratorAdapter(Opcodes.ASM5, isResultDeterministic,
+            writer.visitMethod(Opcodes.ACC_PUBLIC, methodName, isResultDeterministic.getDescriptor(), null, null));
+        deterAdapter.visitCode();
+        deterAdapter.push(scriptRoot.deterministic);
+        deterAdapter.returnValue();
+        deterAdapter.endMethod();
+
+        writer.visitEnd();
         Class<?> factory = loader.defineFactory(className.replace('/', '.'), writer.toByteArray());
 
         try {
@@ -363,19 +377,17 @@ private void writeNeedsMethods(Class<?> clazz, ClassWriter writer, Set<String> e
         }
     }
 
-    void compile(Compiler compiler, Loader loader, Set<String> extractedVariables,
+    ScriptRoot compile(Compiler compiler, Loader loader, Set<String> extractedVariables,
                  String scriptName, String source, Map<String, String> params) {
         final CompilerSettings compilerSettings = buildCompilerSettings(params);
 
         try {
             // Drop all permissions to actually compile the code itself.
-            AccessController.doPrivileged(new PrivilegedAction<Void>() {
+            return AccessController.doPrivileged(new PrivilegedAction<ScriptRoot>() {
                 @Override
-                public Void run() {
+                public ScriptRoot run() {
                     String name = scriptName == null ? source : scriptName;
-                    compiler.compile(loader, extractedVariables, name, source, compilerSettings);
-
-                    return null;
+                    return compiler.compile(loader, extractedVariables, name, source, compilerSettings);
                 }
             }, COMPILATION_CONTEXT);
             // Note that it is safe to catch any of the following errors since Painless is stateless.

modules/lang-painless/src/main/java/org/elasticsearch/painless/ScriptRoot.java

@@ -38,6 +38,7 @@ public class ScriptRoot {
 
     protected final FunctionTable functionTable = new FunctionTable();
     protected int syntheticCounter = 0;
+    protected boolean deterministic = true;
 
     public ScriptRoot(PainlessLookup painlessLookup, CompilerSettings compilerSettings, ScriptClassInfo scriptClassInfo, SClass classRoot) {
         this.painlessLookup = Objects.requireNonNull(painlessLookup);
@@ -72,4 +73,6 @@ public FunctionTable getFunctionTable() {
     public String getNextSyntheticName(String prefix) {
         return prefix + "$synthetic$" + syntheticCounter++;
     }
+
+    public void markNonDeterministic(boolean nondeterministic) { this.deterministic &= !nondeterministic; }
 }

modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessClassBinding.java

@@ -22,6 +22,7 @@
 import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 public class PainlessClassBinding {
@@ -31,13 +32,16 @@ public class PainlessClassBinding {
 
     public final Class<?> returnType;
     public final List<Class<?>> typeParameters;
+    public final Map<Class<?>, Object> annotations;
 
-    PainlessClassBinding(Constructor<?> javaConstructor, Method javaMethod, Class<?> returnType, List<Class<?>> typeParameters) {
+    PainlessClassBinding(Constructor<?> javaConstructor, Method javaMethod, Class<?> returnType, List<Class<?>> typeParameters,
+        Map<Class<?>, Object> annotations) {
         this.javaConstructor = javaConstructor;
         this.javaMethod = javaMethod;
 
         this.returnType = returnType;
         this.typeParameters = typeParameters;
+        this.annotations = annotations;
     }
 
     @Override

modules/lang-painless/src/main/java/org/elasticsearch/painless/lookup/PainlessConstructor.java

@@ -23,6 +23,7 @@
 import java.lang.invoke.MethodType;
 import java.lang.reflect.Constructor;
 import java.util.List;
+import java.util.Map;
 import java.util.Objects;
 
 public class PainlessConstructor {
@@ -31,12 +32,15 @@ public class PainlessConstructor {
     public final List<Class<?>> typeParameters;
     public final MethodHandle methodHandle;
     public final MethodType methodType;
+    public final Map<Class<?>, Object> annotations;
 
-    PainlessConstructor(Constructor<?> javaConstructor, List<Class<?>> typeParameters, MethodHandle methodHandle, MethodType methodType) {
+    PainlessConstructor(Constructor<?> javaConstructor, List<Class<?>> typeParameters, MethodHandle methodHandle, MethodType methodType,
+            Map<Class<?>, Object> annotations) {
         this.javaConstructor = javaConstructor;
         this.typeParameters = typeParameters;
         this.methodHandle = methodHandle;
         this.methodType = methodType;
+        this.annotations = annotations;
     }
 
     @Override
@@ -53,11 +57,12 @@ public boolean equals(Object object) {
 
         return Objects.equals(javaConstructor, that.javaConstructor) &&
                 Objects.equals(typeParameters, that.typeParameters) &&
-                Objects.equals(methodType, that.methodType);
+                Objects.equals(methodType, that.methodType) &&
+                Objects.equals(annotations, that.annotations);
     }
 
     @Override
     public int hashCode() {
-        return Objects.hash(javaConstructor, typeParameters, methodType);
+        return Objects.hash(javaConstructor, typeParameters, methodType, annotations);
     }
 }