Add support for CollectorManager wrapping with profiling with concurr…

…ent aggregation

Signed-off-by: Ticheng Lin <[email protected]>

server/src/main/java/org/opensearch/search/aggregations/AggregationCollectorManager.java

@@ -12,7 +12,6 @@
 import org.apache.lucene.search.CollectorManager;
 import org.opensearch.common.CheckedFunction;
 import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.profile.query.InternalProfileCollector;
 import org.opensearch.search.query.ReduceableSearchResult;
 
 import java.io.IOException;
@@ -42,12 +41,16 @@ class AggregationCollectorManager implements CollectorManager<Collector, Reducea
 
     @Override
     public Collector newCollector() throws IOException {
-        final Collector collector = createCollector(context, aggProvider.apply(context), collectorReason);
+        final Collector collector = createCollector(aggProvider.apply(context));
         // For Aggregations we should not have a NO_OP_Collector
         assert collector != BucketCollector.NO_OP_COLLECTOR;
         return collector;
     }
 
+    public String getCollectorReason() {
+        return collectorReason;
+    }
+
     @Override
     public ReduceableSearchResult reduce(Collection<Collector> collectors) throws IOException {
         final List<Aggregator> aggregators = context.bucketCollectorProcessor().toAggregators(collectors);
@@ -77,17 +80,9 @@ public ReduceableSearchResult reduce(Collection<Collector> collectors) throws IO
         }
     }
 
-    static Collector createCollector(SearchContext context, List<Aggregator> collectors, String reason) throws IOException {
+    static Collector createCollector(List<Aggregator> collectors) throws IOException {
         Collector collector = MultiBucketCollector.wrap(collectors);
         ((BucketCollector) collector).preCollection();
-        if (context.getProfilers() != null) {
-            collector = new InternalProfileCollector(
-                collector,
-                reason,
-                // TODO: report on child aggs as well
-                Collections.emptyList()
-            );
-        }
         return collector;
     }
 }

server/src/main/java/org/opensearch/search/aggregations/ConcurrentAggregationProcessor.java

@@ -65,12 +65,12 @@ public void postProcess(SearchContext context) {
         try {
             if (globalCollectorManager != null) {
                 Query query = context.buildFilteredQuery(Queries.newMatchAllQuery());
-                globalCollectorManager = new InternalProfileCollectorManager(
-                    globalCollectorManager,
-                    CollectorResult.REASON_AGGREGATION_GLOBAL,
-                    Collections.emptyList()
-                );
                 if (context.getProfilers() != null) {
+                    globalCollectorManager = new InternalProfileCollectorManager(
+                        globalCollectorManager,
+                        CollectorResult.REASON_AGGREGATION_GLOBAL,
+                        Collections.emptyList()
+                    );
                     context.getProfilers().addQueryProfiler().setCollector((InternalProfileComponent) globalCollectorManager);
                 }
                 final ReduceableSearchResult result = context.searcher().search(query, globalCollectorManager);

server/src/main/java/org/opensearch/search/aggregations/DefaultAggregationProcessor.java

@@ -11,10 +11,11 @@
 import org.apache.lucene.search.Query;
 import org.opensearch.common.lucene.search.Queries;
 import org.opensearch.search.internal.SearchContext;
-import org.opensearch.search.profile.query.InternalProfileComponent;
+import org.opensearch.search.profile.query.InternalProfileCollector;
 import org.opensearch.search.query.QueryPhaseExecutionException;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.List;
 
 /**
@@ -74,7 +75,13 @@ public void postProcess(SearchContext context) {
                     if (context.getProfilers() != null) {
                         context.getProfilers()
                             .addQueryProfiler()
-                            .setCollector((InternalProfileComponent) globalCollectorManager.newCollector());
+                            .setCollector(
+                                new InternalProfileCollector(
+                                    globalCollectorManager.newCollector(),
+                                    globalCollectorManager.getCollectorReason(),
+                                    Collections.emptyList()
+                                )
+                            );
                     }
                     context.searcher().search(query, globalCollectorManager.newCollector());
                     globalCollectorManager.reduce(List.of()).reduce(context.queryResult());

server/src/main/java/org/opensearch/search/query/QueryCollectorContext.java

@@ -42,6 +42,7 @@
 import org.apache.lucene.search.Weight;
 import org.opensearch.common.lucene.MinimumScoreCollector;
 import org.opensearch.common.lucene.search.FilteredCollector;
+import org.opensearch.search.profile.query.CollectorResult;
 import org.opensearch.search.profile.query.InternalProfileCollector;
 import org.opensearch.search.profile.query.InternalProfileCollectorManager;
 
@@ -99,11 +100,15 @@ protected InternalProfileCollector createWithProfiler(InternalProfileCollector i
      */
     protected InternalProfileCollectorManager createWithProfiler(InternalProfileCollectorManager in) throws IOException {
         final CollectorManager<? extends Collector, ReduceableSearchResult> manager = createManager(in);
-        return new InternalProfileCollectorManager(
-            manager,
-            profilerName,
-            in != null ? Collections.singletonList(in) : Collections.emptyList()
-        );
+        if (manager instanceof InternalProfileCollectorManager) {
+            return (InternalProfileCollectorManager) manager;
+        } else {
+            return new InternalProfileCollectorManager(
+                manager,
+                profilerName,
+                in != null ? Collections.singletonList(in) : Collections.emptyList()
+            );
+        }
     }
 
     /**
@@ -198,16 +203,33 @@ protected InternalProfileCollector createWithProfiler(InternalProfileCollector i
 
                 for (CollectorManager<? extends Collector, ReduceableSearchResult> manager : subs) {
                     final Collector collector = manager.newCollector();
-                    if (!(collector instanceof InternalProfileCollector)) {
-                        throw new IllegalArgumentException("non-profiling collector");
-                    }
-                    subCollectors.add((InternalProfileCollector) collector);
+                    subCollectors.add(new InternalProfileCollector(collector, CollectorResult.REASON_AGGREGATION, Collections.emptyList()));
                 }
 
                 final Collector collector = MultiCollector.wrap(subCollectors);
                 return new InternalProfileCollector(collector, REASON_SEARCH_MULTI, subCollectors);
             }
 
+            @Override
+            protected InternalProfileCollectorManager createWithProfiler(InternalProfileCollectorManager in) {
+                final List<CollectorManager<?, ReduceableSearchResult>> managers = new ArrayList<>();
+                final List<InternalProfileCollectorManager> children = new ArrayList<>();
+                managers.add(in);
+                children.add(in);
+                for (CollectorManager<? extends Collector, ReduceableSearchResult> manager : subs) {
+                    InternalProfileCollectorManager subCollectorManager = new InternalProfileCollectorManager(
+                        manager,
+                        CollectorResult.REASON_AGGREGATION,
+                        Collections.emptyList()
+                    );
+                    managers.add(subCollectorManager);
+                    children.add(subCollectorManager);
+                }
+                CollectorManager<? extends Collector, ReduceableSearchResult> multiCollectorManager = QueryCollectorManagerContext
+                    .createMultiCollectorManager(managers);
+                return new InternalProfileCollectorManager(multiCollectorManager, REASON_SEARCH_MULTI, children);
+            }
+
             @Override
             CollectorManager<? extends Collector, ReduceableSearchResult> createManager(
                 CollectorManager<? extends Collector, ReduceableSearchResult> in