Adjust ExpressionAggregationScript to support inter-segment concurren…

…cy (elastic#99667)

Handling of _value in a script agg does not support search concurrency
when using the expression script engine. The reason is that the value gets set
globally assuming sequential execution. This commit addresses that by setting
the value to the values source associated with the correct leaf reader context,
while it was previosly being set on a shared data structure.

Closes elastic#99156

modules/lang-expression/src/internalClusterTest/java/org/elasticsearch/script/expression/MoreExpressionIT.java

@@ -502,7 +502,6 @@ public void testSpecialValueVariable() throws Exception {
         assertThat(stats.getAvg(), equalTo(3.0));
     }
 
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/99156")
     public void testStringSpecialValueVariable() throws Exception {
         // i.e. expression script for term aggregations, which is not allowed
         assertAcked(indicesAdmin().prepareCreate("test").setMapping("text", "type=keyword").get());

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ExpressionAggregationScript.java

@@ -83,7 +83,7 @@ public void setNextAggregationValue(Object value) {
                 // _value isn't used in script if specialValue == null
                 if (specialValue != null) {
                     if (value instanceof Number) {
-                        specialValue.setValue(((Number) value).doubleValue());
+                        specialValue.setValue(leaf, ((Number) value).doubleValue());
                     } else {
                         throw new GeneralScriptException("Cannot use expression with text variable using " + exprScript);
                     }

modules/lang-expression/src/main/java/org/elasticsearch/script/expression/ReplaceableConstDoubleValueSource.java

@@ -15,20 +15,21 @@
 import org.apache.lucene.search.IndexSearcher;
 
 import java.io.IOException;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * A {@link DoubleValuesSource} which has a stub {@link DoubleValues} that holds a dynamically replaceable constant double.
  */
 final class ReplaceableConstDoubleValueSource extends DoubleValuesSource {
-    final ReplaceableConstDoubleValues fv;
 
-    ReplaceableConstDoubleValueSource() {
-        fv = new ReplaceableConstDoubleValues();
-    }
+    private final Map<LeafReaderContext, ReplaceableConstDoubleValues> specialValues = new ConcurrentHashMap<>();
 
     @Override
     public DoubleValues getValues(LeafReaderContext ctx, DoubleValues scores) throws IOException {
-        return fv;
+        ReplaceableConstDoubleValues replaceableConstDoubleValues = new ReplaceableConstDoubleValues();
+        specialValues.put(ctx, replaceableConstDoubleValues);
+        return replaceableConstDoubleValues;
     }
 
     @Override
@@ -38,8 +39,12 @@ public boolean needsScores() {
 
     @Override
     public Explanation explain(LeafReaderContext ctx, int docId, Explanation scoreExplanation) throws IOException {
-        if (fv.advanceExact(docId)) return Explanation.match((float) fv.doubleValue(), "ReplaceableConstDoubleValues");
-        else return Explanation.noMatch("ReplaceableConstDoubleValues");
+        // TODO where is this explain called? I bet it's never tested, and probably never called.
+        ReplaceableConstDoubleValues fv = specialValues.get(ctx);
+        if (fv.advanceExact(docId)) {
+            return Explanation.match((float) fv.doubleValue(), "ReplaceableConstDoubleValues");
+        }
+        return Explanation.noMatch("ReplaceableConstDoubleValues");
     }
 
     @Override
@@ -52,7 +57,9 @@ public int hashCode() {
         return System.identityHashCode(this);
     }
 
-    public void setValue(double v) {
+    public void setValue(LeafReaderContext ctx, double v) {
+        ReplaceableConstDoubleValues fv = specialValues.get(ctx);
+        assert fv != null : "getValues must be called before setValue for any given leaf reader context";
         fv.setValue(v);
     }