Fix sort related ITs for concurrent search

Signed-off-by: Neetika Singhal <[email protected]>
neetikasinghal · Aug 22, 2023 · b68e08d · b68e08d
1 parent 7195a09
commit b68e08d
Show file tree

Hide file tree

Showing 6 changed files with 70 additions and 16 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -81,6 +81,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Rethrow OpenSearch exception for non-concurrent path while using concurrent search ([#9177](https://github.com/opensearch-project/OpenSearch/pull/9177))
 - Improve performance of encoding composite keys in multi-term aggregations ([#9412](https://github.com/opensearch-project/OpenSearch/pull/9412))
 - Refactor Compressors from CompressorFactory to CompressorRegistry for extensibility ([#9262](https://github.com/opensearch-project/OpenSearch/pull/9262))
+- Fix sort related ITs for concurrent search ([#9177](https://github.com/opensearch-project/OpenSearch/pull/9466)
 
 ### Deprecated
 

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java b/server/src/internalClusterTest/java/org/opensearch/search/sort/FieldSortIT.java
@@ -32,6 +32,8 @@
 
 package org.opensearch.search.sort;
 
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
 import org.apache.lucene.tests.util.TestUtil;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.UnicodeUtil;
@@ -45,6 +47,7 @@
 import org.opensearch.cluster.metadata.IndexMetadata;
 import org.opensearch.common.Numbers;
 import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.core.rest.RestStatus;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
@@ -60,7 +63,7 @@
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.SearchHits;
 import org.opensearch.test.InternalSettingsPlugin;
-import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.ParameterizedOpenSearchIntegTestCase;
 import org.hamcrest.Matchers;
 
 import java.io.IOException;
@@ -86,6 +89,7 @@
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
 import static org.opensearch.index.query.functionscore.ScoreFunctionBuilders.fieldValueFactorFunction;
 import static org.opensearch.script.MockScriptPlugin.NAME;
+import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertFirstHit;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
@@ -105,7 +109,24 @@
 import static org.hamcrest.Matchers.nullValue;
 import static org.hamcrest.Matchers.oneOf;
 
-public class FieldSortIT extends OpenSearchIntegTestCase {
+public class FieldSortIT extends ParameterizedOpenSearchIntegTestCase {
+    public FieldSortIT(Settings dynamicSettings) {
+        super(dynamicSettings);
+    }
+
+    @ParametersFactory
+    public static Collection<Object[]> parameters() {
+        return Arrays.asList(
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), false).build() },
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), true).build() }
+        );
+    }
+
+    @Override
+    protected Settings featureFlagSettings() {
+        return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true").build();
+    }
+
     public static class CustomScriptPlugin extends MockScriptPlugin {
         @Override
         protected Map<String, Function<Map<String, Object>, Object>> pluginScripts() {

diff --git a/server/src/internalClusterTest/java/org/opensearch/search/sort/SimpleSortIT.java b/server/src/internalClusterTest/java/org/opensearch/search/sort/SimpleSortIT.java
@@ -32,11 +32,15 @@
 
 package org.opensearch.search.sort;
 
+import com.carrotsearch.randomizedtesting.annotations.ParametersFactory;
+
 import org.opensearch.action.index.IndexRequestBuilder;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.action.search.ShardSearchFailure;
 import org.opensearch.common.geo.GeoPoint;
 import org.opensearch.common.geo.GeoUtils;
+import org.opensearch.common.settings.Settings;
+import org.opensearch.common.util.FeatureFlags;
 import org.opensearch.index.fielddata.ScriptDocValues;
 import org.opensearch.plugins.Plugin;
 import org.opensearch.script.MockScriptPlugin;
@@ -45,7 +49,7 @@
 import org.opensearch.search.SearchHit;
 import org.opensearch.search.sort.ScriptSortBuilder.ScriptSortType;
 import org.opensearch.test.InternalSettingsPlugin;
-import org.opensearch.test.OpenSearchIntegTestCase;
+import org.opensearch.test.ParameterizedOpenSearchIntegTestCase;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -61,6 +65,7 @@
 import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder;
 import static org.opensearch.index.query.QueryBuilders.matchAllQuery;
 import static org.opensearch.index.query.QueryBuilders.termQuery;
+import static org.opensearch.search.SearchService.CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING;
 import static org.opensearch.search.sort.SortBuilders.scriptSort;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked;
 import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertHitCount;
@@ -70,10 +75,27 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.not;
 
-public class SimpleSortIT extends OpenSearchIntegTestCase {
+public class SimpleSortIT extends ParameterizedOpenSearchIntegTestCase {
 
     private static final String DOUBLE_APOSTROPHE = "\u0027\u0027";
 
+    public SimpleSortIT(Settings dynamicSettings) {
+        super(dynamicSettings);
+    }
+
+    @ParametersFactory
+    public static Collection<Object[]> parameters() {
+        return Arrays.asList(
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), false).build() },
+            new Object[] { Settings.builder().put(CLUSTER_CONCURRENT_SEGMENT_SEARCH_SETTING.getKey(), true).build() }
+        );
+    }
+
+    @Override
+    protected Settings featureFlagSettings() {
+        return Settings.builder().put(super.featureFlagSettings()).put(FeatureFlags.CONCURRENT_SEGMENT_SEARCH, "true").build();
+    }
+
     @Override
     protected Collection<Class<? extends Plugin>> nodePlugins() {
         return Arrays.asList(CustomScriptPlugin.class, InternalSettingsPlugin.class);

diff --git a/...in/java/org/opensearch/index/fielddata/fieldcomparator/BytesRefFieldComparatorSource.java b/...in/java/org/opensearch/index/fielddata/fieldcomparator/BytesRefFieldComparatorSource.java
@@ -91,7 +91,7 @@ protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOEx
         return indexFieldData.load(context).getBytesValues();
     }
 
-    protected void setScorer(Scorable scorer) {}
+    protected void setScorer(Scorable scorer, LeafReaderContext context) {}
 
     @Override
     public FieldComparator<?> newComparator(String fieldname, int numHits, boolean enableSkipping, boolean reversed) {
@@ -133,9 +133,11 @@ protected SortedDocValues getSortedDocValues(LeafReaderContext context, String f
         }
 
         return new FieldComparator.TermValComparator(numHits, null, sortMissingLast) {
+            LeafReaderContext leafReaderContext;
 
             @Override
             protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String field) throws IOException {
+                leafReaderContext = context;
                 final SortedBinaryDocValues values = getValues(context);
                 final BinaryDocValues selectedValues;
                 if (nested == null) {
@@ -151,7 +153,7 @@ protected BinaryDocValues getBinaryDocValues(LeafReaderContext context, String f
 
             @Override
             public void setScorer(Scorable scorer) {
-                BytesRefFieldComparatorSource.this.setScorer(scorer);
+                BytesRefFieldComparatorSource.this.setScorer(scorer, leafReaderContext);
             }
 
         };

diff --git a/...ain/java/org/opensearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java b/...ain/java/org/opensearch/index/fielddata/fieldcomparator/DoubleValuesComparatorSource.java
@@ -95,7 +95,7 @@ private NumericDoubleValues getNumericDocValues(LeafReaderContext context, doubl
         }
     }
 
-    protected void setScorer(Scorable scorer) {}
+    protected void setScorer(Scorable scorer, LeafReaderContext context) {}
 
     @Override
     public FieldComparator<?> newComparator(String fieldname, int numHits, boolean enableSkipping, boolean reversed) {
@@ -115,7 +115,7 @@ protected NumericDocValues getNumericDocValues(LeafReaderContext context, String
 
                     @Override
                     public void setScorer(Scorable scorer) {
-                        DoubleValuesComparatorSource.this.setScorer(scorer);
+                        DoubleValuesComparatorSource.this.setScorer(scorer, context);
                     }
                 };
             }

diff --git a/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java b/server/src/main/java/org/opensearch/search/sort/ScriptSortBuilder.java
@@ -70,7 +70,9 @@
 
 import java.io.IOException;
 import java.util.Locale;
+import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 
 import static org.opensearch.core.xcontent.ConstructingObjectParser.constructorArg;
 import static org.opensearch.search.sort.FieldSortBuilder.validateMaxChildrenExistOnlyInTopLevelNestedSort;
@@ -355,11 +357,14 @@ private IndexFieldData.XFieldComparatorSource fieldComparatorSource(QueryShardCo
                 final StringSortScript.Factory factory = context.compile(script, StringSortScript.CONTEXT);
                 final StringSortScript.LeafFactory searchScript = factory.newFactory(script.getParams(), context.lookup());
                 return new BytesRefFieldComparatorSource(null, null, valueMode, nested) {
-                    StringSortScript leafScript;
+                    // introducing a map to keep a mapping between the leaf reader context and leaf script
+                    // such that the functions of the class are thread safe in case of concurrent search
+                    final Map<LeafReaderContext, StringSortScript> leafContextSortScriptMap = new ConcurrentHashMap<>();
 
                     @Override
                     protected SortedBinaryDocValues getValues(LeafReaderContext context) throws IOException {
-                        leafScript = searchScript.newInstance(context);
+                        final StringSortScript leafScript = searchScript.newInstance(context);
+                        leafContextSortScriptMap.put(context, leafScript);
                         final BinaryDocValues values = new AbstractBinaryDocValues() {
                             final BytesRefBuilder spare = new BytesRefBuilder();
 
@@ -379,8 +384,8 @@ public BytesRef binaryValue() {
                     }
 
                     @Override
-                    protected void setScorer(Scorable scorer) {
-                        leafScript.setScorer(scorer);
+                    protected void setScorer(Scorable scorer, LeafReaderContext context) {
+                        leafContextSortScriptMap.get(context).setScorer(scorer);
                     }
 
                     @Override
@@ -403,11 +408,14 @@ public BucketedSort newBucketedSort(
                 final NumberSortScript.Factory numberSortFactory = context.compile(script, NumberSortScript.CONTEXT);
                 final NumberSortScript.LeafFactory numberSortScript = numberSortFactory.newFactory(script.getParams(), context.lookup());
                 return new DoubleValuesComparatorSource(null, Double.MAX_VALUE, valueMode, nested) {
-                    NumberSortScript leafScript;
+                    // introducing a map to keep a mapping between the leaf reader context and leaf script
+                    // such that the functions of the class are thread safe in case of concurrent search
+                    final Map<LeafReaderContext, NumberSortScript> leafContextSortScriptMap = new ConcurrentHashMap<>();
 
                     @Override
                     protected SortedNumericDoubleValues getValues(LeafReaderContext context) throws IOException {
-                        leafScript = numberSortScript.newInstance(context);
+                        final NumberSortScript leafScript = numberSortScript.newInstance(context);
+                        leafContextSortScriptMap.put(context, leafScript);
                         final NumericDoubleValues values = new NumericDoubleValues() {
                             @Override
                             public boolean advanceExact(int doc) throws IOException {
@@ -424,8 +432,8 @@ public double doubleValue() {
                     }
 
                     @Override
-                    protected void setScorer(Scorable scorer) {
-                        leafScript.setScorer(scorer);
+                    protected void setScorer(Scorable scorer, LeafReaderContext context) {
+                        leafContextSortScriptMap.get(context).setScorer(scorer);
                     }
                 };
             default: