Use cached SourceLookup for FieldScript

Signed-off-by: Jay Deng <jayd0104@gmail.com>

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.
@@ -107,7 +108,6 @@ public SearchLookup(
             mapperService,
             fieldType -> fieldDataLookup.apply(fieldType, () -> forkAndTrackFieldReferences(fieldType.name()))
         );
-        sourceLookup = new SourceLookup();
         fieldsLookup = new FieldsLookup(mapperService);
         this.fieldDataLookup = fieldDataLookup;
         this.shardId = shardId;
@@ -126,7 +126,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 +159,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 +172,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

@@ -31,6 +31,7 @@
 
 package org.opensearch.script;
 
+import org.apache.lucene.index.LeafReaderContext;
 import org.opensearch.ResourceNotFoundException;
 import org.opensearch.action.admin.cluster.storedscripts.GetStoredScriptRequest;
 import org.opensearch.cluster.ClusterName;
@@ -46,6 +47,8 @@
 import org.opensearch.core.common.bytes.BytesReference;
 import org.opensearch.core.xcontent.MediaTypeRegistry;
 import org.opensearch.env.Environment;
+import org.opensearch.index.mapper.MapperService;
+import org.opensearch.search.lookup.SearchLookup;
 import org.opensearch.test.OpenSearchTestCase;
 import org.junit.Before;
 
@@ -67,6 +70,7 @@
 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,25 @@ 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 factoryScript1 = scriptService.compile(script, FieldScript.CONTEXT);
+        FieldScript.LeafFactory leafFactory1 = factoryScript1.newFactory(
+            new HashMap<>(),
+            new SearchLookup(mock(MapperService.class), null)
+        );
+
+        LeafReaderContext ctx = mock(LeafReaderContext.class);
+        FieldScript script1 = leafFactory1.newInstance(ctx);
+        FieldScript script2 = leafFactory1.newInstance(ctx);
+
+        script1.setDocument(1);
+        script2.setDocument(1);
+
+        assertThat(script1.getLeafLookup().source(), sameInstance(script2.getLeafLookup().source()));
+    }
+
     public void testAllowAllScriptTypeSettings() throws IOException {
         buildScriptService(Settings.EMPTY);