LUCENE-9620: Address profiling test failures.

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerTimingType.java

@@ -22,6 +22,7 @@ import java.util.Locale;
 /** This enum breaks down the query into different sections to describe what was timed. */
 public enum QueryProfilerTimingType {
   CREATE_WEIGHT,
+  COUNT,
   BUILD_SCORER,
   NEXT_DOC,
   ADVANCE,

lucene/sandbox/src/java/org/apache/lucene/sandbox/search/QueryProfilerWeight.java

@@ -41,6 +41,17 @@ class QueryProfilerWeight extends Weight {
     this.profile = profile;
   }
 
+  @Override
+  public int count(LeafReaderContext context) throws IOException {
+    QueryProfilerTimer timer = profile.getTimer(QueryProfilerTimingType.COUNT);
+    timer.start();
+    try {
+      return subQueryWeight.count(context);
+    } finally {
+      timer.stop();
+    }
+  }
+
   @Override
   public Scorer scorer(LeafReaderContext context) throws IOException {
     ScorerSupplier supplier = scorerSupplier(context);

lucene/sandbox/src/test/org/apache/lucene/sandbox/search/TestQueryProfilerIndexSearcher.java

@@ -83,6 +83,7 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     Map<String, Long> breakdown = results.get(0).getTimeBreakdown();
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString()), greaterThan(0L));
+    MatcherAssert.assertThat(breakdown.get(QueryProfilerTimingType.COUNT.toString()), equalTo(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString()), greaterThan(0L));
     MatcherAssert.assertThat(
@@ -96,6 +97,8 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString() + "_count"),
         greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString() + "_count"), equalTo(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L));
     MatcherAssert.assertThat(
@@ -156,6 +159,7 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     Map<String, Long> breakdown = results.get(0).getTimeBreakdown();
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString()), greaterThan(0L));
+    MatcherAssert.assertThat(breakdown.get(QueryProfilerTimingType.COUNT.toString()), equalTo(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString()), greaterThan(0L));
     MatcherAssert.assertThat(
@@ -168,6 +172,8 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString() + "_count"),
         greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString() + "_count"), equalTo(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L));
     MatcherAssert.assertThat(
@@ -189,7 +195,36 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     searcher.count(query); // will use index stats
 
     List<QueryProfilerResult> results = searcher.getProfileResult();
-    assertEquals(0, results.size());
+    assertEquals(1, results.size());
+    Map<String, Long> breakdown = results.get(0).getTimeBreakdown();
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString()), greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString()), greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString()), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.NEXT_DOC.toString()), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.ADVANCE.toString()), equalTo(0L));
+    MatcherAssert.assertThat(breakdown.get(QueryProfilerTimingType.SCORE.toString()), equalTo(0L));
+    MatcherAssert.assertThat(breakdown.get(QueryProfilerTimingType.MATCH.toString()), equalTo(0L));
+
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString() + "_count"),
+        greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString() + "_count"), greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString() + "_count"), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.NEXT_DOC.toString() + "_count"), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.ADVANCE.toString() + "_count"), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.SCORE.toString() + "_count"), equalTo(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.MATCH.toString() + "_count"), equalTo(0L));
 
     long rewriteTime = searcher.getRewriteTime();
     MatcherAssert.assertThat(rewriteTime, greaterThan(0L));
@@ -204,6 +239,8 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     Map<String, Long> breakdown = results.get(0).getTimeBreakdown();
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString()), greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString()), greaterThan(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString()), greaterThan(0L));
     MatcherAssert.assertThat(
@@ -217,6 +254,8 @@ public class TestQueryProfilerIndexSearcher extends LuceneTestCase {
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.CREATE_WEIGHT.toString() + "_count"),
         greaterThan(0L));
+    MatcherAssert.assertThat(
+        breakdown.get(QueryProfilerTimingType.COUNT.toString() + "_count"), greaterThan(0L));
     MatcherAssert.assertThat(
         breakdown.get(QueryProfilerTimingType.BUILD_SCORER.toString() + "_count"), greaterThan(0L));
     MatcherAssert.assertThat(