diff --git a/server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java b/server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java index 34f4a1568f21c..ce83b21582ab6 100644 --- a/server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java +++ b/server/src/main/java/org/elasticsearch/search/aggregations/AggregationPhase.java @@ -11,9 +11,6 @@ import org.elasticsearch.action.search.SearchShardTask; import org.elasticsearch.search.aggregations.support.TimeSeriesIndexSearcher; import org.elasticsearch.search.internal.SearchContext; -import org.elasticsearch.search.profile.query.CollectorResult; -import org.elasticsearch.search.profile.query.InternalProfileCollector; -import org.elasticsearch.search.profile.query.InternalProfileCollectorManager; import org.elasticsearch.search.query.QueryPhase; import org.elasticsearch.search.query.SingleThreadCollectorManager; @@ -55,12 +52,7 @@ public static void preProcess(SearchContext context) { } else { collector = bucketCollector.asCollector(); } - if (context.getProfilers() != null) { - InternalProfileCollector profileCollector = new InternalProfileCollector(collector, CollectorResult.REASON_AGGREGATION); - context.aggregations().registerAggsCollectorManager(new InternalProfileCollectorManager(profileCollector)); - } else { - context.aggregations().registerAggsCollectorManager(new SingleThreadCollectorManager(collector)); - } + context.aggregations().registerAggsCollectorManager(new SingleThreadCollectorManager(collector)); } private static List getCancellationChecks(SearchContext context) { diff --git a/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollectorManager.java b/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollectorManager.java deleted file mode 100644 index 90a0891879fc7..0000000000000 --- a/server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollectorManager.java +++ /dev/null @@ -1,26 +0,0 @@ -/* - * 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.profile.query; - -import org.elasticsearch.search.query.SingleThreadCollectorManager; - -/** - * A {@link SingleThreadCollectorManager} that wraps a {@link InternalProfileCollector}. - * It delegates all the profiling to the generated collector via {@link #getCollectorTree()}. - */ -public final class InternalProfileCollectorManager extends SingleThreadCollectorManager { - - public InternalProfileCollectorManager(InternalProfileCollector collector) { - super(collector); - } - - public CollectorResult getCollectorTree() { - return ((InternalProfileCollector) collector).getCollectorTree(); - } -} diff --git a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java index dd627409ab6b9..02c1bcc9c2563 100644 --- a/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java +++ b/server/src/main/java/org/elasticsearch/search/query/QueryPhase.java @@ -36,8 +36,8 @@ import org.elasticsearch.search.internal.ContextIndexSearcher; import org.elasticsearch.search.internal.ScrollContext; import org.elasticsearch.search.internal.SearchContext; -import org.elasticsearch.search.profile.query.InternalProfileCollector; -import org.elasticsearch.search.profile.query.InternalProfileCollectorManager; +import org.elasticsearch.search.profile.query.CollectorResult; +import org.elasticsearch.search.profile.query.ProfileCollectorManager; import org.elasticsearch.search.rank.RankSearchContext; import org.elasticsearch.search.rank.RankShardContext; import org.elasticsearch.search.rescore.RescorePhase; @@ -188,20 +188,24 @@ static void addCollectorsAndSearch(SearchContext searchContext) throws QueryPhas searchContext, searchContext.parsedPostFilter() != null || searchContext.minimumScore() != null ); - - CollectorManager topDocsCollectorManager = topDocsFactory.collectorManager(); + CollectorManager topDocsCollectorManager = topDocsFactory.collectorManager(); + ProfileCollectorManager topDocsProfileCollectorManager = null; if (searchContext.getProfilers() != null) { - Collector topDocsCollector = topDocsCollectorManager.newCollector(); - InternalProfileCollector profileCollector = new InternalProfileCollector(topDocsCollector, topDocsFactory.profilerName); - topDocsCollectorManager = new InternalProfileCollectorManager(profileCollector); + topDocsProfileCollectorManager = new ProfileCollectorManager<>(topDocsCollectorManager, topDocsFactory.profilerName); + topDocsCollectorManager = topDocsProfileCollectorManager; } - Collector topDocsCollector = topDocsCollectorManager.newCollector(); - Collector aggsCollector = null; - Weight postFilterWeight = null; + CollectorManager aggsCollectorManager = null; + ProfileCollectorManager aggsProfileCollectorManager = null; if (searchContext.aggregations() != null) { - aggsCollector = searchContext.aggregations().getAggsCollectorManager().newCollector(); + aggsCollectorManager = searchContext.aggregations().getAggsCollectorManager(); + if (searchContext.getProfilers() != null) { + aggsProfileCollectorManager = new ProfileCollectorManager<>(aggsCollectorManager, CollectorResult.REASON_AGGREGATION); + aggsCollectorManager = aggsProfileCollectorManager; + } } + + Weight postFilterWeight = null; if (searchContext.parsedPostFilter() != null) { postFilterWeight = searcher.createWeight( searcher.rewrite(searchContext.parsedPostFilter().query()), @@ -209,34 +213,26 @@ static void addCollectorsAndSearch(SearchContext searchContext) throws QueryPhas 1f ); } - QueryPhaseCollector queryPhaseCollector = new QueryPhaseCollector( - topDocsCollector, + + QueryPhaseCollector.CollectorManager queryPhaseCollectorManager = QueryPhaseCollector.createManager( + topDocsCollectorManager, postFilterWeight, searchContext.terminateAfter(), - aggsCollector, + aggsCollectorManager, searchContext.minimumScore() ); - - SingleThreadCollectorManager collectorManager; + final CollectorManager collectorManager; if (searchContext.getProfilers() == null) { - collectorManager = new SingleThreadCollectorManager(queryPhaseCollector); + collectorManager = queryPhaseCollectorManager; } else { - InternalProfileCollector profileCollector; - if (aggsCollector == null) { - profileCollector = new InternalProfileCollector( - queryPhaseCollector, - REASON_SEARCH_QUERY_PHASE, - (InternalProfileCollector) topDocsCollector - ); - } else { - profileCollector = new InternalProfileCollector( - queryPhaseCollector, - REASON_SEARCH_QUERY_PHASE, - (InternalProfileCollector) topDocsCollector, - (InternalProfileCollector) aggsCollector - ); - } - collectorManager = new InternalProfileCollectorManager(profileCollector); + ProfileCollectorManager queryPhaseProfileManager = new ProfileCollectorManager<>( + queryPhaseCollectorManager, + REASON_SEARCH_QUERY_PHASE, + topDocsProfileCollectorManager, + aggsProfileCollectorManager + ); + searchContext.getProfilers().getCurrentQueryProfiler().setCollectorManager(queryPhaseProfileManager::getCollectorTree); + collectorManager = queryPhaseProfileManager; } final Runnable timeoutRunnable = getTimeoutCheck(searchContext); @@ -246,7 +242,7 @@ static void addCollectorsAndSearch(SearchContext searchContext) throws QueryPhas try { searchWithCollectorManager(searchContext, searcher, query, collectorManager, timeoutRunnable != null); - if (queryPhaseCollector.isTerminatedAfter()) { + if (queryPhaseCollectorManager.isTerminatedAfter()) { queryResult.terminatedEarly(true); } queryResult.topDocs(topDocsFactory.topDocsAndMaxScore(), topDocsFactory.sortValueFormats); @@ -274,14 +270,9 @@ private static void searchWithCollectorManager( SearchContext searchContext, ContextIndexSearcher searcher, Query query, - CollectorManager collectorManager, + CollectorManager collectorManager, boolean timeoutSet ) throws IOException { - if (searchContext.getProfilers() != null) { - searchContext.getProfilers() - .getCurrentQueryProfiler() - .setCollectorManager(((InternalProfileCollectorManager) collectorManager)::getCollectorTree); - } QuerySearchResult queryResult = searchContext.queryResult(); try { searcher.search(query, collectorManager);