Fix aggregation memory leak for CCS

server/src/internalClusterTest/java/org/elasticsearch/search/ccs/CrossClusterSearchLeakIT.java

@@ -0,0 +1,120 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.search.ccs;
+
+import org.elasticsearch.action.ActionFuture;
+import org.elasticsearch.action.search.ClearScrollRequest;
+import org.elasticsearch.action.search.SearchRequest;
+import org.elasticsearch.action.search.SearchResponse;
+import org.elasticsearch.client.Client;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.node.DiscoveryNode;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.query.MatchAllQueryBuilder;
+import org.elasticsearch.search.builder.SearchSourceBuilder;
+import org.elasticsearch.test.AbstractMultiClustersTestCase;
+import org.elasticsearch.test.InternalTestCluster;
+import org.hamcrest.Matchers;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.StreamSupport;
+
+import static org.elasticsearch.search.aggregations.AggregationBuilders.terms;
+import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
+
+public class CrossClusterSearchLeakIT extends AbstractMultiClustersTestCase {
+    @Override
+    protected Collection<String> remoteClusterAlias() {
+        return List.of("cluster_a");
+    }
+
+    @Override
+    protected boolean reuseClusters() {
+        return false;
+    }
+
+    private int indexDocs(Client client, String index) {
+        int numDocs = between(1, 200);
+        for (int i = 0; i < numDocs; i++) {
+            client.prepareIndex(index).setSource("f", "v" + i).get();
+        }
+        client.admin().indices().prepareRefresh(index).get();
+        return numDocs;
+    }
+
+    /**
+     * This test mainly validates that we do not leak any memory when running CCS in various modes
+     * <ul>
+     *     <li>proxy vs non-proxy</li>
+     *     <li>single-phase query-fetch or multi-phase</li>
+     *     <li>minimize roundtrip vs not</li>
+     *     <li>scroll vs no scroll</li>
+     * </ul>
+     */
+    public void testSearch() throws Exception {
+        assertAcked(client(LOCAL_CLUSTER).admin().indices().prepareCreate("demo")
+            .setMapping("f", "type=keyword")
+            .setSettings(Settings.builder().put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, between(1, 3))));
+        indexDocs(client(LOCAL_CLUSTER), "demo");
+        final InternalTestCluster remoteCluster = cluster("cluster_a");
+        remoteCluster.ensureAtLeastNumDataNodes(3);
+        List<String> remoteDataNodes = StreamSupport.stream(remoteCluster.clusterService().state().nodes().spliterator(), false)
+            .filter(DiscoveryNode::canContainData)
+            .map(DiscoveryNode::getName)
+            .collect(Collectors.toList());
+        assertThat(remoteDataNodes.size(), Matchers.greaterThanOrEqualTo(3));
+        List<String> seedNodes = randomSubsetOf(between(1, remoteDataNodes.size() - 1), remoteDataNodes);
+        disconnectFromRemoteClusters();
+        configureRemoteCluster("cluster_a", seedNodes);
+        final Settings.Builder allocationFilter = Settings.builder();
+        if (randomBoolean()) {
+            // Using proxy connections
+            allocationFilter.put("index.routing.allocation.exclude._name", String.join(",", seedNodes));
+        } else {
+            allocationFilter.put("index.routing.allocation.include._name", String.join(",", seedNodes));
+        }
+        assertAcked(client("cluster_a").admin().indices().prepareCreate("prod")
+            .setMapping("f", "type=keyword")
+            .setSettings(Settings.builder().put(allocationFilter.build()).put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0).put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, between(1, 3))));
+        assertFalse(client("cluster_a").admin().cluster().prepareHealth("prod")
+            .setWaitForYellowStatus().setTimeout(TimeValue.timeValueSeconds(10)).get().isTimedOut());
+        indexDocs(client("cluster_a"), "prod");
+
+        String[] indices = randomBoolean() ? new String[] { "demo", "cluster_a:prod" } : new String[] { "cluster_a:prod" };
+        final SearchRequest templateSearch = new SearchRequest(indices);
+        templateSearch.allowPartialSearchResults(false);
+        templateSearch.setCcsMinimizeRoundtrips(randomBoolean());
+        templateSearch.source(new SearchSourceBuilder().query(new MatchAllQueryBuilder()).aggregation(terms("f").field("f")).size(1000));
+
+        List<ActionFuture<SearchResponse>> futures = new ArrayList<>();
+        for (int i = 0; i < 10; ++i) {
+
+            final SearchRequest searchRequest = new SearchRequest(templateSearch);
+            if (randomBoolean()) {
+                searchRequest.scroll("30s");
+            }
+            searchRequest.setCcsMinimizeRoundtrips(randomBoolean());
+            futures.add(client(LOCAL_CLUSTER).search(searchRequest));
+        }
+
+        for (ActionFuture<SearchResponse> future : futures) {
+            SearchResponse searchResponse = future.get();
+            if (searchResponse.getScrollId() != null) {
+                ClearScrollRequest clearScrollRequest = new ClearScrollRequest();
+                clearScrollRequest.scrollIds(List.of(searchResponse.getScrollId()));
+                client(LOCAL_CLUSTER).clearScroll(clearScrollRequest).get();
+            }
+        }
+        futures.forEach(ActionFuture::actionGet);
+    }
+}

server/src/main/java/org/elasticsearch/action/search/QueryPhaseResultConsumer.java

