Extend Profiler to allow for non-timing info (opensearch-project#18631)

* copied from other branch, ready for pr

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* cleanup from merging

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* another spotless apply and fixed breaking change

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* another spotlessApply

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* forgot to change back concurrent agg stuff

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* forgot to spotless after prev fix

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* added more unit tests for coverage

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* changed to supplier, removed timer ctors

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* made feedback changes

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* forgot javadoc for new class

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* forgot to revert name to toString

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* fix count check in concurrentquerypb

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

* unmodifiable map add

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>

---------

Signed-off-by: TJ Neuenfeldt <tjneu@amazon.com>
Signed-off-by: TJ Neuenfeldt <56700214+neuenfeldttj@users.noreply.github.com>

Author: neuenfeldttj

CHANGELOG.md

@@ -15,6 +15,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Introduce a new pull-based ingestion plugin for file-based indexing (for local testing) ([#18591](https://github.com/opensearch-project/OpenSearch/pull/18591))
 - Add support for search pipeline in search and msearch template ([#18564](https://github.com/opensearch-project/OpenSearch/pull/18564))
 - Add BooleanQuery rewrite moving constant-scoring must clauses to filter clauses ([#18510](https://github.com/opensearch-project/OpenSearch/issues/18510))
+- Add support for non-timing info in profiler ([#18460](https://github.com/opensearch-project/OpenSearch/issues/18460))
 
 ### Changed
 - Update Subject interface to use CheckedRunnable ([#18570](https://github.com/opensearch-project/OpenSearch/issues/18570))

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

@@ -213,7 +213,7 @@ public Weight createWeight(Query query, ScoreMode scoreMode, float boost) throws
             // createWeight() is called for each query in the tree, so we tell the queryProfiler
             // each invocation so that it can build an internal representation of the query
             // tree
-            ContextualProfileBreakdown<QueryTimingType> profile = profiler.getQueryBreakdown(query);
+            ContextualProfileBreakdown profile = profiler.getQueryBreakdown(query);
             Timer timer = profile.getTimer(QueryTimingType.CREATE_WEIGHT);
             timer.start();
             final Weight weight;

server/src/main/java/org/opensearch/search/profile/AbstractInternalProfileTree.java

@@ -45,7 +45,7 @@
  *
  * @opensearch.internal
  */
-public abstract class AbstractInternalProfileTree<PB extends AbstractProfileBreakdown<?>, E> {
+public abstract class AbstractInternalProfileTree<PB extends AbstractProfileBreakdown, E> {
 
     protected ArrayList<PB> breakdowns;
     /** Maps the Query to it's list of children.  This is basically the dependency tree */

server/src/main/java/org/opensearch/search/profile/AbstractProfileBreakdown.java

@@ -32,9 +32,12 @@
 
 package org.opensearch.search.profile;
 
+import java.util.Collection;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.Map;
+import java.util.TreeMap;
+import java.util.function.Supplier;
+import java.util.stream.Collectors;
 
 import static java.util.Collections.emptyMap;
 
@@ -45,58 +48,51 @@
  *
  * @opensearch.internal
  */
-public abstract class AbstractProfileBreakdown<T extends Enum<T>> {
+public abstract class AbstractProfileBreakdown {
 
-    /**
-     * The accumulated timings for this query node
-     */
-    protected final Timer[] timings;
-    protected final T[] timingTypes;
-    public static final String TIMING_TYPE_COUNT_SUFFIX = "_count";
-    public static final String TIMING_TYPE_START_TIME_SUFFIX = "_start_time";
+    protected final Map<String, ProfileMetric> metrics;
 
     /** Sole constructor. */
-    public AbstractProfileBreakdown(Class<T> clazz) {
-        this.timingTypes = clazz.getEnumConstants();
-        timings = new Timer[timingTypes.length];
-        for (int i = 0; i < timings.length; ++i) {
-            timings[i] = new Timer();
-        }
+    public AbstractProfileBreakdown(Collection<Supplier<ProfileMetric>> metricSuppliers) {
+        this.metrics = metricSuppliers.stream().map(Supplier::get).collect(Collectors.toMap(ProfileMetric::getName, metric -> metric));
     }
 
-    public Timer getTimer(T timing) {
-        return timings[timing.ordinal()];
+    public Timer getTimer(Enum<?> type) {
+        ProfileMetric metric = metrics.get(type.toString());
+        assert metric instanceof Timer : "Metric " + type + " is not a timer";
+        return (Timer) metric;
     }
 
-    public void setTimer(T timing, Timer timer) {
-        timings[timing.ordinal()] = timer;
+    public ProfileMetric getMetric(String name) {
+        return metrics.get(name);
     }
 
     /**
-     * Build a timing count breakdown for current instance
+     * Build a breakdown for current instance
      */
     public Map<String, Long> toBreakdownMap() {
-        Map<String, Long> map = new HashMap<>(this.timings.length * 3);
-        for (T timingType : this.timingTypes) {
-            map.put(timingType.toString(), this.timings[timingType.ordinal()].getApproximateTiming());
-            map.put(timingType + TIMING_TYPE_COUNT_SUFFIX, this.timings[timingType.ordinal()].getCount());
-            map.put(timingType + TIMING_TYPE_START_TIME_SUFFIX, this.timings[timingType.ordinal()].getEarliestTimerStartTime());
+        Map<String, Long> map = new TreeMap<>();
+        for (Map.Entry<String, ProfileMetric> entry : metrics.entrySet()) {
+            map.putAll(entry.getValue().toBreakdownMap());
         }
         return Collections.unmodifiableMap(map);
     }
 
+    public long toNodeTime() {
+        long total = 0;
+        for (Map.Entry<String, ProfileMetric> entry : metrics.entrySet()) {
+            if (entry.getValue() instanceof Timer t) {
+                total += t.getApproximateTiming();
+            }
+        }
+        return total;
+    }
+
     /**
      * Fetch extra debugging information.
      */
     public Map<String, Object> toDebugMap() {
         return emptyMap();
     }
 
-    public long toNodeTime() {
-        long total = 0;
-        for (T timingType : timingTypes) {
-            total += timings[timingType.ordinal()].getApproximateTiming();
-        }
-        return total;
-    }
 }

server/src/main/java/org/opensearch/search/profile/AbstractProfiler.java

@@ -39,7 +39,7 @@
  *
  * @opensearch.internal
  */
-public class AbstractProfiler<PB extends AbstractProfileBreakdown<?>, E> {
+public abstract class AbstractProfiler<PB extends AbstractProfileBreakdown, E> {
 
     protected final AbstractInternalProfileTree<PB, E> profileTree;
 

server/src/main/java/org/opensearch/search/profile/ContextualProfileBreakdown.java

@@ -11,26 +11,29 @@
 import org.apache.lucene.index.LeafReaderContext;
 import org.apache.lucene.search.Collector;
 
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Supplier;
 
 /**
  * Provide contextual profile breakdowns which are associated with freestyle context. Used when concurrent
  * search over segments is activated and each collector needs own non-shareable profile breakdown instance.
  *
  * @opensearch.internal
  */
-public abstract class ContextualProfileBreakdown<T extends Enum<T>> extends AbstractProfileBreakdown<T> {
-    public ContextualProfileBreakdown(Class<T> clazz) {
-        super(clazz);
+public abstract class ContextualProfileBreakdown extends AbstractProfileBreakdown {
+
+    public ContextualProfileBreakdown(Collection<Supplier<ProfileMetric>> metrics) {
+        super(metrics);
     }
 
     /**
      * Return (or create) contextual profile breakdown instance
      * @param context freestyle context
      * @return contextual profile breakdown instance
      */
-    public abstract AbstractProfileBreakdown<T> context(Object context);
+    public abstract AbstractProfileBreakdown context(Object context);
 
     public void associateCollectorToLeaves(Collector collector, LeafReaderContext leaf) {}
 

server/src/main/java/org/opensearch/search/profile/ProfileMetric.java

@@ -0,0 +1,40 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.profile;
+
+import org.opensearch.common.annotation.ExperimentalApi;
+
+import java.util.Map;
+
+/**
+ * A metric for profiling.
+ */
+@ExperimentalApi
+public abstract class ProfileMetric {
+
+    private final String name;
+
+    public ProfileMetric(String name) {
+        this.name = name;
+    }
+
+    /**
+     *
+     * @return name of the metric
+     */
+    public String getName() {
+        return name;
+    }
+
+    /**
+     *
+     * @return map representation of breakdown
+     */
+    abstract public Map<String, Long> toBreakdownMap();
+}

server/src/main/java/org/opensearch/search/profile/ProfileMetricUtil.java

@@ -0,0 +1,38 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.search.profile;
+
+import org.opensearch.search.profile.aggregation.AggregationTimingType;
+import org.opensearch.search.profile.query.QueryTimingType;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.function.Supplier;
+
+/**
+ * Utility class to provide profile metrics to breakdowns.
+ */
+public class ProfileMetricUtil {
+
+    public static Collection<Supplier<ProfileMetric>> getDefaultQueryProfileMetrics() {
+        Collection<Supplier<ProfileMetric>> metrics = new ArrayList<>();
+        for (QueryTimingType type : QueryTimingType.values()) {
+            metrics.add(() -> new Timer(type.toString()));
+        }
+        return metrics;
+    }
+
+    public static Collection<Supplier<ProfileMetric>> getAggregationProfileMetrics() {
+        Collection<Supplier<ProfileMetric>> metrics = new ArrayList<>();
+        for (AggregationTimingType type : AggregationTimingType.values()) {
+            metrics.add(() -> new Timer(type.toString()));
+        }
+        return metrics;
+    }
+}

server/src/main/java/org/opensearch/search/profile/ProfileResult.java

@@ -271,7 +271,7 @@ static void removeStartTimeFields(Map<String, Long> modifiedBreakdown) {
         Iterator<Map.Entry<String, Long>> iterator = modifiedBreakdown.entrySet().iterator();
         while (iterator.hasNext()) {
             Map.Entry<String, Long> entry = iterator.next();
-            if (entry.getKey().endsWith(AbstractProfileBreakdown.TIMING_TYPE_START_TIME_SUFFIX)) {
+            if (entry.getKey().endsWith(Timer.TIMING_TYPE_START_TIME_SUFFIX)) {
                 iterator.remove();
             }
         }

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

@@ -79,7 +79,7 @@ public QueryProfiler addQueryProfiler() {
 
     /** Get the current profiler. */
     public QueryProfiler getCurrentQueryProfiler() {
-        return queryProfilers.get(queryProfilers.size() - 1);
+        return queryProfilers.getLast();
     }
 
     /** Return the list of all created {@link QueryProfiler}s so far. */