Adjust ExpressionAggregationScript to support inter-segment concurrency (#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 #99156

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

@@ -502,7 +502,6 @@ public class MoreExpressionIT extends ESIntegTestCase {
         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 @@ class ExpressionAggregationScript implements AggregationScript.LeafFactory {
                 // _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.Explanation;
 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() {
+    private final Map<LeafReaderContext, ReplaceableConstDoubleValues> specialValues = new ConcurrentHashMap<>();
-        fv = new ReplaceableConstDoubleValues();
-    }
 
     @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 @@ final class ReplaceableConstDoubleValueSource extends DoubleValuesSource {
 
     @Override
     public Explanation explain(LeafReaderContext ctx, int docId, Explanation scoreExplanation) throws IOException {
-        if (fv.advanceExact(docId)) return Explanation.match((float) fv.doubleValue(), "ReplaceableConstDoubleValues");
+        // TODO where is this explain called? I bet it's never tested, and probably never called.
-        else return Explanation.noMatch("ReplaceableConstDoubleValues");
+        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 @@ final class ReplaceableConstDoubleValueSource extends DoubleValuesSource {
         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);
     }