Supports additional query timing types for profiling plugin query components

CHANGELOG.md

@@ -31,6 +31,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Added new Setting property UnmodifiableOnRestore to prevent updating settings on restore snapshot ([#16957](https://github.com/opensearch-project/OpenSearch/pull/16957))
 - Introduce Template query ([#16818](https://github.com/opensearch-project/OpenSearch/pull/16818))
 - Propagate the sourceIncludes and excludes fields from fetchSourceContext to FieldsVisitor. ([#17080](https://github.com/opensearch-project/OpenSearch/pull/17080))
+- Adds KNN related enums to support additional components for knn query during profiling ([#17146](https://github.com/opensearch-project/OpenSearch/pull/17146))
 
 ### Dependencies
 - Bump `com.google.cloud:google-cloud-core-http` from 2.23.0 to 2.47.0 ([#16504](https://github.com/opensearch-project/OpenSearch/pull/16504))

server/src/main/java/org/opensearch/index/query/QueryBuilder.java

@@ -38,6 +38,8 @@
 import org.opensearch.core.xcontent.ToXContentObject;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.Set;
 
 /**
  * Foundation class for all OpenSearch query builders
@@ -98,6 +100,10 @@ default QueryBuilder rewrite(QueryRewriteContext queryShardContext) throws IOExc
         return this;
     }
 
+    default Set<String> queryProfilerTimingTypes() {
+        return Collections.emptySet();
+    }
+
     /**
      * Recurse through the QueryBuilder tree, visiting any child QueryBuilder.
      * @param visitor a query builder visitor to be called by each query builder in the tree.

server/src/main/java/org/opensearch/node/Node.java

@@ -34,6 +34,7 @@
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.util.Constants;
 import org.opensearch.Build;
 import org.opensearch.ExceptionsHelper;
@@ -1386,7 +1387,8 @@ protected Node(
                 circuitBreakerService,
                 searchModule.getIndexSearcherExecutor(threadPool),
                 taskResourceTrackingService,
-                searchModule.getConcurrentSearchRequestDeciderFactories()
+                searchModule.getConcurrentSearchRequestDeciderFactories(),
+                searchModule.getQueryProfilerTimingTypes()
             );
 
             final List<PersistentTasksExecutor<?>> tasksExecutors = pluginsService.filterPlugins(PersistentTaskPlugin.class)
@@ -2050,7 +2052,8 @@ protected SearchService newSearchService(
         CircuitBreakerService circuitBreakerService,
         Executor indexSearcherExecutor,
         TaskResourceTrackingService taskResourceTrackingService,
-        Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories
+        Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories,
+        Map<Class<? extends Query>, Set<String>> profilerTimingsPerQuery
     ) {
         return new SearchService(
             clusterService,
@@ -2064,7 +2067,8 @@ protected SearchService newSearchService(
             circuitBreakerService,
             indexSearcherExecutor,
             taskResourceTrackingService,
-            concurrentSearchDeciderFactories
+            concurrentSearchDeciderFactories,
+            profilerTimingsPerQuery
         );
     }
 

server/src/main/java/org/opensearch/plugins/SearchPlugin.java

@@ -82,6 +82,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.ExecutorService;
 import java.util.function.BiFunction;
@@ -227,6 +228,13 @@ default Optional<ExecutorServiceProvider> getIndexSearcherExecutorProvider() {
         return Optional.empty();
     }
 
+    /**
+     * Register any additional profiler timing types that is needed across all your queries
+     */
+    default Map<Class<? extends Query>, Set<String>> registerProfilerTimingTypes() {
+        return emptyMap();
+    }
+
     /**
      * Executor service provider
      */

server/src/main/java/org/opensearch/search/SearchModule.java

@@ -33,6 +33,7 @@
 package org.opensearch.search;
 
 import org.apache.lucene.search.BooleanQuery;
+import org.apache.lucene.search.Query;
 import org.opensearch.common.NamedRegistry;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.geo.GeoShapeType;
@@ -300,9 +301,11 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.function.Consumer;
 import java.util.function.Function;
@@ -335,6 +338,7 @@ public class SearchModule {
     private final SearchPlugin.ExecutorServiceProvider indexSearcherExecutorProvider;
 
     private final Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories;
+    private final Map<Class<? extends Query>, Set<String>> profilerTimingsPerQuery = new HashMap<>();
 
     /**
      * Constructs a new SearchModule object
@@ -361,12 +365,19 @@ public SearchModule(Settings settings, List<SearchPlugin> plugins) {
         registerSearchExts(plugins);
         registerShapes();
         registerIntervalsSourceProviders();
+        registerProfilerTimingTypes(plugins);
         queryPhaseSearcher = registerQueryPhaseSearcher(plugins);
         indexSearcherExecutorProvider = registerIndexSearcherExecutorProvider(plugins);
         namedWriteables.addAll(SortValue.namedWriteables());
         concurrentSearchDeciderFactories = registerConcurrentSearchDeciderFactories(plugins);
     }
 
+    private void registerProfilerTimingTypes(List<SearchPlugin> plugins) {
+        for (SearchPlugin plugin : plugins) {
+            profilerTimingsPerQuery.putAll(plugin.registerProfilerTimingTypes());
+        }
+    }
+
     private Collection<ConcurrentSearchRequestDecider.Factory> registerConcurrentSearchDeciderFactories(List<SearchPlugin> plugins) {
         List<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories = new ArrayList<>();
         for (SearchPlugin plugin : plugins) {
@@ -383,6 +394,10 @@ public Collection<ConcurrentSearchRequestDecider.Factory> getConcurrentSearchReq
         return concurrentSearchDeciderFactories;
     }
 
+    public Map<Class<? extends Query>, Set<String>> getQueryProfilerTimingTypes() {
+        return profilerTimingsPerQuery;
+    }
+
     public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         return namedWriteables;
     }

server/src/main/java/org/opensearch/search/SearchService.java

@@ -36,6 +36,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.lucene.search.FieldDoc;
 import org.apache.lucene.search.IndexSearcher;
+import org.apache.lucene.search.Query;
 import org.apache.lucene.search.TopDocs;
 import org.opensearch.LegacyESVersion;
 import org.opensearch.OpenSearchException;
@@ -401,6 +402,7 @@ public class SearchService extends AbstractLifecycleComponent implements IndexEv
     private final String sessionId = UUIDs.randomBase64UUID();
     private final Executor indexSearcherExecutor;
     private final TaskResourceTrackingService taskResourceTrackingService;
+    private final Map<Class<? extends Query>, Set<String>> profilerTimingsPerQuery;
 
     public SearchService(
         ClusterService clusterService,
@@ -414,7 +416,8 @@ public SearchService(
         CircuitBreakerService circuitBreakerService,
         Executor indexSearcherExecutor,
         TaskResourceTrackingService taskResourceTrackingService,
-        Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories
+        Collection<ConcurrentSearchRequestDecider.Factory> concurrentSearchDeciderFactories,
+        Map<Class<? extends Query>, Set<String>> profilerTimingsPerQuery
     ) {
         Settings settings = clusterService.getSettings();
         this.threadPool = threadPool;
@@ -471,6 +474,7 @@ public SearchService(
         clusterService.getClusterSettings().addSettingsUpdateConsumer(CLUSTER_ALLOW_DERIVED_FIELD_SETTING, this::setAllowDerivedField);
 
         this.concurrentSearchDeciderFactories = concurrentSearchDeciderFactories;
+        this.profilerTimingsPerQuery = profilerTimingsPerQuery;
     }
 
     private void validateKeepAlives(TimeValue defaultKeepAlive, TimeValue maxKeepAlive) {
@@ -1541,7 +1545,7 @@ private void parseSource(DefaultSearchContext context, SearchSourceBuilder sourc
         }
         context.evaluateRequestShouldUseConcurrentSearch();
         if (source.profile()) {
-            context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch()));
+            context.setProfilers(new Profilers(context.searcher(), context.shouldUseConcurrentSearch(), this.profilerTimingsPerQuery));
         }
 
         if (this.indicesService.getCompositeIndexSettings() != null

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

@@ -154,6 +154,10 @@ public void setProfiler(QueryProfiler profiler) {
         this.profiler = profiler;
     }
 
+    public QueryProfiler getProfiler() {
+        return profiler;
+    }
+
     /**
      * Add a {@link Runnable} that will be run on a regular basis while accessing documents in the
      * DirectoryReader but also while collecting them and check for query cancellation or timeout.

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

@@ -39,6 +39,8 @@
 import java.util.Collections;
 import java.util.Deque;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * Base class for a profiling tree.
@@ -57,13 +59,16 @@ public abstract class AbstractInternalProfileTree<PB extends AbstractProfileBrea
     /** A temporary stack used to record where we are in the dependency tree. */
     protected Deque<Integer> stack;
     private int currentToken = 0;
+    /** Register of additional custom timings per element **/
+    private Map<Class<? extends E>, Set<String>> profilerTimingsPerElement;
 
-    public AbstractInternalProfileTree() {
+    public AbstractInternalProfileTree(Map<Class<? extends E>, Set<String>> profilerTimingsPerElement) {
         breakdowns = new ArrayList<>(10);
         stack = new ArrayDeque<>(10);
         tree = new ArrayList<>(10);
         elements = new ArrayList<>(10);
         roots = new ArrayList<>(10);
+        this.profilerTimingsPerElement = profilerTimingsPerElement == null ? Collections.emptyMap() : profilerTimingsPerElement;
     }
 
     /**
@@ -128,12 +133,17 @@ private PB addDependencyNode(E element, int token) {
         // Save our query for lookup later
         elements.add(element);
 
-        PB breakdown = createProfileBreakdown();
+        final Set<String> timings = profilerTimingsPerElement.get(element.getClass());
+        PB breakdown = createProfileBreakdown(timings == null ? Collections.emptySet() : timings);
         breakdowns.add(token, breakdown);
         return breakdown;
     }
 
-    protected abstract PB createProfileBreakdown();
+    public Map<Class<? extends E>, Set<String>> getProfilerTimingsPerElement() {
+        return profilerTimingsPerElement;
+    }
+
+    protected abstract PB createProfileBreakdown(Set<String> profilerTimingsPerElement);
 
     /**
      * Removes the last (e.g. most recent) value on the stack

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

@@ -32,9 +32,16 @@
 
 package org.opensearch.search.profile;
 
+import java.util.Arrays;
 import java.util.Collections;
-import java.util.HashMap;
+import java.util.Locale;
 import java.util.Map;
+import java.util.Objects;
+import java.util.Set;
+import java.util.TreeMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import static java.util.Collections.emptyMap;
 
@@ -50,37 +57,41 @@ public abstract class AbstractProfileBreakdown<T extends Enum<T>> {
     /**
      * The accumulated timings for this query node
      */
-    protected final Timer[] timings;
-    protected final T[] timingTypes;
+    protected final Map<String, Timer> timings;
     public static final String TIMING_TYPE_COUNT_SUFFIX = "_count";
     public static final String TIMING_TYPE_START_TIME_SUFFIX = "_start_time";
 
     /** 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(final Class<T> timingType, final Set<String> additionalProfilerTimings) {
+        Set<String> additionalTimings = additionalProfilerTimings == null ? Collections.emptySet() : additionalProfilerTimings;
+        timings = Stream.of(Arrays.stream(timingType.getEnumConstants()).map(Enum::name), additionalTimings.stream())
+            .flatMap(Function.identity())
+            .filter(Objects::nonNull)
+            .map(val -> val.toLowerCase(Locale.ROOT))
+            .collect(Collectors.toUnmodifiableMap(value -> value, value -> new Timer(), (a, b) -> a));
     }
 
     public Timer getTimer(T timing) {
-        return timings[timing.ordinal()];
+        return timings.get(timing.name().toLowerCase(Locale.ROOT));
+    }
+
+    public Timer getTimer(String timingName) {
+        return timings.get(timingName.toLowerCase(Locale.ROOT));
     }
 
     public void setTimer(T timing, Timer timer) {
-        timings[timing.ordinal()] = timer;
+        timings.put(timing.name().toLowerCase(Locale.ROOT), timer);
     }
 
     /**
      * Build a timing count 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 (String timingType : this.timings.keySet()) {
+            map.put(timingType, this.timings.get(timingType).getApproximateTiming());
+            map.put(timingType + TIMING_TYPE_COUNT_SUFFIX, this.timings.get(timingType).getCount());
+            map.put(timingType + TIMING_TYPE_START_TIME_SUFFIX, this.timings.get(timingType).getEarliestTimerStartTime());
         }
         return Collections.unmodifiableMap(map);
     }
@@ -94,8 +105,8 @@ public Map<String, Object> toDebugMap() {
 
     public long toNodeTime() {
         long total = 0;
-        for (T timingType : timingTypes) {
-            total += timings[timingType.ordinal()].getApproximateTiming();
+        for (String timingType : timings.keySet()) {
+            total += timings.get(timingType).getApproximateTiming();
         }
         return total;
     }

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

@@ -13,6 +13,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * Provide contextual profile breakdowns which are associated with freestyle context. Used when concurrent
@@ -21,8 +22,8 @@
  * @opensearch.internal
  */
 public abstract class ContextualProfileBreakdown<T extends Enum<T>> extends AbstractProfileBreakdown<T> {
-    public ContextualProfileBreakdown(Class<T> clazz) {
-        super(clazz);
+    public ContextualProfileBreakdown(Class<T> clazz, Set<String> additionalProfilerTimings) {
+        super(clazz, additionalProfilerTimings);
     }
 
     /**