Scripting: Add StatefulFactoryType as optional intermediate factory i…

…n script contexts (#24974)

ScriptContexts currently understand a FactoryType that can produce
instances of the script InstanceType. However, for search scripts, this
does not work as we have the concept of LeafSearchScript that is created
per lucene segment. This commit effectively renames the existing
SearchScript class into SearchScript.LeafFactory, which is a new,
optional, class that can be defined within a ScriptContext.
LeafSearchScript is effectively renamed back into SearchScript. This
change allows the model of stateless factory -> stateful factory ->
script instance to continue, but in a generic way that any script
context may take advantage of.

relates #20426

core/src/main/java/org/elasticsearch/common/lucene/search/function/ScriptScoreFunction.java

@@ -24,7 +24,6 @@
 import org.apache.lucene.search.Explanation;
 import org.apache.lucene.search.Scorer;
 import org.elasticsearch.script.ExplainableSearchScript;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
 import org.elasticsearch.script.GeneralScriptException;
 import org.elasticsearch.script.SearchScript;
@@ -65,18 +64,18 @@ public DocIdSetIterator iterator() {
 
     private final Script sScript;
 
-    private final SearchScript script;
+    private final SearchScript.LeafFactory script;
 
 
-    public ScriptScoreFunction(Script sScript, SearchScript script) {
+    public ScriptScoreFunction(Script sScript, SearchScript.LeafFactory script) {
         super(CombineFunction.REPLACE);
         this.sScript = sScript;
         this.script = script;
     }
 
     @Override
     public LeafScoreFunction getLeafScoreFunction(LeafReaderContext ctx) throws IOException {
-        final LeafSearchScript leafScript = script.getLeafSearchScript(ctx);
+        final SearchScript leafScript = script.newInstance(ctx);
         final CannedScorer scorer = new CannedScorer();
         leafScript.setScorer(scorer);
         return new LeafScoreFunction() {

core/src/main/java/org/elasticsearch/index/query/InnerHitContextBuilder.java

@@ -74,7 +74,7 @@ protected void setupInnerHitsContext(QueryShardContext queryShardContext,
         }
         if (innerHitBuilder.getScriptFields() != null) {
             for (SearchSourceBuilder.ScriptField field : innerHitBuilder.getScriptFields()) {
-                SearchScript searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
+                SearchScript.LeafFactory searchScript = innerHitsContext.getQueryShardContext().getSearchScript(field.script(),
                     SearchScript.CONTEXT);
                 innerHitsContext.scriptFields().add(new org.elasticsearch.search.fetch.subphase.ScriptFieldsContext.ScriptField(
                     field.fieldName(), searchScript, field.ignoreFailure()));

core/src/main/java/org/elasticsearch/index/query/QueryShardContext.java

@@ -329,21 +329,21 @@ public final Index index() {
      * Compiles (or retrieves from cache) and binds the parameters to the
      * provided script
      */
-    public final SearchScript getSearchScript(Script script, ScriptContext<SearchScript.Factory> context) {
+    public final SearchScript.LeafFactory getSearchScript(Script script, ScriptContext<SearchScript.Factory> context) {
         failIfFrozen();
         SearchScript.Factory factory = scriptService.compile(script, context);
-        return factory.newInstance(script.getParams(), lookup());
+        return factory.newFactory(script.getParams(), lookup());
     }
     /**
      * Returns a lazily created {@link SearchScript} that is compiled immediately but can be pulled later once all
      * parameters are available.
      */
-    public final Function<Map<String, Object>, SearchScript> getLazySearchScript(
+    public final Function<Map<String, Object>, SearchScript.LeafFactory> getLazySearchScript(
         Script script, ScriptContext<SearchScript.Factory> context) {
         // TODO: this "lazy" binding can be removed once scripted metric aggs have their own contexts, which take _agg/_aggs as a parameter
         failIfFrozen();
         SearchScript.Factory factory = scriptService.compile(script, context);
-        return (p) -> factory.newInstance(p, lookup());
+        return (p) -> factory.newFactory(p, lookup());
     }
 
     /**

core/src/main/java/org/elasticsearch/index/query/ScriptQueryBuilder.java

@@ -34,9 +34,7 @@
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.xcontent.XContentBuilder;
 import org.elasticsearch.common.xcontent.XContentParser;
-import org.elasticsearch.script.LeafSearchScript;
 import org.elasticsearch.script.Script;
-import org.elasticsearch.script.ScriptContext;
 import org.elasticsearch.script.SearchScript;
 
 import java.io.IOException;
@@ -137,9 +135,9 @@ protected Query doToQuery(QueryShardContext context) throws IOException {
     static class ScriptQuery extends Query {
 
         final Script script;
-        final SearchScript searchScript;
+        final SearchScript.LeafFactory searchScript;
 
-        ScriptQuery(Script script, SearchScript searchScript) {
+        ScriptQuery(Script script, SearchScript.LeafFactory searchScript) {
             this.script = script;
             this.searchScript = searchScript;
         }
@@ -181,7 +179,7 @@ public Weight createWeight(IndexSearcher searcher, boolean needsScores, float bo
                 @Override
                 public Scorer scorer(LeafReaderContext context) throws IOException {
                     DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc());
-                    final LeafSearchScript leafScript = searchScript.getLeafSearchScript(context);
+                    final SearchScript leafScript = searchScript.newInstance(context);
                     TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) {
 
                         @Override

core/src/main/java/org/elasticsearch/index/query/functionscore/ScriptScoreFunctionBuilder.java

@@ -94,7 +94,7 @@ protected int doHashCode() {
     @Override
     protected ScoreFunction doToFunction(QueryShardContext context) {
         try {
-            SearchScript searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
+            SearchScript.LeafFactory searchScript = context.getSearchScript(script, SearchScript.CONTEXT);
             return new ScriptScoreFunction(script, searchScript);
         } catch (Exception e) {
             throw new QueryShardException(context, "script_score: the script could not be loaded", e);

core/src/main/java/org/elasticsearch/script/ExplainableSearchScript.java

@@ -47,7 +47,7 @@
  * This is currently not used inside elasticsearch but it is used, see for example here:
  * https://github.com/elastic/elasticsearch/issues/8561
  */
-public interface ExplainableSearchScript extends LeafSearchScript {
+public interface ExplainableSearchScript {
 
     /**
      * Build the explanation of the current document being scored

core/src/main/java/org/elasticsearch/script/ScriptContext.java

@@ -27,48 +27,76 @@
  * A {@link ScriptContext} contains the information related to a single use case and the interfaces
  * and methods necessary for a {@link ScriptEngine} to implement.
  * <p>
- * There are two related classes which must be supplied to construct a {@link ScriptContext}.
+ * There are at least two (and optionally a third) related classes which must be defined.
  * <p>
- * The <i>FactoryType</i> is a factory class for constructing instances of a script. The
- * {@link ScriptService} returns an instance of <i>FactoryType</i> when compiling a script. This class
- * must be stateless so it is cacheable by the {@link ScriptService}. It must have an abstract method
- * named {@code newInstance} which {@link ScriptEngine} implementations will define.
- * <p>
- * The <i>InstanceType</i> is a class returned by the {@code newInstance} method of the
- * <i>FactoryType</i>. It is an instance of a script and may be stateful. Instances of
+ * The <i>InstanceType</i> is a class which users of the script api call to execute a script. It
+ * may be stateful. Instances of
  * the <i>InstanceType</i> may be executed multiple times by a caller with different arguments. This
  * class must have an abstract method named {@code execute} which {@link ScriptEngine} implementations
  * will define.
+ * <p>
+ * The <i>FactoryType</i> is a factory class returned by the {@link ScriptService} when compiling
+ * a script. This class must be stateless so it is cacheable by the {@link ScriptService}. It must
+ * have one of the following:
+ * <ul>
+ *     <li>An abstract method named {@code newInstance} which returns an instance of <i>InstanceType</i></li>
+ *     <li>An abstract method named {@code newFactory} which returns an instance of <i>StatefulFactoryType</i></li>
+ * </ul>
+ * <p>
+ * The <i>StatefulFactoryType</i> is an optional class which allows a stateful factory from the
+ * stateless factory type required by the {@link ScriptService}. If defined, the <i>StatefulFactoryType</i>
+ * must have a method named {@code newInstance} which returns an instance of <i>InstanceType</i>.
  */
 public final class ScriptContext<FactoryType> {
 
     /** A unique identifier for this context. */
     public final String name;
 
-    /** A factory class for constructing instances of a script. */
+    /** A factory class for constructing script or stateful factory instances. */
     public final Class<FactoryType> factoryClazz;
 
+    /** A factory class for construct script instances. */
+    public final Class<?> statefulFactoryClazz;
+
     /** A class that is an instance of a script. */
     public final Class<?> instanceClazz;
 
     /** Construct a context with the related instance and compiled classes. */
     public ScriptContext(String name, Class<FactoryType> factoryClazz) {
         this.name = name;
         this.factoryClazz = factoryClazz;
-        Method newInstanceMethod = null;
-        for (Method method : factoryClazz.getMethods()) {
-            if (method.getName().equals("newInstance")) {
-                if (newInstanceMethod != null) {
-                    throw new IllegalArgumentException("Cannot have multiple newInstance methods on FactoryType class ["
-                        + factoryClazz.getName() + "] for script context [" + name + "]");
-                }
-                newInstanceMethod = method;
+        Method newInstanceMethod = findMethod("FactoryType", factoryClazz, "newInstance");
+        Method newFactoryMethod = findMethod("FactoryType", factoryClazz, "newFactory");
+        if (newFactoryMethod != null) {
+            assert newInstanceMethod == null;
+            statefulFactoryClazz = newFactoryMethod.getReturnType();
+            newInstanceMethod = findMethod("StatefulFactoryType", statefulFactoryClazz, "newInstance");
+            if (newInstanceMethod == null) {
+                throw new IllegalArgumentException("Could not find method newInstance StatefulFactoryType class ["
+                    + statefulFactoryClazz.getName() + "] for script context [" + name + "]");
             }
-        }
-        if (newInstanceMethod == null) {
-            throw new IllegalArgumentException("Could not find method newInstance on FactoryType class ["
+        } else if (newInstanceMethod != null) {
+            assert newFactoryMethod == null;
+            statefulFactoryClazz = null;
+        } else {
+            throw new IllegalArgumentException("Could not find method newInstance or method newFactory on FactoryType class ["
                 + factoryClazz.getName() + "] for script context [" + name + "]");
         }
         instanceClazz = newInstanceMethod.getReturnType();
     }
+
+    /** Returns a method with the given name, or throws an exception if multiple are found. */
+    private Method findMethod(String type, Class<?> clazz, String methodName) {
+        Method foundMethod = null;
+        for (Method method : clazz.getMethods()) {
+            if (method.getName().equals(methodName)) {
+                if (foundMethod != null) {
+                    throw new IllegalArgumentException("Cannot have multiple " + methodName + " methods on " + type + " class ["
+                        + clazz.getName() + "] for script context [" + name + "]");
+                }
+                foundMethod = method;
+            }
+        }
+        return foundMethod;
+    }
 }