From fc6a3f7992224cb913f9074a97b27a81a10692c5 Mon Sep 17 00:00:00 2001
From: Ticheng Lin <ticheng@amazon.com>
Date: Tue, 7 Nov 2023 21:36:08 -0800
Subject: [PATCH] Fix slice collectors to leaves association with profile
 enabled (#11134)

Signed-off-by: Ticheng Lin <ticheng@amazon.com>
---
 .../ConcurrentQueryProfileBreakdown.java      | 20 ++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java b/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
index 59ef01f9f947a..e11086f165719 100644
--- a/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
+++ b/server/src/main/java/org/opensearch/search/profile/query/ConcurrentQueryProfileBreakdown.java
@@ -65,6 +65,24 @@ public AbstractProfileBreakdown<QueryTimingType> context(Object context) {
     @Override
     public Map<String, Long> toBreakdownMap() {
         final Map<String, Long> topLevelBreakdownMapWithWeightTime = super.toBreakdownMap();
+        int totalLeafCount = 0;
+        for (List<LeafReaderContext> leaves : sliceCollectorsToLeaves.values()) {
+            totalLeafCount += leaves.size();
+        }
+        if (contexts.size() != totalLeafCount) {
+            // If the leaf counts in contexts map don't match those in sliceCollectorsToLeaves, this means that the current
+            // node is outside the concurrent search path, which occurred before the contextIndexSearcher::searchLeaf.
+            // For example, in the post filter query, we may create a profile node during filtered collector context creation.
+            // In this case, as there is no concurrent implementation, we can simply return the accumulated breakdown result
+            // instead of the statistical result.
+            final Map<String, Long> queryBreakdownMap = new HashMap<>(topLevelBreakdownMapWithWeightTime);
+            for (final AbstractProfileBreakdown<QueryTimingType> context : contexts.values()) {
+                for (final Map.Entry<String, Long> entry : context.toBreakdownMap().entrySet()) {
+                    queryBreakdownMap.merge(entry.getKey(), entry.getValue(), Long::sum);
+                }
+            }
+            return queryBreakdownMap;
+        }
         final long createWeightStartTime = topLevelBreakdownMapWithWeightTime.get(
             QueryTimingType.CREATE_WEIGHT + TIMING_TYPE_START_TIME_SUFFIX
         );
@@ -83,7 +101,7 @@ public Map<String, Long> toBreakdownMap() {
             // creates a new weight and breakdown map for each rewritten query. This new breakdown map captures the timing information for
             // the new rewritten query. The sliceCollectorsToLeaves is empty because this breakdown for rewritten query gets created later
             // in search leaf path which doesn't have collector. Also, this is not needed since this breakdown is per leaf and there is no
-            // concurrency involved. An empty sliceCollectorsToLeaves could also happen in the case of early termination.
+            // concurrency involved.
             AbstractProfileBreakdown<QueryTimingType> breakdown = contexts.values().iterator().next();
             queryNodeTime = breakdown.toNodeTime() + createWeightTime;
             maxSliceNodeTime = 0L;