tidy up post collection. (#98345)

server/src/main/java/org/elasticsearch/search/internal/ContextIndexSearcher.java

@@ -40,7 +40,6 @@ import org.apache.lucene.util.ThreadInterruptedException;
 import org.elasticsearch.common.util.concurrent.FutureUtils;
 import org.elasticsearch.core.Releasable;
 import org.elasticsearch.lucene.util.CombinedBitSet;
-import org.elasticsearch.search.aggregations.BucketCollector;
 import org.elasticsearch.search.dfs.AggregatedDfs;
 import org.elasticsearch.search.profile.Timer;
 import org.elasticsearch.search.profile.query.InternalProfileCollector;
@@ -471,11 +470,8 @@ public class ContextIndexSearcher extends IndexSearcher implements Releasable {
 
     private void doAggregationPostCollection(Collector collector) throws IOException {
         if (collector instanceof QueryPhaseCollector queryPhaseCollector) {
-            if (queryPhaseCollector.getAggsCollector() instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
-                aggsCollector.bucketCollector().postCollection();
-            }
-        }
-        if (collector instanceof InternalProfileCollector profilerCollector) {
+            queryPhaseCollector.doPostCollection();
+        } else if (collector instanceof InternalProfileCollector profilerCollector) {
             profilerCollector.doPostCollection();
         }
     }

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

@@ -82,11 +82,7 @@ public class InternalProfileCollector extends ProfilerCollector {
         if (wrappedCollector instanceof InternalProfileCollector profileCollector) {
             profileCollector.doPostCollection();
         } else if (wrappedCollector instanceof QueryPhaseCollector queryPhaseCollector) {
-            if (queryPhaseCollector.getAggsCollector() instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
-                aggsCollector.bucketCollector().postCollection();
-            } else if (queryPhaseCollector.getAggsCollector() instanceof InternalProfileCollector profileCollector) {
-                profileCollector.doPostCollection();
-            }
+            queryPhaseCollector.doPostCollection();
         } else if (wrappedCollector instanceof BucketCollector.BucketCollectorWrapper aggsCollector) {
             aggsCollector.bucketCollector().postCollection();
         }

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

@@ -22,6 +22,8 @@ import org.apache.lucene.search.ScorerSupplier;
 import org.apache.lucene.search.Weight;
 import org.apache.lucene.util.Bits;
 import org.elasticsearch.common.lucene.Lucene;
+import org.elasticsearch.search.aggregations.BucketCollector;
+import org.elasticsearch.search.profile.query.InternalProfileCollector;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -458,7 +460,11 @@ public final class QueryPhaseCollector implements Collector {
         }
     }
 
-    public Collector getAggsCollector() {
-        return aggsCollector;
+    public void doPostCollection() throws IOException {
+        if (aggsCollector instanceof BucketCollector.BucketCollectorWrapper bucketCollectorWrapper) {
+            bucketCollectorWrapper.bucketCollector().postCollection();
+        } else if (aggsCollector instanceof InternalProfileCollector profileCollector) {
+            profileCollector.doPostCollection();
+        }
     }
 }