Use cached SourceLookup for FieldScript (#17927)

Signed-off-by: Jay Deng <[email protected]>

Author: jed326

server/src/main/java/org/opensearch/search/lookup/SearchLookup.java

@@ -42,6 +42,7 @@
 import java.util.LinkedHashSet;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
 
@@ -77,10 +78,10 @@ public class SearchLookup {
      */
     private final Set<String> fieldChain;
     private final DocLookup docMap;
-    private final SourceLookup sourceLookup;
     private final FieldsLookup fieldsLookup;
     private final BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup;
     private final int shardId;
+    private final ConcurrentHashMap<Long, SourceLookup> sourceLookupMap = new ConcurrentHashMap<>();
 
     /**
      * Constructor for backwards compatibility. Use the one with explicit shardId argument.
@@ -101,14 +102,22 @@ public SearchLookup(
         MapperService mapperService,
         BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup,
         int shardId
+    ) {
+        this(mapperService, fieldDataLookup, shardId, new FieldsLookup(mapperService));
+    }
+
+    public SearchLookup(
+        MapperService mapperService,
+        BiFunction<MappedFieldType, Supplier<SearchLookup>, IndexFieldData<?>> fieldDataLookup,
+        int shardId,
+        FieldsLookup fieldsLookup
     ) {
         this.fieldChain = Collections.emptySet();
         docMap = new DocLookup(
             mapperService,
             fieldType -> fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()))
         );
-        sourceLookup = new SourceLookup();
-        fieldsLookup = new FieldsLookup(mapperService);
+        this.fieldsLookup = fieldsLookup;
         this.fieldDataLookup = fieldDataLookup;
         this.shardId = shardId;
     }
@@ -126,7 +135,6 @@ private SearchLookup(SearchLookup searchLookup, Set<String> fieldChain) {
             searchLookup.docMap.mapperService(),
             fieldType -> searchLookup.fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()))
         );
-        this.sourceLookup = searchLookup.sourceLookup;
         this.fieldsLookup = searchLookup.fieldsLookup;
         this.fieldDataLookup = searchLookup.fieldDataLookup;
         this.shardId = searchLookup.shardId;
@@ -160,7 +168,7 @@ public LeafSearchLookup getLeafSearchLookup(LeafReaderContext context) {
         return new LeafSearchLookup(
             context,
             docMap.getLeafDocLookup(context),
-            new SourceLookup(),
+            sourceLookupMap.computeIfAbsent(Thread.currentThread().threadId(), K -> new SourceLookup()),
             fieldsLookup.getLeafFieldsLookup(context)
         );
     }
@@ -173,7 +181,7 @@ public DocLookup doc() {
      * Returned SourceLookup will be unrelated to any created LeafSearchLookups. Instead, use {@link LeafSearchLookup#source()} to access the related {@link SearchLookup}.
      */
     public SourceLookup source() {
-        return sourceLookup;
+        return sourceLookupMap.computeIfAbsent(Thread.currentThread().threadId(), K -> new SourceLookup());
     }
 
     public int shardId() {

server/src/test/java/org/opensearch/script/ScriptServiceTests.java

@@ -46,6 +46,8 @@
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.env.Environment;
+import org.opensearch.search.lookup.FieldsLookup;
+import org.opensearch.search.lookup.SearchLookup;
 import org.opensearch.test.OpenSearchTestCase;
 import org.junit.Before;
 
@@ -63,10 +65,12 @@
 import static org.opensearch.script.ScriptService.SCRIPT_GENERAL_CACHE_SIZE_SETTING;
 import static org.opensearch.script.ScriptService.SCRIPT_GENERAL_MAX_COMPILATIONS_RATE_SETTING;
 import static org.opensearch.script.ScriptService.SCRIPT_MAX_COMPILATIONS_RATE_SETTING;
+import static org.opensearch.search.lookup.SearchLookup.UNKNOWN_SHARD_ID;
 import static org.hamcrest.CoreMatchers.containsString;
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.hamcrest.Matchers.sameInstance;
+import static org.mockito.Mockito.mock;
 
 public class ScriptServiceTests extends OpenSearchTestCase {
 
@@ -180,6 +184,22 @@ public void testInlineScriptCompiledOnceCache() throws IOException {
         assertThat(factoryScript1, sameInstance(factoryScript2));
     }
 
+    public void testScriptsUseCachedSourceLookup() throws IOException {
+        buildScriptService(Settings.EMPTY);
+        Script script = new Script(ScriptType.INLINE, "test", "1+1", Collections.emptyMap());
+        FieldScript.Factory factoryScript = scriptService.compile(script, FieldScript.CONTEXT);
+
+        FieldScript.LeafFactory leafFactory = factoryScript.newFactory(
+            new HashMap<>(),
+            new SearchLookup(null, null, UNKNOWN_SHARD_ID, mock(FieldsLookup.class))
+        );
+
+        FieldScript script1 = leafFactory.newInstance(null);
+        FieldScript script2 = leafFactory.newInstance(null);
+
+        assertThat(script1.getLeafLookup().source(), sameInstance(script2.getLeafLookup().source()));
+    }
+
     public void testAllowAllScriptTypeSettings() throws IOException {
         buildScriptService(Settings.EMPTY);