@@ -102,6 +102,7 @@ public void close() {
     public void consumeResult(SearchPhaseResult result, Runnable next) {
         super.consumeResult(result, () -> {});
         QuerySearchResult querySearchResult = result.queryResult();
+        querySearchResult.retainAggregationsUntilConsumed();
         progressListener.notifyQueryResult(querySearchResult.getShardIndex());
         pendingMerges.consume(querySearchResult, next);
     }

server/src/main/java/org/elasticsearch/search/fetch/QueryFetchSearchResult.java

@@ -72,4 +72,26 @@ public void writeTo(StreamOutput out) throws IOException {
         queryResult.writeTo(out);
         fetchResult.writeTo(out);
     }
+
+
+    @Override
+    public void incRef() {
+        queryResult.incRef();
+    }
+
+    @Override
+    public boolean tryIncRef() {
+        return queryResult.tryIncRef();
+    }
+
+    @Override
+    public boolean decRef() {
+        return queryResult.decRef();
+    }
+
+    @Override
+    public boolean hasReferences() {
+        return queryResult.hasReferences();
+    }
+
 }

server/src/main/java/org/elasticsearch/search/fetch/ScrollQueryFetchSearchResult.java

@@ -63,4 +63,25 @@ public void writeTo(StreamOutput out) throws IOException {
         getSearchShardTarget().writeTo(out);
         result.writeTo(out);
     }
+
+
+    @Override
+    public void incRef() {
+        result.incRef();
+    }
+
+    @Override
+    public boolean tryIncRef() {
+        return result.tryIncRef();
+    }
+
+    @Override
+    public boolean decRef() {
+        return result.decRef();
+    }
+
+    @Override
+    public boolean hasReferences() {
+        return result.hasReferences();
+    }
 }

server/src/main/java/org/elasticsearch/search/query/QuerySearchResult.java

@@ -10,11 +10,14 @@
 
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.TotalHits;
+import org.elasticsearch.Assertions;
 import org.elasticsearch.Version;
+import org.elasticsearch.common.bytes.ReleasableBytesReference;
 import org.elasticsearch.common.io.stream.DelayableWriteable;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
 import org.elasticsearch.common.lucene.search.TopDocsAndMaxScore;
+import org.elasticsearch.core.AbstractRefCounted;
 import org.elasticsearch.core.Releasables;
 import org.elasticsearch.search.DocValueFormat;
 import org.elasticsearch.search.RescoreDocIds;
@@ -33,7 +36,7 @@
 import static org.elasticsearch.common.lucene.Lucene.writeTopDocs;
 
 public final class QuerySearchResult extends SearchPhaseResult {
-
+    private final static org.apache.logging.log4j.Logger logger = org.apache.logging.log4j.LogManager.getLogger(QuerySearchResult.class);
     private int from;
     private int size;
     private TopDocsAndMaxScore topDocsAndMaxScore;
@@ -60,6 +63,13 @@ public final class QuerySearchResult extends SearchPhaseResult {
 
     private final boolean isNull;
 
+    private final AbstractRefCounted refCounted = AbstractRefCounted.of(() -> {
+        if (aggregations != null) {
+            aggregations.close();
+            aggregations = null;
+        }
+    });
+
     public QuerySearchResult() {
         this(false);
     }
@@ -187,6 +197,11 @@ public boolean hasAggs() {
         return hasAggs;
     }
 
+    public void retainAggregationsUntilConsumed() {
+        if (aggregations != null) {
+            incRef();
+        }
+    }
     /**
      * Returns and nulls out the aggregation for this search results. This allows to free up memory once the aggregation is consumed.
      * @throws IllegalStateException if the aggregations have already been consumed.
@@ -200,13 +215,15 @@ public InternalAggregations consumeAggs() {
         } finally {
             aggregations.close();
             aggregations = null;
+            decRef();
         }
     }
 
     public void releaseAggs() {
         if (aggregations != null) {
             aggregations.close();
             aggregations = null;
+            decRef();
         }
     }
 
@@ -405,4 +422,37 @@ public TotalHits getTotalHits() {
     public float getMaxScore() {
         return maxScore;
     }
+
+    @Override
+    public void incRef() {
+        refCounted.incRef();
+    }
+
+    @Override
+    public boolean tryIncRef() {
+        return refCounted.tryIncRef();
+    }
+
+    @Override
+    public boolean decRef() {
+        return refCounted.decRef();
+    }
+
+    @Override
+    public boolean hasReferences() {
+        return refCounted.hasReferences();
+    }
+
+    // todo: remove this and the logger before merging
+    @Override
+    protected void finalize() throws Throwable {
+        if (Assertions.ENABLED) {
+            if (refCounted.hasReferences()) {
+                // cannot really assert this.
+                logger.info("LEAK: " + this);
+//                 finalizer thread errors are not reported anywhere.
+//                new Thread(() -> { assert refCounted.hasReferences() == false : where; }).start();
+            }
+        }
+    }
 }

server/src/main/java/org/elasticsearch/search/query/ScrollQuerySearchResult.java

@@ -53,4 +53,24 @@ public void writeTo(StreamOutput out) throws IOException {
         getSearchShardTarget().writeTo(out);
         result.writeTo(out);
     }
+
+    @Override
+    public void incRef() {
+        result.incRef();
+    }
+
+    @Override
+    public boolean tryIncRef() {
+        return result.tryIncRef();
+    }
+
+    @Override
+    public boolean decRef() {
+        return result.decRef();
+    }
+
+    @Override
+    public boolean hasReferences() {
+        return result.hasReferences();
+    }
 }