Reduce overhead of (source-less) SearchHit (#105659)

Reduce overhead of `SearchHit` a little. No need for any real ref-counting if there's neither source nor nested hits. Same goes for `SearchHits` which don't have to be ref-counted if their contents aren't. Also, don't create pointless unmodifiable maps wrapping the empty singleton for highlight fields and use the singleton for the empty search sort values.
elastic · Mar 6, 2024 · c21b23c · c21b23c
1 parent 61a5033
commit c21b23c
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 14 deletions.
diff --git a/server/src/main/java/org/elasticsearch/search/SearchHit.java b/server/src/main/java/org/elasticsearch/search/SearchHit.java
@@ -204,7 +204,7 @@ public SearchHit(
         this.innerHits = innerHits;
         this.documentFields = documentFields;
         this.metaFields = metaFields;
-        this.refCounted = refCounted == null ? LeakTracker.wrap(new SimpleRefCounted()) : ALWAYS_REFERENCED;
+        this.refCounted = refCounted == null ? LeakTracker.wrap(new SimpleRefCounted()) : refCounted;
     }
 
     public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOException {
@@ -233,8 +233,10 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept
         }
         final Map<String, DocumentField> documentFields = in.readMap(DocumentField::new);
         final Map<String, DocumentField> metaFields = in.readMap(DocumentField::new);
-        final Map<String, HighlightField> highlightFields = in.readMapValues(HighlightField::new, HighlightField::name);
-        final SearchSortValues sortValues = new SearchSortValues(in);
+        Map<String, HighlightField> highlightFields = in.readMapValues(HighlightField::new, HighlightField::name);
+        highlightFields = highlightFields.isEmpty() ? null : unmodifiableMap(highlightFields);
+
+        final SearchSortValues sortValues = SearchSortValues.readFrom(in);
 
         final Map<String, Float> matchedQueries;
         if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_8_0)) {
@@ -257,12 +259,17 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept
             index = shardTarget.getIndex();
             clusterAlias = shardTarget.getClusterAlias();
         }
+
+        boolean isPooled = pooled && source != null;
         final Map<String, SearchHits> innerHits;
         int size = in.readVInt();
         if (size > 0) {
             innerHits = Maps.newMapWithExpectedSize(size);
             for (int i = 0; i < size; i++) {
-                innerHits.put(in.readString(), SearchHits.readFrom(in, pooled));
+                var key = in.readString();
+                var nestedHits = SearchHits.readFrom(in, pooled);
+                innerHits.put(key, nestedHits);
+                isPooled = isPooled || nestedHits.isPooled();
             }
         } else {
             innerHits = null;
@@ -277,7 +284,7 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept
             seqNo,
             primaryTerm,
             source,
-            unmodifiableMap(highlightFields),
+            highlightFields,
             sortValues,
             matchedQueries,
             explanation,
@@ -288,7 +295,7 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept
             innerHits,
             documentFields,
             metaFields,
-            pooled ? null : ALWAYS_REFERENCED
+            isPooled ? null : ALWAYS_REFERENCED
         );
     }
 

diff --git a/server/src/main/java/org/elasticsearch/search/SearchHits.java b/server/src/main/java/org/elasticsearch/search/SearchHits.java
@@ -132,24 +132,31 @@ public static SearchHits readFrom(StreamInput in, boolean pooled) throws IOExcep
         final float maxScore = in.readFloat();
         int size = in.readVInt();
         final SearchHit[] hits;
+        boolean isPooled = false;
         if (size == 0) {
             hits = EMPTY;
         } else {
             hits = new SearchHit[size];
             for (int i = 0; i < hits.length; i++) {
-                hits[i] = SearchHit.readFrom(in, pooled);
+                var hit = SearchHit.readFrom(in, pooled);
+                hits[i] = hit;
+                isPooled = isPooled || hit.isPooled();
             }
         }
         var sortFields = in.readOptionalArray(Lucene::readSortField, SortField[]::new);
         var collapseField = in.readOptionalString();
         var collapseValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new);
-        if (pooled) {
+        if (isPooled) {
             return new SearchHits(hits, totalHits, maxScore, sortFields, collapseField, collapseValues);
         } else {
             return unpooled(hits, totalHits, maxScore, sortFields, collapseField, collapseValues);
         }
     }
 
+    public boolean isPooled() {
+        return refCounted != ALWAYS_REFERENCED;
+    }
+
     @Override
     public void writeTo(StreamOutput out) throws IOException {
         assert hasReferences();

diff --git a/server/src/main/java/org/elasticsearch/search/SearchSortValues.java b/server/src/main/java/org/elasticsearch/search/SearchSortValues.java
@@ -32,8 +32,7 @@ public class SearchSortValues implements ToXContentFragment, Writeable {
     private final Object[] rawSortValues;
 
     SearchSortValues(Object[] sortValues) {
-        this.formattedSortValues = Objects.requireNonNull(sortValues, "sort values must not be empty");
-        this.rawSortValues = EMPTY_ARRAY;
+        this(Objects.requireNonNull(sortValues, "sort values must not be empty"), EMPTY_ARRAY);
     }
 
     public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormats) {
@@ -52,9 +51,18 @@ public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormat
         }
     }
 
-    SearchSortValues(StreamInput in) throws IOException {
-        this.formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
-        this.rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
+    public static SearchSortValues readFrom(StreamInput in) throws IOException {
+        Object[] formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
+        Object[] rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new);
+        if (formattedSortValues.length == 0 && rawSortValues.length == 0) {
+            return EMPTY;
+        }
+        return new SearchSortValues(formattedSortValues, rawSortValues);
+    }
+
+    private SearchSortValues(Object[] formattedSortValues, Object[] rawSortValues) {
+        this.formattedSortValues = formattedSortValues;
+        this.rawSortValues = rawSortValues;
     }
 
     @Override

diff --git a/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java b/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java
@@ -90,7 +90,7 @@ protected SearchSortValues createTestInstance() {
 
     @Override
     protected Writeable.Reader<SearchSortValues> instanceReader() {
-        return SearchSortValues::new;
+        return SearchSortValues::readFrom;
     }
 
     @Override