Introduce ProfileCollectorManager (#96828)

This change adds a new Collector manager that can warp other CollectorManagers and subsequently
wraps all collectors generated by the underlying collector manager into an InternalProfileCollector.
This will be used to enable inter-segment search concurrency.

Author: cbuescher

server/src/main/java/org/elasticsearch/search/profile/Profilers.java

@@ -65,7 +65,7 @@ public SearchProfileQueryPhaseResult buildQueryPhaseResults() {
         QueryProfileShardResult result = new QueryProfileShardResult(
             queryProfiler.getTree(),
             queryProfiler.getRewriteTime(),
-            queryProfiler.getCollector()
+            queryProfiler.getCollectorResult()
         );
         AggregationProfileShardResult aggResults = new AggregationProfileShardResult(aggProfiler.getTree());
         return new SearchProfileQueryPhaseResult(Collections.singletonList(result), aggResults);

server/src/main/java/org/elasticsearch/search/profile/dfs/DfsProfiler.java

@@ -54,7 +54,7 @@ public void stopTimer(DfsTimingType dfsTimingType) {
 
     public QueryProfiler addQueryProfiler(InternalProfileCollectorManager collectorManager) {
         QueryProfiler queryProfiler = new QueryProfiler();
-        queryProfiler.setCollectorManager(collectorManager);
+        queryProfiler.setCollectorManager(collectorManager::getCollectorTree);
         knnQueryProfilers.add(queryProfiler);
         collectorManagerSet = true;
         return queryProfiler;
@@ -74,7 +74,7 @@ public SearchProfileDfsPhaseResult buildDfsPhaseResults() {
             queryProfileShardResult = new ArrayList<>(knnQueryProfilers.size());
             for (QueryProfiler queryProfiler : knnQueryProfilers) {
                 queryProfileShardResult.add(
-                    new QueryProfileShardResult(queryProfiler.getTree(), queryProfiler.getRewriteTime(), queryProfiler.getCollector())
+                    new QueryProfileShardResult(queryProfiler.getTree(), queryProfiler.getRewriteTime(), queryProfiler.getCollectorResult())
                 );
             }
         } else {

server/src/main/java/org/elasticsearch/search/profile/query/InternalProfileCollector.java

@@ -21,21 +21,27 @@
  * - collect()
  * - doSetNextReader()
  * - needsScores()
- *
+ * <p>
  * InternalProfiler facilitates the linking of the Collector graph
  */
 public class InternalProfileCollector extends ProfilerCollector {
 
     private final InternalProfileCollector[] children;
+    private final Collector wrappedCollector;
 
     public InternalProfileCollector(Collector collector, String reason, InternalProfileCollector... children) {
         super(collector, reason, Arrays.asList(children));
+        this.wrappedCollector = collector;
         this.children = children;
     }
 
+    public Collector getWrappedCollector() {
+        return wrappedCollector;
+    }
+
     /**
      * Creates a human-friendly representation of the Collector name.
-     *
+     * <p>
      * InternalBucket Collectors use the aggregation name in their toString() method,
      * which makes the profiled output a bit nicer.
      *
@@ -46,7 +52,7 @@ public InternalProfileCollector(Collector collector, String reason, InternalProf
     protected String deriveCollectorName(Collector c) {
         String s = c.getClass().getSimpleName();
 
-        // MutiCollector which wraps multiple BucketCollectors is generated
+        // MultiCollector which wraps multiple BucketCollectors is generated
         // via an anonymous class, so this corrects the lack of a name by
         // asking the enclosingClass
         if (s.equals("")) {

server/src/main/java/org/elasticsearch/search/profile/query/ProfileCollectorManager.java

@@ -0,0 +1,60 @@
+/*
+ * 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.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * A {@link CollectorManager} that takes another CollectorManager as input and wraps all Collectors generated by it
+ * in an {@link InternalProfileCollector}. It delegates all the profiling to the generated collectors via {@link #getCollectorTree()}
+ * and joins them up when its {@link #reduce} method is called. The profile result can
+ */
+public final class ProfileCollectorManager implements CollectorManager<InternalProfileCollector, Void> {
+
+    private final CollectorManager<Collector, ?> collectorManager;
+    private final String reason;
+    private CollectorResult collectorTree;
+
+    @SuppressWarnings("unchecked")
+    public ProfileCollectorManager(CollectorManager<? extends Collector, ?> collectorManager, String reason) {
+        this.collectorManager = (CollectorManager<Collector, ?>) collectorManager;
+        this.reason = reason;
+    }
+
+    @Override
+    public InternalProfileCollector newCollector() throws IOException {
+        return new InternalProfileCollector(collectorManager.newCollector(), reason);
+    }
+
+    public Void reduce(Collection<InternalProfileCollector> profileCollectors) throws IOException {
+        List<Collector> unwrapped = profileCollectors.stream()
+            .map(InternalProfileCollector::getWrappedCollector)
+            .collect(Collectors.toList());
+        collectorManager.reduce(unwrapped);
+
+        List<CollectorResult> resultsPerProfiler = profileCollectors.stream()
+            .map(ipc -> ipc.getCollectorTree())
+            .collect(Collectors.toList());
+        this.collectorTree = new CollectorResult(this.getClass().getSimpleName(), "segment_search", 0, resultsPerProfiler);
+        return null;
+    }
+
+    public CollectorResult getCollectorTree() {
+        if (this.collectorTree == null) {
+            throw new IllegalStateException("A collectorTree hasn't been set yet, call reduce() before attempting to retrieve it");
+        }
+        return this.collectorTree;
+    }
+}

server/src/main/java/org/elasticsearch/search/profile/query/QueryProfiler.java

@@ -12,6 +12,7 @@
 import org.elasticsearch.search.profile.AbstractProfiler;
 
 import java.util.Objects;
+import java.util.function.Supplier;
 
 /**
  * This class acts as a thread-local storage for profiling a query.  It also
@@ -29,18 +30,18 @@ public final class QueryProfiler extends AbstractProfiler<QueryProfileBreakdown,
     /**
      * The root Collector used in the search
      */
-    private InternalProfileCollectorManager collectorManager;
+    private Supplier<CollectorResult> collectorResultSupplier;
 
     public QueryProfiler() {
         super(new InternalQueryProfileTree());
     }
 
     /** Set the collector manager that is associated with this profiler. */
-    public void setCollectorManager(InternalProfileCollectorManager collectorManager) {
-        if (this.collectorManager != null) {
-            throw new IllegalStateException("The collector manager can only be set once.");
+    public void setCollectorManager(Supplier<CollectorResult> collectorResultSupplier) {
+        if (this.collectorResultSupplier != null) {
+            throw new IllegalStateException("The collector result supplier can only be set once.");
         }
-        this.collectorManager = Objects.requireNonNull(collectorManager);
+        this.collectorResultSupplier = Objects.requireNonNull(collectorResultSupplier);
     }
 
     /**
@@ -71,8 +72,8 @@ public long getRewriteTime() {
     /**
      * Return the current root Collector for this search
      */
-    public CollectorResult getCollector() {
-        return collectorManager.getCollectorTree();
+    public CollectorResult getCollectorResult() {
+        return this.collectorResultSupplier.get();
     }
 
 }

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

@@ -312,7 +312,9 @@ private static void searchWithCollectorManager(
         boolean timeoutSet
     ) throws IOException {
         if (searchContext.getProfilers() != null) {
-            searchContext.getProfilers().getCurrentQueryProfiler().setCollectorManager((InternalProfileCollectorManager) collectorManager);
+            searchContext.getProfilers()
+                .getCurrentQueryProfiler()
+                .setCollectorManager(((InternalProfileCollectorManager) collectorManager)::getCollectorTree);
         }
         QuerySearchResult queryResult = searchContext.queryResult();
         try {

server/src/test/java/org/elasticsearch/search/profile/query/ProfileCollectorManagerTests.java

@@ -0,0 +1,118 @@
+/*
+ * 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.apache.lucene.document.Document;
+import org.apache.lucene.document.Field;
+import org.apache.lucene.document.StringField;
+import org.apache.lucene.index.IndexReader;
+import org.apache.lucene.sandbox.search.ProfilerCollectorResult;
+import org.apache.lucene.search.Collector;
+import org.apache.lucene.search.CollectorManager;
+import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.MatchAllDocsQuery;
+import org.apache.lucene.search.TopDocs;
+import org.apache.lucene.search.TopScoreDocCollector;
+import org.apache.lucene.search.similarities.BM25Similarity;
+import org.apache.lucene.store.Directory;
+import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.apache.lucene.tests.search.DummyTotalHitCountCollector;
+import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.test.ESTestCase;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.List;
+
+public class ProfileCollectorManagerTests extends ESTestCase {
+
+    private static class TestCollector extends DummyTotalHitCountCollector {
+
+        private final int id;
+
+        TestCollector(int id) {
+            this.id = id;
+        }
+    }
+
+    /**
+     * This test checks that each new collector is a different instance on each call and that
+     * the call to reduce() is forwarded to the wrapped collector manager.
+     */
+    public void testBasic() throws IOException {
+        final SetOnce<Boolean> reduceCalled = new SetOnce<>();
+        ProfileCollectorManager pcm = new ProfileCollectorManager(new CollectorManager<>() {
+
+            private static int counter = 0;
+
+            @Override
+            public Collector newCollector() {
+                return new TestCollector(counter++);
+            }
+
+            @Override
+            public Void reduce(Collection<Collector> collectors) {
+                reduceCalled.set(true);
+                return null;
+            }
+        }, CollectorResult.REASON_SEARCH_TOP_HITS);
+        for (int i = 0; i < randomIntBetween(5, 10); i++) {
+            InternalProfileCollector internalProfileCollector = pcm.newCollector();
+            assertEquals(i, ((TestCollector) internalProfileCollector.getWrappedCollector()).id);
+        }
+        pcm.reduce(Collections.emptyList());
+        assertTrue(reduceCalled.get());
+    }
+
+    /**
+     * This test checks functionality with potentially more than one slice on a real searcher,
+     * wrapping a {@link TopScoreDocCollector} into  {@link ProfileCollectorManager} and checking the
+     * result from calling the collector tree contains profile results for each slice.
+     */
+    public void testManagerWithSearcher() throws IOException {
+        Directory directory = newDirectory();
+        try (RandomIndexWriter writer = new RandomIndexWriter(random(), directory, newIndexWriterConfig())) {
+            int numDocs = randomIntBetween(900, 1000);
+            for (int i = 0; i < numDocs; i++) {
+                Document doc = new Document();
+                doc.add(new StringField("field1", "value", Field.Store.NO));
+                writer.addDocument(doc);
+            }
+            writer.flush();
+            IndexReader reader = writer.getReader();
+            IndexSearcher searcher = newSearcher(reader);
+            int numSlices = searcher.getSlices() == null ? 1 : searcher.getSlices().length;
+            searcher.setSimilarity(new BM25Similarity());
+
+            CollectorManager<TopScoreDocCollector, TopDocs> topDocsManager = TopScoreDocCollector.createSharedManager(10, null, 1000);
+            TopDocs topDocs = searcher.search(new MatchAllDocsQuery(), topDocsManager);
+            assertEquals(numDocs, topDocs.totalHits.value);
+
+            String profileReason = "profiler_reason";
+            ProfileCollectorManager profileCollectorManager = new ProfileCollectorManager(topDocsManager, profileReason);
+
+            searcher.search(new MatchAllDocsQuery(), profileCollectorManager);
+
+            CollectorResult parent = profileCollectorManager.getCollectorTree();
+            assertEquals("ProfileCollectorManager", parent.getName());
+            assertEquals("segment_search", parent.getReason());
+            assertEquals(0, parent.getTime());
+            List<ProfilerCollectorResult> delegateCollectorResults = parent.getProfiledChildren();
+            assertEquals(numSlices, delegateCollectorResults.size());
+            for (ProfilerCollectorResult pcr : delegateCollectorResults) {
+                assertEquals("SimpleTopScoreDocCollector", pcr.getName());
+                assertEquals(profileReason, pcr.getReason());
+                assertTrue(pcr.getTime() > 0);
+            }
+            reader.close();
+        }
+        directory.close();
+    }
+}