From e317430a74c57c7bce13f5c470e398ac4ef1767c Mon Sep 17 00:00:00 2001
From: Armin Braun <me@obrown.io>
Date: Tue, 16 Jul 2024 18:31:01 +0200
Subject: [PATCH] Fix leak in collapsing search results (#110927)

Fixing this case for now by enforcing unpooled to plug the leak, this
needs a little more work to function well pooled.
---
 docs/changelog/110927.yaml                    |  5 ++++
 .../search/CollapseSearchResultsIT.java       | 23 +++++++++++++++++++
 .../action/search/ExpandSearchPhase.java      |  4 ++++
 3 files changed, 32 insertions(+)
 create mode 100644 docs/changelog/110927.yaml

diff --git a/docs/changelog/110927.yaml b/docs/changelog/110927.yaml
new file mode 100644
index 0000000000000..3602ce3e811fa
--- /dev/null
+++ b/docs/changelog/110927.yaml
@@ -0,0 +1,5 @@
+pr: 110927
+summary: Fix leak in collapsing search results
+area: Search
+type: bug
+issues: []
diff --git a/server/src/internalClusterTest/java/org/elasticsearch/search/CollapseSearchResultsIT.java b/server/src/internalClusterTest/java/org/elasticsearch/search/CollapseSearchResultsIT.java
index f5fdd752a6f57..aa721122c2160 100644
--- a/server/src/internalClusterTest/java/org/elasticsearch/search/CollapseSearchResultsIT.java
+++ b/server/src/internalClusterTest/java/org/elasticsearch/search/CollapseSearchResultsIT.java
@@ -61,4 +61,27 @@ public void testCollapseWithDocValueFields() {
             }
         );
     }
+
+    public void testCollapseWithFields() {
+        final String indexName = "test_collapse";
+        createIndex(indexName);
+        final String collapseField = "collapse_field";
+        final String otherField = "other_field";
+        assertAcked(indicesAdmin().preparePutMapping(indexName).setSource(collapseField, "type=keyword", otherField, "type=keyword"));
+        index(indexName, "id_1_0", Map.of(collapseField, "value1", otherField, "other_value1"));
+        index(indexName, "id_1_1", Map.of(collapseField, "value1", otherField, "other_value2"));
+        index(indexName, "id_2_0", Map.of(collapseField, "value2", otherField, "other_value3"));
+        refresh(indexName);
+
+        assertNoFailuresAndResponse(
+            prepareSearch(indexName).setQuery(new MatchAllQueryBuilder())
+                .setFetchSource(false)
+                .addFetchField(otherField)
+                .setCollapse(new CollapseBuilder(collapseField).setInnerHits(new InnerHitBuilder("ih").setSize(2))),
+            searchResponse -> {
+                assertEquals(collapseField, searchResponse.getHits().getCollapseField());
+                assertEquals(Set.of(new BytesRef("value1"), new BytesRef("value2")), Set.of(searchResponse.getHits().getCollapseValues()));
+            }
+        );
+    }
 }
diff --git a/server/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java b/server/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java
index e8470ba77632f..e2385745149c1 100644
--- a/server/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java
+++ b/server/src/main/java/org/elasticsearch/action/search/ExpandSearchPhase.java
@@ -100,6 +100,10 @@ private void doRun() {
                     if (hit.getInnerHits() == null) {
                         hit.setInnerHits(Maps.newMapWithExpectedSize(innerHitBuilders.size()));
                     }
+                    if (hit.isPooled() == false) {
+                        // TODO: make this work pooled by forcing the hit itself to become pooled as needed here
+                        innerHits = innerHits.asUnpooled();
+                    }
                     hit.getInnerHits().put(innerHitBuilder.getName(), innerHits);
                     assert innerHits.isPooled() == false || hit.isPooled() : "pooled inner hits can only be added to a pooled hit";
                     innerHits.mustIncRef();